Add datastore namespace option in configs (#1581)

Tsotne Tabidze · web-flow · commit dfb029df97ab · 2021-05-24T14:25:20.000-07:00
Signed-off-by: Tsotne Tabidze &lt;tsotne@tecton.ai&gt;
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])
