Skip to content

Commit

Permalink
Add datastore namespace option in configs (#1581)
Browse files Browse the repository at this point in the history
Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai>
  • Loading branch information
Tsotne Tabidze authored May 24, 2021
1 parent 6e243b8 commit dfb029d
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 11 deletions.
16 changes: 7 additions & 9 deletions sdk/python/feast/infra/gcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 4 additions & 1 deletion sdk/python/feast/repo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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,
Expand Down
7 changes: 6 additions & 1 deletion sdk/python/tests/test_offline_online_store_consistency.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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])
Expand Down

0 comments on commit dfb029d

Please sign in to comment.