From 7553c0f19fb814b1688fbeabb873ec9f61bb29b1 Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Mon, 13 Dec 2021 16:22:44 -0800 Subject: [PATCH 1/3] Add DatastoreTable infra object Signed-off-by: Felix Wang --- protos/feast/core/DatastoreTable.proto | 43 ++++++ protos/feast/core/InfraObject.proto | 2 + .../feast/infra/online_stores/datastore.py | 124 +++++++++++++++--- 3 files changed, 152 insertions(+), 17 deletions(-) create mode 100644 protos/feast/core/DatastoreTable.proto diff --git a/protos/feast/core/DatastoreTable.proto b/protos/feast/core/DatastoreTable.proto new file mode 100644 index 00000000000..3e0d2225d5f --- /dev/null +++ b/protos/feast/core/DatastoreTable.proto @@ -0,0 +1,43 @@ +// +// * Copyright 2021 The Feast Authors +// * +// * Licensed under the Apache License, Version 2.0 (the "License"); +// * you may not use this file except in compliance with the License. +// * You may obtain a copy of the License at +// * +// * https://www.apache.org/licenses/LICENSE-2.0 +// * +// * Unless required by applicable law or agreed to in writing, software +// * distributed under the License is distributed on an "AS IS" BASIS, +// * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// * See the License for the specific language governing permissions and +// * limitations under the License. +// + +syntax = "proto3"; + +package feast.core; +option java_package = "feast.proto.core"; +option java_outer_classname = "DatastoreTableProto"; +option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/core"; + +// Represents a Datastore table +message DatastoreTable { + // Feast project of the table + string project = 1; + + // Name of the table + string name = 2; + + // GCP project id, which can be null + oneof project_id { + bool project_id_null = 3; + string project_id_value = 4; + } + + // Datastore namespace, which can be null + oneof namespace { + bool namespace_null = 5; + string namespace_value = 6; + } +} \ No newline at end of file diff --git a/protos/feast/core/InfraObject.proto b/protos/feast/core/InfraObject.proto index ded4c3ed68c..a0f3541dec6 100644 --- a/protos/feast/core/InfraObject.proto +++ b/protos/feast/core/InfraObject.proto @@ -22,6 +22,7 @@ option java_outer_classname = "InfraObjectProto"; option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/core"; import "feast/core/DynamoDBTable.proto"; +import "feast/core/DatastoreTable.proto"; // Represents a set of infrastructure objects managed by Feast message Infra { @@ -37,6 +38,7 @@ message InfraObject { // The infrastructure object oneof infra_object { DynamoDBTable dynamodb_table = 2; + DatastoreTable datastore_table = 3; CustomInfra custom_infra = 100; } diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py index e9e5973eddd..65143eebb49 100644 --- a/sdk/python/feast/infra/online_stores/datastore.py +++ b/sdk/python/feast/infra/online_stores/datastore.py @@ -21,8 +21,13 @@ from feast import Entity, FeatureTable, utils from feast.feature_view import FeatureView +from feast.infra.infra_object import InfraObject from feast.infra.online_stores.helpers import compute_entity_id from feast.infra.online_stores.online_store import OnlineStore +from feast.protos.feast.core.DatastoreTable_pb2 import ( + DatastoreTable as DatastoreTableProto, +) +from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto from feast.repo_config import FeastConfigBaseModel, RepoConfig @@ -80,8 +85,6 @@ def update( entities_to_keep: Sequence[Entity], partial: bool, ): - """ - """ online_config = config.online_store assert isinstance(online_config, DatastoreOnlineStoreConfig) client = self._get_client(online_config) @@ -110,9 +113,6 @@ def teardown( tables: Sequence[Union[FeatureTable, FeatureView]], entities: Sequence[Entity], ): - """ - There's currently no teardown done for Datastore. - """ online_config = config.online_store assert isinstance(online_config, DatastoreOnlineStoreConfig) client = self._get_client(online_config) @@ -128,18 +128,10 @@ def teardown( client.delete(key) def _get_client(self, online_config: DatastoreOnlineStoreConfig): - if not self._client: - try: - self._client = datastore.Client( - project=online_config.project_id, namespace=online_config.namespace, - ) - except DefaultCredentialsError as e: - raise FeastProviderLoginError( - str(e) - + '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your ' - "local Google Cloud account " - ) + self._client = _initialize_client( + online_config.project_id, online_config.namespace + ) return self._client @log_exceptions_and_usage(online_store="datastore") @@ -267,7 +259,7 @@ def online_read( return result -def _delete_all_values(client, key) -> None: +def _delete_all_values(client, key): """ Delete all data under the key path in datastore. """ @@ -279,3 +271,101 @@ def _delete_all_values(client, key) -> None: for entity in entities: client.delete(entity.key) + + +def _initialize_client( + project_id: Optional[str], namespace: Optional[str] +) -> datastore.Client: + try: + client = datastore.Client(project=project_id, namespace=namespace,) + return client + except DefaultCredentialsError as e: + raise FeastProviderLoginError( + str(e) + + '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your ' + "local Google Cloud account " + ) + + +class DatastoreTable(InfraObject): + """ + A Datastore table managed by Feast. + + Attributes: + project: The Feast project of the table. + name: The name of the table. + project_id (optional): The GCP project id. + namespace (optional): Datastore namespace. + """ + + project: str + name: str + project_id: Optional[str] + namespace: Optional[str] + + def __init__( + self, + project: str, + name: str, + project_id: Optional[str], + namespace: Optional[str], + ): + self.project = project + self.name = name + self.project_id = project_id + self.namespace = namespace + + def to_proto(self) -> InfraObjectProto: + datastore_table_proto = DatastoreTableProto() + datastore_table_proto.project = self.project + datastore_table_proto.name = self.name + if self.project_id is None: + datastore_table_proto.project_id_null = True + else: + datastore_table_proto.project_id_value = self.project_id + if self.namespace is None: + datastore_table_proto.namespace_null = True + else: + datastore_table_proto.namespace_value = self.namespace + + return InfraObjectProto( + infra_object_class_type="feast.infra.online_stores.datastore.DatastoreTable", + datastore_table=datastore_table_proto, + ) + + @staticmethod + def from_proto(infra_object_proto: InfraObjectProto) -> Any: + project_id_value = ( + None + if infra_object_proto.datastore_table.project_id_null + else infra_object_proto.datastore_table.project_id_value + ) + namespace_value = ( + None + if infra_object_proto.datastore_table.namespace_null + else infra_object_proto.datastore_table.namespace_value + ) + + return DatastoreTable( + project=infra_object_proto.datastore_table.project, + name=infra_object_proto.datastore_table.name, + project_id=project_id_value, + namespace=namespace_value, + ) + + def update(self): + client = _initialize_client(self.project_id, self.namespace) + key = client.key("Project", self.project, "Table", self.name) + entity = datastore.Entity( + key=key, exclude_from_indexes=("created_ts", "event_ts", "values") + ) + entity.update({"created_ts": datetime.utcnow()}) + client.put(entity) + + def teardown(self): + client = _initialize_client(self.project_id, self.namespace) + key = client.key("Project", self.project, "Table", self.name) + _delete_all_values(client, key) + + # Delete the table metadata datastore entity + client.delete(key) From 0bc433abec41c3b34ab29edce9f433f80e6bcad4 Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Tue, 14 Dec 2021 15:24:11 -0800 Subject: [PATCH 2/3] Switch to StringValue Signed-off-by: Felix Wang --- protos/feast/core/DatastoreTable.proto | 16 +++---- .../feast/infra/online_stores/datastore.py | 42 ++++++++----------- 2 files changed, 24 insertions(+), 34 deletions(-) diff --git a/protos/feast/core/DatastoreTable.proto b/protos/feast/core/DatastoreTable.proto index 3e0d2225d5f..15720ad809c 100644 --- a/protos/feast/core/DatastoreTable.proto +++ b/protos/feast/core/DatastoreTable.proto @@ -21,6 +21,8 @@ option java_package = "feast.proto.core"; option java_outer_classname = "DatastoreTableProto"; option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/core"; +import "google/protobuf/wrappers.proto"; + // Represents a Datastore table message DatastoreTable { // Feast project of the table @@ -29,15 +31,9 @@ message DatastoreTable { // Name of the table string name = 2; - // GCP project id, which can be null - oneof project_id { - bool project_id_null = 3; - string project_id_value = 4; - } + // GCP project id + google.protobuf.StringValue project_id = 3; - // Datastore namespace, which can be null - oneof namespace { - bool namespace_null = 5; - string namespace_value = 6; - } + // Datastore namespace + google.protobuf.StringValue namespace = 4; } \ No newline at end of file diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py index 65143eebb49..6d56421edf1 100644 --- a/sdk/python/feast/infra/online_stores/datastore.py +++ b/sdk/python/feast/infra/online_stores/datastore.py @@ -307,8 +307,8 @@ def __init__( self, project: str, name: str, - project_id: Optional[str], - namespace: Optional[str], + project_id: Optional[str] = None, + namespace: Optional[str] = None, ): self.project = project self.name = name @@ -319,14 +319,10 @@ def to_proto(self) -> InfraObjectProto: datastore_table_proto = DatastoreTableProto() datastore_table_proto.project = self.project datastore_table_proto.name = self.name - if self.project_id is None: - datastore_table_proto.project_id_null = True - else: - datastore_table_proto.project_id_value = self.project_id - if self.namespace is None: - datastore_table_proto.namespace_null = True - else: - datastore_table_proto.namespace_value = self.namespace + if self.project_id: + datastore_table_proto.project_id.FromString(bytes(self.project_id, "utf-8")) + if self.namespace: + datastore_table_proto.namespace.FromString(bytes(self.namespace, "utf-8")) return InfraObjectProto( infra_object_class_type="feast.infra.online_stores.datastore.DatastoreTable", @@ -335,24 +331,22 @@ def to_proto(self) -> InfraObjectProto: @staticmethod def from_proto(infra_object_proto: InfraObjectProto) -> Any: - project_id_value = ( - None - if infra_object_proto.datastore_table.project_id_null - else infra_object_proto.datastore_table.project_id_value - ) - namespace_value = ( - None - if infra_object_proto.datastore_table.namespace_null - else infra_object_proto.datastore_table.namespace_value - ) - - return DatastoreTable( + datastore_table = DatastoreTable( project=infra_object_proto.datastore_table.project, name=infra_object_proto.datastore_table.name, - project_id=project_id_value, - namespace=namespace_value, ) + if infra_object_proto.datastore_table.HasField("project_id"): + datastore_table.project_id = ( + infra_object_proto.datastore_table.project_id.SerializeToString() + ).decode("utf-8") + if infra_object_proto.datastore_table.HasField("namespace"): + datastore_table.namespace = ( + infra_object_proto.datastore_table.namespace.SerializeToString() + ).decode("utf-8") + + return datastore_table + def update(self): client = _initialize_client(self.project_id, self.namespace) key = client.key("Project", self.project, "Table", self.name) From fad30d48ae5718ce7818a7fb7447c9112b6b3ff0 Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Tue, 14 Dec 2021 15:26:09 -0800 Subject: [PATCH 3/3] Initialize Datastore client in __init__ Signed-off-by: Felix Wang --- sdk/python/feast/infra/online_stores/datastore.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/datastore.py b/sdk/python/feast/infra/online_stores/datastore.py index 6d56421edf1..3aa55f6be41 100644 --- a/sdk/python/feast/infra/online_stores/datastore.py +++ b/sdk/python/feast/infra/online_stores/datastore.py @@ -296,12 +296,14 @@ class DatastoreTable(InfraObject): name: The name of the table. project_id (optional): The GCP project id. namespace (optional): Datastore namespace. + client: Datastore client. """ project: str name: str project_id: Optional[str] namespace: Optional[str] + client: datastore.Client def __init__( self, @@ -314,6 +316,7 @@ def __init__( self.name = name self.project_id = project_id self.namespace = namespace + self.client = _initialize_client(self.project_id, self.namespace) def to_proto(self) -> InfraObjectProto: datastore_table_proto = DatastoreTableProto() @@ -348,18 +351,16 @@ def from_proto(infra_object_proto: InfraObjectProto) -> Any: return datastore_table def update(self): - client = _initialize_client(self.project_id, self.namespace) - key = client.key("Project", self.project, "Table", self.name) + key = self.client.key("Project", self.project, "Table", self.name) entity = datastore.Entity( key=key, exclude_from_indexes=("created_ts", "event_ts", "values") ) entity.update({"created_ts": datetime.utcnow()}) - client.put(entity) + self.client.put(entity) def teardown(self): - client = _initialize_client(self.project_id, self.namespace) - key = client.key("Project", self.project, "Table", self.name) - _delete_all_values(client, key) + key = self.client.key("Project", self.project, "Table", self.name) + _delete_all_values(self.client, key) # Delete the table metadata datastore entity - client.delete(key) + self.client.delete(key)