From fad30d48ae5718ce7818a7fb7447c9112b6b3ff0 Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Tue, 14 Dec 2021 15:26:09 -0800 Subject: [PATCH] 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 6d56421edf..3aa55f6be4 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)