Skip to content

Commit b7de4b2

Browse files
Lint: Raise exceptions if cursor returned no columns or rows
Add log statement Lint: Fix _to_arrow_internal Lint: Fix _get_entity_df_event_timestamp_range Update exception Use ZeroColumnQueryResult Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>
1 parent d0f77ca commit b7de4b2

File tree

3 files changed

+32
-11
lines changed

3 files changed

+32
-11
lines changed

sdk/python/feast/errors.py

+10
Original file line numberDiff line numberDiff line change
@@ -401,3 +401,13 @@ def __init__(self, input_dict: dict):
401401
super().__init__(
402402
f"Failed to serialize the provided dictionary into a pandas DataFrame: {input_dict.keys()}"
403403
)
404+
405+
406+
class ZeroRowsQueryResult(Exception):
407+
def __init__(self, query: str):
408+
super().__init__(f"This query returned zero rows:\n{query}")
409+
410+
411+
class ZeroColumnQueryResult(Exception):
412+
def __init__(self, query: str):
413+
super().__init__(f"This query returned zero columns:\n{query}")

sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py

+15-9
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import contextlib
2+
import logging
23
from dataclasses import asdict
34
from datetime import datetime
45
from typing import (
@@ -23,7 +24,7 @@
2324
from pytz import utc
2425

2526
from feast.data_source import DataSource
26-
from feast.errors import InvalidEntityType
27+
from feast.errors import InvalidEntityType, ZeroColumnQueryResult, ZeroRowsQueryResult
2728
from feast.feature_view import DUMMY_ENTITY_ID, DUMMY_ENTITY_VAL, FeatureView
2829
from feast.infra.offline_stores import offline_utils
2930
from feast.infra.offline_stores.contrib.postgres_offline_store.postgres_source import (
@@ -276,6 +277,8 @@ def _to_arrow_internal(self, timeout: Optional[int] = None) -> pa.Table:
276277
with _get_conn(self.config.offline_store) as conn, conn.cursor() as cur:
277278
conn.read_only = True
278279
cur.execute(query)
280+
if not cur.description:
281+
raise ZeroColumnQueryResult(query)
279282
fields = [
280283
(c.name, pg_type_code_to_arrow(c.type_code))
281284
for c in cur.description
@@ -331,16 +334,19 @@ def _get_entity_df_event_timestamp_range(
331334
entity_df_event_timestamp.max().to_pydatetime(),
332335
)
333336
elif isinstance(entity_df, str):
334-
# If the entity_df is a string (SQL query), determine range
335-
# from table
337+
# If the entity_df is a string (SQL query), determine range from table
336338
with _get_conn(config.offline_store) as conn, conn.cursor() as cur:
337-
(
338-
cur.execute(
339-
f"SELECT MIN({entity_df_event_timestamp_col}) AS min, MAX({entity_df_event_timestamp_col}) AS max FROM ({entity_df}) as tmp_alias"
340-
),
341-
)
339+
query = f"""
340+
SELECT
341+
MIN({entity_df_event_timestamp_col}) AS min,
342+
MAX({entity_df_event_timestamp_col}) AS max
343+
FROM ({entity_df}) AS tmp_alias
344+
"""
345+
cur.execute(query)
342346
res = cur.fetchone()
343-
entity_df_event_timestamp_range = (res[0], res[1])
347+
if not res:
348+
raise ZeroRowsQueryResult(query)
349+
entity_df_event_timestamp_range = (res[0], res[1])
344350
else:
345351
raise InvalidEntityType(type(entity_df))
346352

sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres_source.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import json
2+
import logging
23
from typing import Callable, Dict, Iterable, Optional, Tuple
34

45
from typeguard import typechecked
56

67
from feast.data_source import DataSource
7-
from feast.errors import DataSourceNoNameException
8+
from feast.errors import DataSourceNoNameException, ZeroColumnQueryResult
89
from feast.infra.utils.postgres.connection_utils import _get_conn
910
from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
1011
from feast.protos.feast.core.SavedDataset_pb2 import (
@@ -111,7 +112,11 @@ def get_table_column_names_and_types(
111112
self, config: RepoConfig
112113
) -> Iterable[Tuple[str, str]]:
113114
with _get_conn(config.offline_store) as conn, conn.cursor() as cur:
114-
cur.execute(f"SELECT * FROM {self.get_table_query_string()} AS sub LIMIT 0")
115+
query = f"SELECT * FROM {self.get_table_query_string()} AS sub LIMIT 0"
116+
cur.execute(query)
117+
if not cur.description:
118+
raise ZeroColumnQueryResult(query)
119+
115120
return (
116121
(c.name, pg_type_code_to_pg_type(c.type_code)) for c in cur.description
117122
)

0 commit comments

Comments
 (0)