fix: Fix online pg import (feast-dev#4581)

HaoXuAI · web-flow · commit 1f17caacdaa5 · 2024-09-27T17:22:46.000-07:00
* fix postgres import issue

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;

* fix postgres import issue

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;

* fix postgres import issue

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;

* fix postgres import issue

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;

* fix lint

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;

---------

Signed-off-by: cmuhao &lt;sduxuhao@gmail.com&gt;
diff --git 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
@@ -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
@@ -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)
