From 1f17caacdaa573d08dbf8dc68b20e73a187ed8a4 Mon Sep 17 00:00:00 2001 From: Hao Xu Date: Fri, 27 Sep 2024 17:22:46 -0700 Subject: [PATCH] fix: Fix online pg import (#4581) * fix postgres import issue Signed-off-by: cmuhao * fix postgres import issue Signed-off-by: cmuhao * fix postgres import issue Signed-off-by: cmuhao * fix postgres import issue Signed-off-by: cmuhao * fix lint Signed-off-by: cmuhao --------- Signed-off-by: cmuhao --- .../feast/infra/online_stores/contrib/postgres.py | 4 +--- .../contrib/singlestore_online_store/singlestore.py | 10 ++-------- sdk/python/feast/infra/online_stores/helpers.py | 8 ++++++++ 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py index 8125da33be..036c0b6b92 100644 --- a/sdk/python/feast/infra/online_stores/contrib/postgres.py +++ b/sdk/python/feast/infra/online_stores/contrib/postgres.py @@ -23,9 +23,7 @@ from feast import Entity from feast.feature_view import FeatureView from feast.infra.key_encoding_utils import get_list_val_str, serialize_entity_key -from feast.infra.online_stores.contrib.singlestore_online_store.singlestore import ( - _to_naive_utc, -) +from feast.infra.online_stores.helpers import _to_naive_utc from feast.infra.online_stores.online_store import OnlineStore from feast.infra.utils.postgres.connection_utils import ( _get_conn, diff --git a/sdk/python/feast/infra/online_stores/contrib/singlestore_online_store/singlestore.py b/sdk/python/feast/infra/online_stores/contrib/singlestore_online_store/singlestore.py index 3e921afcea..d78289c867 100644 --- a/sdk/python/feast/infra/online_stores/contrib/singlestore_online_store/singlestore.py +++ b/sdk/python/feast/infra/online_stores/contrib/singlestore_online_store/singlestore.py @@ -1,7 +1,7 @@ from __future__ import absolute_import from collections import defaultdict -from datetime import datetime, timezone +from datetime import datetime from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple import singlestoredb @@ -11,6 +11,7 @@ from feast import Entity, FeatureView, RepoConfig from feast.infra.key_encoding_utils import serialize_entity_key +from feast.infra.online_stores.helpers import _to_naive_utc from feast.infra.online_stores.online_store import OnlineStore from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto @@ -225,10 +226,3 @@ def _drop_table_and_index(cur: Cursor, project: str, table: FeatureView) -> None def _table_id(project: str, table: FeatureView) -> str: return f"{project}_{table.name}" - - -def _to_naive_utc(ts: datetime) -> datetime: - if ts.tzinfo is None: - return ts - else: - return ts.astimezone(tz=timezone.utc).replace(tzinfo=None) diff --git a/sdk/python/feast/infra/online_stores/helpers.py b/sdk/python/feast/infra/online_stores/helpers.py index 0e2fdb3500..409a1eb212 100644 --- a/sdk/python/feast/infra/online_stores/helpers.py +++ b/sdk/python/feast/infra/online_stores/helpers.py @@ -1,4 +1,5 @@ import struct +from datetime import datetime, timezone from typing import Any, List import mmh3 @@ -62,3 +63,10 @@ def compute_entity_id( entity_key_serialization_version=entity_key_serialization_version, ) ).hex() + + +def _to_naive_utc(ts: datetime) -> datetime: + if ts.tzinfo is None: + return ts + else: + return ts.astimezone(tz=timezone.utc).replace(tzinfo=None)