Add datastore namespace option in configs

Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai>
feast-dev · May 24, 2021 · c179f88 · c179f88
1 parent 6e243b8
commit c179f88
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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])