Skip to content

Commit 92ffbb9

Browse files
jialsjials
authored
fix: Filter empty SQL commands at execute_snowflake_statement call sites (#6249)
* fix: Handle empty query in execute_snowflake_statement Signed-off-by: Jia Le <[email protected]> * fix: Filter empty SQL commands at call sites instead of silently ignoring in execute_snowflake_statement Signed-off-by: Jia Le <[email protected]> * fix: Add missing semicolons in snowflake_table_deletion.sql Signed-off-by: Jia Le <[email protected]> * fix: Add missing PROJECTS table drop to snowflake_table_deletion.sql Signed-off-by: Jia Le <[email protected]> --------- Signed-off-by: Jia Le <[email protected]>
1 parent 9107a43 commit 92ffbb9

4 files changed

Lines changed: 53 additions & 10 deletions

File tree

‎.secrets.baseline‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎sdk/python/feast/infra/registry/snowflake.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -139,8 +139,9 @@ def __init__(
139139
with GetSnowflakeConnection(self.registry_config) as conn:
140140
sql_function_file = f"{os.path.dirname(feast.__file__)}/infra/utils/snowflake/registry/snowflake_table_creation.sql"
141141
with open(sql_function_file, "r") as file:
142-
sql_file = file.read()
143-
sql_cmds = sql_file.split(";")
142+
sql_cmds = [
143+
cmd.strip() for cmd in file.read().split(";") if cmd.strip()
144+
]
144145
for command in sql_cmds:
145146
query = command.replace("REGISTRY_PATH", f"{self.registry_path}")
146147
execute_snowflake_statement(conn, query)
@@ -224,9 +225,10 @@ def teardown(self):
224225
with GetSnowflakeConnection(self.registry_config) as conn:
225226
sql_function_file = f"{os.path.dirname(feast.__file__)}/infra/utils/snowflake/registry/snowflake_table_deletion.sql"
226227
with open(sql_function_file, "r") as file:
227-
sqlFile = file.read()
228-
sqlCommands = sqlFile.split(";")
229-
for command in sqlCommands:
228+
sql_cmds = [
229+
cmd.strip() for cmd in file.read().split(";") if cmd.strip()
230+
]
231+
for command in sql_cmds:
230232
query = command.replace("REGISTRY_PATH", f"{self.registry_path}")
231233
execute_snowflake_statement(conn, query)
232234

‎sdk/python/feast/infra/utils/snowflake/registry/snowflake_table_deletion.sql‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
DROP TABLE IF EXISTS REGISTRY_PATH."PROJECTS";
2+
13
DROP TABLE IF EXISTS REGISTRY_PATH."DATA_SOURCES";
24

35
DROP TABLE IF EXISTS REGISTRY_PATH."ENTITIES";
@@ -16,6 +18,6 @@ DROP TABLE IF EXISTS REGISTRY_PATH."SAVED_DATASETS";
1618

1719
DROP TABLE IF EXISTS REGISTRY_PATH."STREAM_FEATURE_VIEWS";
1820

19-
DROP TABLE IF EXISTS REGISTRY_PATH."VALIDATION_REFERENCES"
21+
DROP TABLE IF EXISTS REGISTRY_PATH."VALIDATION_REFERENCES";
2022

21-
DROP TABLE IF EXISTS REGISTRY_PATH."PERMISSIONS"
23+
DROP TABLE IF EXISTS REGISTRY_PATH."PERMISSIONS";

‎sdk/python/tests/unit/infra/utils/snowflake/test_snowflake_utils.py‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
import tempfile
22
from typing import Optional
3+
from unittest.mock import MagicMock
34

45
import pytest
56
from cryptography.hazmat.primitives import serialization
67
from cryptography.hazmat.primitives.asymmetric import rsa
78

8-
from feast.infra.utils.snowflake.snowflake_utils import parse_private_key_path
9+
from feast.infra.utils.snowflake.snowflake_utils import (
10+
execute_snowflake_statement,
11+
parse_private_key_path,
12+
)
913

1014
PRIVATE_KEY_PASSPHRASE = "test"
1115

@@ -69,3 +73,38 @@ def test_parse_private_key_path_key_path_encrypted(encrypted_private_key):
6973
f.name,
7074
None,
7175
)
76+
77+
78+
class TestExecuteSnowflakeStatement:
79+
def test_empty_query_is_passed_through_to_execute(self):
80+
mock_conn = MagicMock()
81+
mock_cursor = MagicMock()
82+
mock_executed_cursor = MagicMock()
83+
mock_conn.cursor.return_value = mock_cursor
84+
mock_cursor.execute.return_value = mock_executed_cursor
85+
86+
result = execute_snowflake_statement(mock_conn, "")
87+
88+
assert result is mock_executed_cursor
89+
mock_cursor.execute.assert_called_once_with("")
90+
91+
def test_valid_query_executes_and_returns_cursor(self):
92+
mock_conn = MagicMock()
93+
mock_cursor = MagicMock()
94+
mock_executed_cursor = MagicMock()
95+
mock_conn.cursor.return_value = mock_cursor
96+
mock_cursor.execute.return_value = mock_executed_cursor
97+
98+
result = execute_snowflake_statement(mock_conn, "SELECT 1")
99+
100+
assert result is mock_executed_cursor
101+
mock_cursor.execute.assert_called_once_with("SELECT 1")
102+
103+
def test_valid_query_raises_on_none_cursor(self):
104+
mock_conn = MagicMock()
105+
mock_cursor = MagicMock()
106+
mock_conn.cursor.return_value = mock_cursor
107+
mock_cursor.execute.return_value = None
108+
109+
with pytest.raises(Exception, match="Snowflake query failed"):
110+
execute_snowflake_statement(mock_conn, "SELECT 1")

0 commit comments

Comments
 (0)