Skip to content

Commit 940e0f0

Browse files
YassinNouh21YassinNouh21
authored
feat: Add feature view versioning support to PostgreSQL and MySQL online stores (#6193)
* feat: Add feature view versioning support to PostgreSQL and MySQL online stores Add versioned read/write support so that version-qualified feature references (e.g., driver_stats@v2:trips_today) resolve to the correct versioned table in both PostgreSQL and MySQL online stores. Changes: - PostgreSQL: Updated _table_id() and all callers to support enable_versioning - MySQL: Updated _table_id(), _execute_batch(), write_to_table(), and _drop_table_and_index() to thread versioning flag through - online_store.py: Registered PostgreSQLOnlineStore and MySQLOnlineStore in _check_versioned_read_support() - errors.py: Updated VersionedOnlineReadNotSupported message - Unit tests split per store in tests/unit/infra/online_store/ - Integration tests in tests/integration/online_store/ (Docker, testcontainers) Closes #6168 Closes #6169 Part of #2728 Signed-off-by: yassinnouh21 <[email protected]> * fix: guard optional MySQL/Postgres imports in _check_versioned_read_support Unconditional imports of pymysql and psycopg broke online reads for users on other stores (Redis, DynamoDB, etc.) that don't have these optional dependencies installed. Signed-off-by: yassinnouh21 <[email protected]> * fix: drop all versioned tables on teardown and extract shared _table_id Teardown was only dropping the current version's table, orphaning older versioned tables (e.g. _v1, _v2). Now queries the database catalog to find and drop all versioned tables for each feature view. Also extracts the duplicated _table_id logic from SQLite, MySQL, and Postgres into a shared compute_table_id helper. Signed-off-by: yassinnouh21 <[email protected]> * style: fix ruff formatting for _table_id signatures Signed-off-by: yassinnouh21 <[email protected]> * style: collapse remaining multi-line signatures to single line Signed-off-by: yassinnouh21 <[email protected]> * fix: use regex for precise version table matching in teardown LIKE _v% could match unrelated tables (e.g. driver_verbose when dropping versions for driver). Switch to REGEXP/~ with _v[0-9]+$ pattern. Also scope Postgres pg_tables query by schemaname. Signed-off-by: yassinnouh21 <[email protected]> * style: wrap long _drop_all_version_tables signature Signed-off-by: yassinnouh21 <[email protected]> * fix: remove redundant DROP INDEX in MySQL _drop_table_and_index MySQL DROP TABLE already removes all indexes. The separate DROP INDEX was failing on versioned tables that were created via online_write_batch (which doesn't create the _ek index). Also reuse _drop_table_and_index in _drop_all_version_tables instead of duplicating logic. Signed-off-by: yassinnouh21 <[email protected]> --------- Signed-off-by: yassinnouh21 <[email protected]>
1 parent b46a62b commit 940e0f0

10 files changed

Lines changed: 743 additions & 44 deletions

File tree

‎sdk/python/feast/errors.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ class VersionedOnlineReadNotSupported(FeastError):
142142
def __init__(self, store_name: str, version: int):
143143
super().__init__(
144144
f"Versioned feature reads (@v{version}) are not yet supported by {store_name}. "
145-
f"Currently only SQLite supports version-qualified feature references. "
145+
f"Currently only SQLite, PostgreSQL, and MySQL support version-qualified feature references. "
146146
)
147147

148148

‎sdk/python/feast/infra/online_stores/helpers.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,3 +70,15 @@ def _to_naive_utc(ts: datetime) -> datetime:
7070
return ts
7171
else:
7272
return ts.astimezone(tz=timezone.utc).replace(tzinfo=None)
73+
74+
75+
def compute_table_id(project: str, table: Any, enable_versioning: bool = False) -> str:
76+
"""Build the online-store table name, appending a version suffix when versioning is enabled."""
77+
name = table.name
78+
if enable_versioning:
79+
version = getattr(table.projection, "version_tag", None)
80+
if version is None:
81+
version = getattr(table, "current_version_number", None)
82+
if version is not None and version > 0:
83+
name = f"{table.name}_v{version}"
84+
return f"{project}_{name}"

‎sdk/python/feast/infra/online_stores/mysql_online_store/mysql.py‎

Lines changed: 53 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
from feast import Entity, FeatureView, RepoConfig
1212
from feast.infra.key_encoding_utils import serialize_entity_key
13+
from feast.infra.online_stores.helpers import compute_table_id
1314
from feast.infra.online_stores.online_store import OnlineStore
1415
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
1516
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
@@ -70,6 +71,7 @@ def online_write_batch(
7071
cur = conn.cursor()
7172

7273
project = config.project
74+
versioning = config.registry.enable_online_feature_view_versioning
7375

7476
batch_write = config.online_store.batch_write
7577
if not batch_write:
@@ -92,6 +94,7 @@ def online_write_batch(
9294
table,
9395
timestamp,
9496
val,
97+
versioning,
9598
)
9699
conn.commit()
97100
if progress:
@@ -124,7 +127,9 @@ def online_write_batch(
124127

125128
if len(insert_values) >= batch_size:
126129
try:
127-
self._execute_batch(cur, project, table, insert_values)
130+
self._execute_batch(
131+
cur, project, table, insert_values, versioning
132+
)
128133
conn.commit()
129134
if progress:
130135
progress(len(insert_values))
@@ -135,17 +140,20 @@ def online_write_batch(
135140

136141
if insert_values:
137142
try:
138-
self._execute_batch(cur, project, table, insert_values)
143+
self._execute_batch(cur, project, table, insert_values, versioning)
139144
conn.commit()
140145
if progress:
141146
progress(len(insert_values))
142147
except Exception as e:
143148
conn.rollback()
144149
raise e
145150

146-
def _execute_batch(self, cur, project, table, insert_values):
147-
sql = f"""
148-
INSERT INTO {_table_id(project, table)}
151+
def _execute_batch(
152+
self, cur, project, table, insert_values, enable_versioning=False
153+
):
154+
table_name = _table_id(project, table, enable_versioning)
155+
stmt = f"""
156+
INSERT INTO {table_name}
149157
(entity_key, feature_name, value, event_ts, created_ts)
150158
values (%s, %s, %s, %s, %s)
151159
ON DUPLICATE KEY UPDATE
@@ -154,22 +162,29 @@ def _execute_batch(self, cur, project, table, insert_values):
154162
created_ts = VALUES(created_ts);
155163
"""
156164
try:
157-
cur.executemany(sql, insert_values)
165+
cur.executemany(stmt, insert_values)
158166
except Exception as e:
159-
# Log SQL info for debugging without leaking sensitive data
160167
first_sample = insert_values[0] if insert_values else None
161168
raise RuntimeError(
162-
f"Failed to execute batch insert into table '{_table_id(project, table)}' "
169+
f"Failed to execute batch insert into table '{table_name}' "
163170
f"(rows={len(insert_values)}, sample={first_sample}): {e}"
164171
) from e
165172

166173
@staticmethod
167174
def write_to_table(
168-
created_ts, cur, entity_key_bin, feature_name, project, table, timestamp, val
175+
created_ts,
176+
cur,
177+
entity_key_bin,
178+
feature_name,
179+
project,
180+
table,
181+
timestamp,
182+
val,
183+
enable_versioning=False,
169184
) -> None:
170185
cur.execute(
171186
f"""
172-
INSERT INTO {_table_id(project, table)}
187+
INSERT INTO {_table_id(project, table, enable_versioning)}
173188
(entity_key, feature_name, value, event_ts, created_ts)
174189
values (%s, %s, %s, %s, %s)
175190
ON DUPLICATE KEY UPDATE
@@ -204,14 +219,15 @@ def online_read(
204219
result: List[Tuple[Optional[datetime], Optional[Dict[str, Any]]]] = []
205220

206221
project = config.project
222+
versioning = config.registry.enable_online_feature_view_versioning
207223
for entity_key in entity_keys:
208224
entity_key_bin = serialize_entity_key(
209225
entity_key,
210226
entity_key_serialization_version=3,
211227
).hex()
212228

213229
cur.execute(
214-
f"SELECT feature_name, value, event_ts FROM {_table_id(project, table)} WHERE entity_key = %s",
230+
f"SELECT feature_name, value, event_ts FROM {_table_id(project, table, versioning)} WHERE entity_key = %s",
215231
(entity_key_bin,),
216232
)
217233

@@ -243,10 +259,11 @@ def update(
243259
conn = self._get_conn(config)
244260
cur = conn.cursor()
245261
project = config.project
262+
versioning = config.registry.enable_online_feature_view_versioning
246263

247264
# We don't create any special state for the entities in this implementation.
248265
for table in tables_to_keep:
249-
table_name = _table_id(project, table)
266+
table_name = _table_id(project, table, versioning)
250267
index_name = f"{table_name}_ek"
251268
cur.execute(
252269
f"""CREATE TABLE IF NOT EXISTS {table_name} (entity_key VARCHAR(512),
@@ -269,7 +286,10 @@ def update(
269286
)
270287

271288
for table in tables_to_delete:
272-
_drop_table_and_index(cur, project, table)
289+
if versioning:
290+
_drop_all_version_tables(cur, project, table)
291+
else:
292+
_drop_table_and_index(cur, _table_id(project, table))
273293

274294
def teardown(
275295
self,
@@ -280,16 +300,30 @@ def teardown(
280300
conn = self._get_conn(config)
281301
cur = conn.cursor()
282302
project = config.project
303+
versioning = config.registry.enable_online_feature_view_versioning
283304

284305
for table in tables:
285-
_drop_table_and_index(cur, project, table)
306+
if versioning:
307+
_drop_all_version_tables(cur, project, table)
308+
else:
309+
_drop_table_and_index(cur, _table_id(project, table))
286310

287311

288-
def _drop_table_and_index(cur: Cursor, project: str, table: FeatureView) -> None:
289-
table_name = _table_id(project, table)
290-
cur.execute(f"DROP INDEX {table_name}_ek ON {table_name};")
312+
def _drop_table_and_index(cur: Cursor, table_name: str) -> None:
291313
cur.execute(f"DROP TABLE IF EXISTS {table_name}")
292314

293315

294-
def _table_id(project: str, table: FeatureView) -> str:
295-
return f"{project}_{table.name}"
316+
def _drop_all_version_tables(cur: Cursor, project: str, table: FeatureView) -> None:
317+
"""Drop the base table and all versioned tables (e.g. _v1, _v2, ...)."""
318+
base = f"{project}_{table.name}"
319+
cur.execute(
320+
"SELECT table_name FROM information_schema.tables "
321+
"WHERE table_schema = DATABASE() AND (table_name = %s OR table_name REGEXP %s)",
322+
(base, f"^{base}_v[0-9]+$"),
323+
)
324+
for (name,) in cur.fetchall():
325+
_drop_table_and_index(cur, name)
326+
327+
328+
def _table_id(project: str, table: FeatureView, enable_versioning: bool = False) -> str:
329+
return compute_table_id(project, table, enable_versioning)

‎sdk/python/feast/infra/online_stores/online_store.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,25 @@ def _check_versioned_read_support(self, grouped_refs):
257257
"""Raise an error if versioned reads are attempted on unsupported stores."""
258258
from feast.infra.online_stores.sqlite import SqliteOnlineStore
259259

260-
if isinstance(self, SqliteOnlineStore):
260+
supported_types: list[type] = [SqliteOnlineStore]
261+
try:
262+
from feast.infra.online_stores.mysql_online_store.mysql import (
263+
MySQLOnlineStore,
264+
)
265+
266+
supported_types.append(MySQLOnlineStore)
267+
except ImportError:
268+
pass
269+
try:
270+
from feast.infra.online_stores.postgres_online_store.postgres import (
271+
PostgreSQLOnlineStore,
272+
)
273+
274+
supported_types.append(PostgreSQLOnlineStore)
275+
except ImportError:
276+
pass
277+
278+
if isinstance(self, tuple(supported_types)):
261279
return
262280
for table, _ in grouped_refs:
263281
version_tag = getattr(table.projection, "version_tag", None)

‎sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py‎

Lines changed: 68 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222

2323
from feast import Entity, FeatureView, ValueType
2424
from feast.infra.key_encoding_utils import get_list_val_str, serialize_entity_key
25-
from feast.infra.online_stores.helpers import _to_naive_utc
25+
from feast.infra.online_stores.helpers import _to_naive_utc, compute_table_id
2626
from feast.infra.online_stores.online_store import OnlineStore
2727
from feast.infra.online_stores.vector_store import VectorStoreConfig
2828
from feast.infra.utils.postgres.connection_utils import (
@@ -152,7 +152,15 @@ def online_write_batch(
152152
event_ts = EXCLUDED.event_ts,
153153
created_ts = EXCLUDED.created_ts;
154154
"""
155-
).format(sql.Identifier(_table_id(config.project, table)))
155+
).format(
156+
sql.Identifier(
157+
_table_id(
158+
config.project,
159+
table,
160+
config.registry.enable_online_feature_view_versioning,
161+
)
162+
)
163+
)
156164

157165
# Push data into the online store
158166
with self._get_conn(config) as conn, conn.cursor() as cur:
@@ -214,7 +222,13 @@ def _construct_query_and_params(
214222
FROM {} WHERE entity_key = ANY(%s) AND feature_name = ANY(%s);
215223
"""
216224
).format(
217-
sql.Identifier(_table_id(config.project, table)),
225+
sql.Identifier(
226+
_table_id(
227+
config.project,
228+
table,
229+
config.registry.enable_online_feature_view_versioning,
230+
)
231+
),
218232
)
219233
params = (keys, requested_features)
220234
else:
@@ -224,7 +238,13 @@ def _construct_query_and_params(
224238
FROM {} WHERE entity_key = ANY(%s);
225239
"""
226240
).format(
227-
sql.Identifier(_table_id(config.project, table)),
241+
sql.Identifier(
242+
_table_id(
243+
config.project,
244+
table,
245+
config.registry.enable_online_feature_view_versioning,
246+
)
247+
),
228248
)
229249
params = (keys, [])
230250
return query, params
@@ -304,12 +324,16 @@ def update(
304324
),
305325
)
306326

327+
versioning = config.registry.enable_online_feature_view_versioning
307328
for table in tables_to_delete:
308-
table_name = _table_id(project, table)
309-
cur.execute(_drop_table_and_index(table_name))
329+
if versioning:
330+
_drop_all_version_tables(cur, project, table, schema_name)
331+
else:
332+
table_name = _table_id(project, table)
333+
cur.execute(_drop_table_and_index(table_name))
310334

311335
for table in tables_to_keep:
312-
table_name = _table_id(project, table)
336+
table_name = _table_id(project, table, versioning)
313337
if config.online_store.vector_enabled:
314338
vector_value_type = "vector"
315339
else:
@@ -363,11 +387,16 @@ def teardown(
363387
entities: Sequence[Entity],
364388
):
365389
project = config.project
390+
schema_name = config.online_store.db_schema or config.online_store.user
391+
versioning = config.registry.enable_online_feature_view_versioning
366392
try:
367393
with self._get_conn(config) as conn, conn.cursor() as cur:
368394
for table in tables:
369-
table_name = _table_id(project, table)
370-
cur.execute(_drop_table_and_index(table_name))
395+
if versioning:
396+
_drop_all_version_tables(cur, project, table, schema_name)
397+
else:
398+
table_name = _table_id(project, table)
399+
cur.execute(_drop_table_and_index(table_name))
371400
conn.commit()
372401
except Exception:
373402
logging.exception("Teardown failed")
@@ -432,7 +461,9 @@ def retrieve_online_documents(
432461
]
433462
] = []
434463
with self._get_conn(config, autocommit=True) as conn, conn.cursor() as cur:
435-
table_name = _table_id(project, table)
464+
table_name = _table_id(
465+
project, table, config.registry.enable_online_feature_view_versioning
466+
)
436467

437468
# Search query template to find the top k items that are closest to the given embedding
438469
# SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;
@@ -533,7 +564,11 @@ def retrieve_online_documents_v2(
533564
and feature.name in requested_features
534565
]
535566

536-
table_name = _table_id(config.project, table)
567+
table_name = _table_id(
568+
config.project,
569+
table,
570+
config.registry.enable_online_feature_view_versioning,
571+
)
537572

538573
with self._get_conn(config, autocommit=True) as conn, conn.cursor() as cur:
539574
query = None
@@ -794,8 +829,8 @@ def retrieve_online_documents_v2(
794829
return result
795830

796831

797-
def _table_id(project: str, table: FeatureView) -> str:
798-
return f"{project}_{table.name}"
832+
def _table_id(project: str, table: FeatureView, enable_versioning: bool = False) -> str:
833+
return compute_table_id(project, table, enable_versioning)
799834

800835

801836
def _drop_table_and_index(table_name):
@@ -808,3 +843,23 @@ def _drop_table_and_index(table_name):
808843
sql.Identifier(table_name),
809844
sql.Identifier(f"{table_name}_ek"),
810845
)
846+
847+
848+
def _drop_all_version_tables(
849+
cur, project: str, table: FeatureView, schema_name: Optional[str] = None
850+
) -> None:
851+
"""Drop the base table and all versioned tables (e.g. _v1, _v2, ...)."""
852+
base = f"{project}_{table.name}"
853+
if schema_name:
854+
cur.execute(
855+
"SELECT tablename FROM pg_tables "
856+
"WHERE schemaname = %s AND (tablename = %s OR tablename ~ %s)",
857+
(schema_name, base, f"^{base}_v[0-9]+$"),
858+
)
859+
else:
860+
cur.execute(
861+
"SELECT tablename FROM pg_tables WHERE tablename = %s OR tablename ~ %s",
862+
(base, f"^{base}_v[0-9]+$"),
863+
)
864+
for (name,) in cur.fetchall():
865+
cur.execute(_drop_table_and_index(name))

0 commit comments

Comments
 (0)