diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py index 515a6c39b11..8fccb8aaffb 100644 --- a/sdk/python/feast/errors.py +++ b/sdk/python/feast/errors.py @@ -142,7 +142,7 @@ class VersionedOnlineReadNotSupported(FeastError): def __init__(self, store_name: str, version: int): super().__init__( f"Versioned feature reads (@v{version}) are not yet supported by {store_name}. " - f"Currently only SQLite, PostgreSQL, MySQL, FAISS, Redis, and DynamoDB support version-qualified feature references. " + f"Currently only SQLite, PostgreSQL, MySQL, FAISS, Redis, DynamoDB, Milvus, and Couchbase support version-qualified feature references. " ) diff --git a/sdk/python/feast/infra/online_stores/couchbase_online_store/couchbase.py b/sdk/python/feast/infra/online_stores/couchbase_online_store/couchbase.py index c80f9e1285c..76de26b85bb 100644 --- a/sdk/python/feast/infra/online_stores/couchbase_online_store/couchbase.py +++ b/sdk/python/feast/infra/online_stores/couchbase_online_store/couchbase.py @@ -17,6 +17,7 @@ from feast import Entity, FeatureView, RepoConfig from feast.infra.key_encoding_utils import serialize_entity_key +from feast.infra.online_stores.helpers import compute_table_id from feast.infra.online_stores.online_store import OnlineStore from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto @@ -105,8 +106,7 @@ def online_write_batch( RuntimeWarning, ) project = config.project - scope_name = f"{project}_{table.name}_scope" - collection_name = f"{project}_{table.name}_collection" + scope_name, collection_name = _scope_and_collection(config, table) collection = self._get_conn(config, scope_name, collection_name) for entity_key, values, timestamp, created_ts in data: @@ -171,9 +171,7 @@ def online_read( RuntimeWarning, ) project = config.project - - scope_name = f"{project}_{table.name}_scope" - collection_name = f"{project}_{table.name}_collection" + scope_name, collection_name = _scope_and_collection(config, table) collection = self._get_conn(config, scope_name, collection_name) @@ -239,11 +237,8 @@ def update( "Some functionality may still be unstable so functionality can change in the future.", RuntimeWarning, ) - project = config.project - for table in tables_to_keep: - scope_name = f"{project}_{table.name}_scope" - collection_name = f"{project}_{table.name}_collection" + scope_name, collection_name = _scope_and_collection(config, table) self._get_conn(config, scope_name, collection_name) cm = self.bucket.collections() @@ -288,11 +283,8 @@ def teardown( "Some functionality may still be unstable so functionality can change in the future.", RuntimeWarning, ) - project = config.project - for table in tables: - scope_name = f"{project}_{table.name}_scope" - collection_name = f"{project}_{table.name}_collection" + scope_name, collection_name = _scope_and_collection(config, table) self._get_conn(config, scope_name, collection_name) cm = self.bucket.collections() try: @@ -302,6 +294,21 @@ def teardown( logger.error(f"Error removing collection or scope: {e}") +def _scope_and_collection(config: RepoConfig, table: FeatureView) -> Tuple[str, str]: + """Couchbase scope and collection backing a feature view. + + Built on ``compute_table_id``, so with + ``registry.enable_online_feature_view_versioning`` enabled each version gets its + own scope and collection (``{project}_{name}_v2_scope``). Document ids stay + unversioned: a collection is already a namespace, so partitioning at the + collection level is enough to keep versions from colliding. + """ + base = compute_table_id( + config.project, table, config.registry.enable_online_feature_view_versioning + ) + return f"{base}_scope", f"{base}_collection" + + def _document_id( project: str, table: FeatureView, entity_key_str: str, feature_name: str ) -> str: diff --git a/sdk/python/feast/infra/online_stores/online_store.py b/sdk/python/feast/infra/online_stores/online_store.py index cdf06639fe0..b3ae5b5e94d 100644 --- a/sdk/python/feast/infra/online_stores/online_store.py +++ b/sdk/python/feast/infra/online_stores/online_store.py @@ -315,6 +315,10 @@ def _is_versioned_read_supported(self) -> bool: "feast.infra.online_stores.milvus_online_store.milvus", "MilvusOnlineStore", ), + ( + "feast.infra.online_stores.couchbase_online_store.couchbase", + "CouchbaseOnlineStore", + ), ): try: import importlib diff --git a/sdk/python/tests/unit/infra/online_store/test_couchbase_versioning.py b/sdk/python/tests/unit/infra/online_store/test_couchbase_versioning.py new file mode 100644 index 00000000000..21eba2bae70 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_couchbase_versioning.py @@ -0,0 +1,140 @@ +"""Unit tests for Couchbase online store feature view versioning.""" + +from datetime import timedelta +from unittest.mock import MagicMock, patch + +import pytest + +from feast import Entity, FeatureView +from feast.field import Field +from feast.types import Float32 +from feast.value_type import ValueType + +pytest.importorskip("couchbase") + +from feast.infra.online_stores.couchbase_online_store.couchbase import ( # noqa: E402 + CouchbaseOnlineStore, + _scope_and_collection, +) + + +def _make_feature_view(name="driver_stats", version_number=None, version_tag=None): + entity = Entity( + name="driver_id", join_keys=["driver_id"], value_type=ValueType.INT64 + ) + fv = FeatureView( + name=name, + entities=[entity], + ttl=timedelta(days=1), + schema=[Field(name="trips_today", dtype=Float32)], + ) + if version_number is not None: + fv.current_version_number = version_number + if version_tag is not None: + fv.projection.version_tag = version_tag + return fv + + +def _make_config(project="test_project", versioning=False): + config = MagicMock() + config.project = project + config.entity_key_serialization_version = 2 + config.registry.enable_online_feature_view_versioning = versioning + return config + + +class TestScopeAndCollection: + def test_no_versioning(self): + scope, collection = _scope_and_collection( + _make_config(versioning=False), _make_feature_view() + ) + assert scope == "test_project_driver_stats_scope" + assert collection == "test_project_driver_stats_collection" + + def test_versioning_disabled_ignores_version(self): + """A version on the feature view must not change names while the flag is off.""" + scope, collection = _scope_and_collection( + _make_config(versioning=False), _make_feature_view(version_number=2) + ) + assert scope == "test_project_driver_stats_scope" + assert collection == "test_project_driver_stats_collection" + + def test_versioning_enabled_with_version(self): + scope, collection = _scope_and_collection( + _make_config(versioning=True), _make_feature_view(version_number=2) + ) + assert scope == "test_project_driver_stats_v2_scope" + assert collection == "test_project_driver_stats_v2_collection" + + def test_projection_version_tag_takes_priority(self): + scope, _ = _scope_and_collection( + _make_config(versioning=True), + _make_feature_view(version_number=1, version_tag=3), + ) + assert scope == "test_project_driver_stats_v3_scope" + + def test_version_zero_no_suffix(self): + scope, _ = _scope_and_collection( + _make_config(versioning=True), _make_feature_view(version_number=0) + ) + assert scope == "test_project_driver_stats_scope" + + def test_versions_do_not_collide(self): + config = _make_config(versioning=True) + v1, _ = _scope_and_collection(config, _make_feature_view(version_number=1)) + v2, _ = _scope_and_collection(config, _make_feature_view(version_number=2)) + assert v1 != v2 + + +class TestVersionedNamesReachCouchbase: + """The resolved names must be what the store actually connects with.""" + + @pytest.mark.parametrize( + "versioning, expected_scope", + [ + (False, "test_project_driver_stats_scope"), + (True, "test_project_driver_stats_v2_scope"), + ], + ) + def test_update_creates_the_versioned_scope(self, versioning, expected_scope): + store = CouchbaseOnlineStore() + config = _make_config(versioning=versioning) + fv = _make_feature_view(version_number=2) + + with patch.object(CouchbaseOnlineStore, "_get_conn") as get_conn: + store.bucket = MagicMock() + store.update(config, [], [fv], [], [], partial=False) + + get_conn.assert_called_once_with( + config, expected_scope, expected_scope.replace("_scope", "_collection") + ) + store.bucket.collections().create_scope.assert_called_once_with( + expected_scope + ) + + def test_teardown_drops_the_versioned_scope(self): + store = CouchbaseOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + + with patch.object(CouchbaseOnlineStore, "_get_conn"): + store.bucket = MagicMock() + store.teardown(config, [fv], []) + + store.bucket.collections().drop_scope.assert_called_once_with( + "test_project_driver_stats_v2_scope" + ) + + +class TestVersionedReadSupport: + def test_couchbase_is_allowlisted_for_versioned_reads(self): + """Couchbase's online_read honours the contract, so it must not raise.""" + store = CouchbaseOnlineStore() + store._versioned_read_supported = None + assert store._is_versioned_read_supported() is True + + def test_versioned_ref_does_not_raise(self): + store = CouchbaseOnlineStore() + store._versioned_read_supported = None + fv = _make_feature_view(version_tag=2) + store._check_versioned_read_support([(fv, ["trips_today"])])