Skip to content

Commit 3cd9678

Browse files
felixwang9817felixwang9817
authored
Add DatastoreTable infra object (#2140)
* Add DatastoreTable infra object Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Switch to StringValue Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Initialize Datastore client in __init__ Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 107eddc commit 3cd9678

3 files changed

Lines changed: 143 additions & 17 deletions

File tree

‎protos/feast/core/DatastoreTable.proto‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
//
2+
// * Copyright 2021 The Feast Authors
3+
// *
4+
// * Licensed under the Apache License, Version 2.0 (the "License");
5+
// * you may not use this file except in compliance with the License.
6+
// * You may obtain a copy of the License at
7+
// *
8+
// * https://www.apache.org/licenses/LICENSE-2.0
9+
// *
10+
// * Unless required by applicable law or agreed to in writing, software
11+
// * distributed under the License is distributed on an "AS IS" BASIS,
12+
// * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// * See the License for the specific language governing permissions and
14+
// * limitations under the License.
15+
//
16+
17+
syntax = "proto3";
18+
19+
package feast.core;
20+
option java_package = "feast.proto.core";
21+
option java_outer_classname = "DatastoreTableProto";
22+
option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/core";
23+
24+
import "google/protobuf/wrappers.proto";
25+
26+
// Represents a Datastore table
27+
message DatastoreTable {
28+
// Feast project of the table
29+
string project = 1;
30+
31+
// Name of the table
32+
string name = 2;
33+
34+
// GCP project id
35+
google.protobuf.StringValue project_id = 3;
36+
37+
// Datastore namespace
38+
google.protobuf.StringValue namespace = 4;
39+
}

‎protos/feast/core/InfraObject.proto‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ option java_outer_classname = "InfraObjectProto";
2222
option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/core";
2323

2424
import "feast/core/DynamoDBTable.proto";
25+
import "feast/core/DatastoreTable.proto";
2526

2627
// Represents a set of infrastructure objects managed by Feast
2728
message Infra {
@@ -37,6 +38,7 @@ message InfraObject {
3738
// The infrastructure object
3839
oneof infra_object {
3940
DynamoDBTable dynamodb_table = 2;
41+
DatastoreTable datastore_table = 3;
4042
CustomInfra custom_infra = 100;
4143
}
4244

‎sdk/python/feast/infra/online_stores/datastore.py‎

Lines changed: 102 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,13 @@
2121

2222
from feast import Entity, utils
2323
from feast.feature_view import FeatureView
24+
from feast.infra.infra_object import InfraObject
2425
from feast.infra.online_stores.helpers import compute_entity_id
2526
from feast.infra.online_stores.online_store import OnlineStore
27+
from feast.protos.feast.core.DatastoreTable_pb2 import (
28+
DatastoreTable as DatastoreTableProto,
29+
)
30+
from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto
2631
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
2732
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
2833
from feast.repo_config import FeastConfigBaseModel, RepoConfig
@@ -80,8 +85,6 @@ def update(
8085
entities_to_keep: Sequence[Entity],
8186
partial: bool,
8287
):
83-
"""
84-
"""
8588
online_config = config.online_store
8689
assert isinstance(online_config, DatastoreOnlineStoreConfig)
8790
client = self._get_client(online_config)
@@ -110,9 +113,6 @@ def teardown(
110113
tables: Sequence[FeatureView],
111114
entities: Sequence[Entity],
112115
):
113-
"""
114-
There's currently no teardown done for Datastore.
115-
"""
116116
online_config = config.online_store
117117
assert isinstance(online_config, DatastoreOnlineStoreConfig)
118118
client = self._get_client(online_config)
@@ -128,18 +128,10 @@ def teardown(
128128
client.delete(key)
129129

130130
def _get_client(self, online_config: DatastoreOnlineStoreConfig):
131-
132131
if not self._client:
133-
try:
134-
self._client = datastore.Client(
135-
project=online_config.project_id, namespace=online_config.namespace,
136-
)
137-
except DefaultCredentialsError as e:
138-
raise FeastProviderLoginError(
139-
str(e)
140-
+ '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your '
141-
"local Google Cloud account "
142-
)
132+
self._client = _initialize_client(
133+
online_config.project_id, online_config.namespace
134+
)
143135
return self._client
144136

145137
@log_exceptions_and_usage(online_store="datastore")
@@ -267,7 +259,7 @@ def online_read(
267259
return result
268260

269261

270-
def _delete_all_values(client, key) -> None:
262+
def _delete_all_values(client, key):
271263
"""
272264
Delete all data under the key path in datastore.
273265
"""
@@ -279,3 +271,96 @@ def _delete_all_values(client, key) -> None:
279271

280272
for entity in entities:
281273
client.delete(entity.key)
274+
275+
276+
def _initialize_client(
277+
project_id: Optional[str], namespace: Optional[str]
278+
) -> datastore.Client:
279+
try:
280+
client = datastore.Client(project=project_id, namespace=namespace,)
281+
return client
282+
except DefaultCredentialsError as e:
283+
raise FeastProviderLoginError(
284+
str(e)
285+
+ '\nIt may be necessary to run "gcloud auth application-default login" if you would like to use your '
286+
"local Google Cloud account "
287+
)
288+
289+
290+
class DatastoreTable(InfraObject):
291+
"""
292+
A Datastore table managed by Feast.
293+
294+
Attributes:
295+
project: The Feast project of the table.
296+
name: The name of the table.
297+
project_id (optional): The GCP project id.
298+
namespace (optional): Datastore namespace.
299+
client: Datastore client.
300+
"""
301+
302+
project: str
303+
name: str
304+
project_id: Optional[str]
305+
namespace: Optional[str]
306+
client: datastore.Client
307+
308+
def __init__(
309+
self,
310+
project: str,
311+
name: str,
312+
project_id: Optional[str] = None,
313+
namespace: Optional[str] = None,
314+
):
315+
self.project = project
316+
self.name = name
317+
self.project_id = project_id
318+
self.namespace = namespace
319+
self.client = _initialize_client(self.project_id, self.namespace)
320+
321+
def to_proto(self) -> InfraObjectProto:
322+
datastore_table_proto = DatastoreTableProto()
323+
datastore_table_proto.project = self.project
324+
datastore_table_proto.name = self.name
325+
if self.project_id:
326+
datastore_table_proto.project_id.FromString(bytes(self.project_id, "utf-8"))
327+
if self.namespace:
328+
datastore_table_proto.namespace.FromString(bytes(self.namespace, "utf-8"))
329+
330+
return InfraObjectProto(
331+
infra_object_class_type="feast.infra.online_stores.datastore.DatastoreTable",
332+
datastore_table=datastore_table_proto,
333+
)
334+
335+
@staticmethod
336+
def from_proto(infra_object_proto: InfraObjectProto) -> Any:
337+
datastore_table = DatastoreTable(
338+
project=infra_object_proto.datastore_table.project,
339+
name=infra_object_proto.datastore_table.name,
340+
)
341+
342+
if infra_object_proto.datastore_table.HasField("project_id"):
343+
datastore_table.project_id = (
344+
infra_object_proto.datastore_table.project_id.SerializeToString()
345+
).decode("utf-8")
346+
if infra_object_proto.datastore_table.HasField("namespace"):
347+
datastore_table.namespace = (
348+
infra_object_proto.datastore_table.namespace.SerializeToString()
349+
).decode("utf-8")
350+
351+
return datastore_table
352+
353+
def update(self):
354+
key = self.client.key("Project", self.project, "Table", self.name)
355+
entity = datastore.Entity(
356+
key=key, exclude_from_indexes=("created_ts", "event_ts", "values")
357+
)
358+
entity.update({"created_ts": datetime.utcnow()})
359+
self.client.put(entity)
360+
361+
def teardown(self):
362+
key = self.client.key("Project", self.project, "Table", self.name)
363+
_delete_all_values(self.client, key)
364+
365+
# Delete the table metadata datastore entity
366+
self.client.delete(key)

0 commit comments

Comments
 (0)