diff --git a/sdk/python/feast/infra/gcp.py b/sdk/python/feast/infra/gcp.py index aa5284c441..94e5639da5 100644 --- a/sdk/python/feast/infra/gcp.py +++ b/sdk/python/feast/infra/gcp.py @@ -36,23 +36,21 @@ class GcpProvider(Provider): _gcp_project_id: Optional[str] + _namespace: Optional[str] def __init__(self, config: RepoConfig): assert isinstance(config.online_store, DatastoreOnlineStoreConfig) - assert config.offline_store is not None - if config and config.online_store and config.online_store.project_id: - self._gcp_project_id = config.online_store.project_id - else: - self._gcp_project_id = None + self._gcp_project_id = config.online_store.project_id + self._namespace = config.online_store.namespace + assert config.offline_store is not None self.offline_store = get_offline_store_from_config(config.offline_store) def _initialize_client(self): try: - if self._gcp_project_id is not None: - return datastore.Client(self._gcp_project_id) - else: - return datastore.Client() + return datastore.Client( + project=self._gcp_project_id, namespace=self._namespace + ) except DefaultCredentialsError as e: raise FeastProviderLoginError( str(e) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index e1696324b9..472e0861df 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -35,6 +35,9 @@ class DatastoreOnlineStoreConfig(FeastBaseModel): project_id: Optional[StrictStr] = None """ (optional) GCP Project Id """ + namespace: Optional[StrictStr] = None + """ (optional) Datastore namespace """ + OnlineStoreConfig = Union[DatastoreOnlineStoreConfig, SqliteOnlineStoreConfig] @@ -138,7 +141,7 @@ def _validate_online_store_config(cls, values): elif online_store_type == "datastore": DatastoreOnlineStoreConfig(**values["online_store"]) else: - raise ValidationError(f"Invalid online store type {online_store_type}") + raise ValueError(f"Invalid online store type {online_store_type}") except ValidationError as e: raise ValidationError( [ErrorWrapper(e, loc="online_store")], model=SqliteOnlineStoreConfig, diff --git a/sdk/python/tests/test_offline_online_store_consistency.py b/sdk/python/tests/test_offline_online_store_consistency.py index e0d5a645f4..c39fb42754 100644 --- a/sdk/python/tests/test_offline_online_store_consistency.py +++ b/sdk/python/tests/test_offline_online_store_consistency.py @@ -17,7 +17,11 @@ from feast.feature import Feature from feast.feature_store import FeatureStore from feast.feature_view import FeatureView -from feast.repo_config import RepoConfig, SqliteOnlineStoreConfig +from feast.repo_config import ( + DatastoreOnlineStoreConfig, + RepoConfig, + SqliteOnlineStoreConfig, +) from feast.value_type import ValueType @@ -98,6 +102,7 @@ def prep_bq_fs_and_fv( registry=str(Path(repo_dir_name) / "registry.db"), project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}", provider="gcp", + online_store=DatastoreOnlineStoreConfig(namespace="integration_test"), ) fs = FeatureStore(config=config) fs.apply([fv, e])