Skip to content

Commit 0d04574

Browse files
treysptreysp
committed
Fix tests and snowflake regex
1 parent 2fb3daf commit 0d04574

File tree

5 files changed

+57
-38
lines changed

5 files changed

+57
-38
lines changed

‎sqlmesh/core/engine_adapter/snowflake.py‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -676,23 +676,32 @@ def _record_execution_stats(
676676
They neither report the sentinel value -1 nor do they report 0 rows. Instead, they return a single data row
677677
containing the string "Table <table_name> successfully created." and a row count of 1.
678678
679-
We do not want to record the row count of 1 for CTAS operations, so we check for that data pattern and return
680-
early if it is detected.
681-
682-
Regex explanation - Snowflake identifiers may be:
683-
- An unquoted contiguous set of [a-zA-Z0-9_$] characters
684-
- A double-quoted string that may contain spaces and nested double-quotes represented by `""`
685-
- Example: " my ""table"" name "
686-
- Pattern: "(?:[^"]|"")+"
687-
- ?: is a non-capturing group
688-
- [^"] matches any single character except a double-quote
689-
- "" matches two sequential double-quotes
679+
We do not want to record the incorrect row count of 1, so we check whether:
680+
- There is exactly one row to fetch (in general, DML operations should return no rows to fetch from the cursor)
681+
- That row contains the table successfully created string
682+
683+
If so, we return early and do not record the row count.
690684
"""
691685
if rowcount == 1:
692686
results = self.cursor.fetchall()
693687
if results and len(results) == 1:
688+
try:
689+
results_str = str(results[0][0])
690+
except (ValueError, TypeError):
691+
return
692+
693+
# Snowflake identifiers may be:
694+
# - An unquoted contiguous set of [a-zA-Z0-9_$] characters
695+
# - A double-quoted string that may contain spaces and nested double-quotes represented by `""`. Example: " my ""table"" name "
696+
# - Regex:
697+
# - [a-zA-Z0-9_$]+ matches one or more character in the set
698+
# - "(?:[^"]|"")+" matches a double-quoted string that may contain spaces and nested double-quotes
699+
# - ?: non-capturing group
700+
# - [^"] matches any single character except a double-quote
701+
# - | or
702+
# - "" matches two sequential double-quotes
694703
is_ctas = re.match(
695-
r'Table ([a-zA-Z0-9_$]+|"(?:[^"]|"")+") successfully created\.', results[0][0]
704+
r'Table ([a-zA-Z0-9_$]+|"(?:[^"]|"")+") successfully created\.', results_str
696705
)
697706
if is_ctas:
698707
return

‎sqlmesh/core/snapshot/execution_tracker.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,15 +41,15 @@ def __post_init__(self) -> None:
4141
def add_execution(
4242
self, sql: str, row_count: t.Optional[int], bytes_processed: t.Optional[int]
4343
) -> None:
44-
if row_count is not None and row_count >= 0:
44+
if row_count is not None:
4545
if self.stats.total_rows_processed is None:
4646
self.stats.total_rows_processed = row_count
4747
else:
4848
self.stats.total_rows_processed += row_count
4949

5050
# conditional on row_count because we should only count bytes corresponding to
5151
# DML actions whose rows were captured
52-
if bytes_processed is not None and bytes_processed >= 0:
52+
if bytes_processed is not None:
5353
if self.stats.total_bytes_processed is None:
5454
self.stats.total_bytes_processed = bytes_processed
5555
else:

‎tests/core/test_snapshot_evaluator.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -606,8 +606,10 @@ def test_evaluate_materialized_view_with_partitioned_by_cluster_by(
606606

607607
execute_mock.assert_has_calls(
608608
[
609+
call("CREATE SCHEMA IF NOT EXISTS `sqlmesh__test_schema`", False),
609610
call(
610-
f"CREATE MATERIALIZED VIEW `sqlmesh__test_schema`.`test_schema__test_model__{snapshot.version}` PARTITION BY `a` CLUSTER BY `b` AS SELECT `a` AS `a`, `b` AS `b` FROM `tbl` AS `tbl`"
611+
f"CREATE MATERIALIZED VIEW `sqlmesh__test_schema`.`test_schema__test_model__{snapshot.version}` PARTITION BY `a` CLUSTER BY `b` AS SELECT `a` AS `a`, `b` AS `b` FROM `tbl` AS `tbl`",
612+
False,
611613
),
612614
]
613615
)

‎tests/core/test_table_diff.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -360,11 +360,11 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
360360
temp_schema="sqlmesh_temp_test",
361361
)
362362

363-
spy_execute.assert_any_call(query_sql)
364-
spy_execute.assert_any_call(summary_query_sql)
365-
spy_execute.assert_any_call(compare_sql)
366-
spy_execute.assert_any_call(sample_query_sql)
367-
spy_execute.assert_any_call(drop_sql)
363+
spy_execute.assert_any_call(query_sql, False)
364+
spy_execute.assert_any_call(summary_query_sql, False)
365+
spy_execute.assert_any_call(compare_sql, False)
366+
spy_execute.assert_any_call(sample_query_sql, False)
367+
spy_execute.assert_any_call(drop_sql, False)
368368

369369
spy_execute.reset_mock()
370370

@@ -378,7 +378,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
378378
)
379379

380380
query_sql_where = 'CREATE TABLE IF NOT EXISTS "memory"."sqlmesh_temp"."__temp_diff_abcdefgh" AS WITH "__source" AS (SELECT "s"."key", "s"."value", "s"."key" AS "__sqlmesh_join_key" FROM "table_diff_source" AS "s" WHERE "s"."key" = 2), "__target" AS (SELECT "t"."key", "t"."value", "t"."key" AS "__sqlmesh_join_key" FROM "table_diff_target" AS "t" WHERE "t"."key" = 2), "__stats" AS (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "s"."__sqlmesh_join_key" AS "s____sqlmesh_join_key", "t"."key" AS "t__key", "t"."value" AS "t__value", "t"."__sqlmesh_join_key" AS "t____sqlmesh_join_key", CASE WHEN NOT "s"."__sqlmesh_join_key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."__sqlmesh_join_key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key" THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "__source" AS "s" FULL JOIN "__target" AS "t" ON "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key") SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM "__stats"'
381-
spy_execute.assert_any_call(query_sql_where)
381+
spy_execute.assert_any_call(query_sql_where, False)
382382

383383

384384
def test_tables_and_grain_inferred_from_model(sushi_context_fixed_date: Context):

‎tests/core/test_test.py‎

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -874,7 +874,8 @@ def test_partially_inferred_schemas(sushi_context: Context, mocker: MockerFixtur
874874
'CAST("s" AS STRUCT("d" DATE)) AS "s", '
875875
'CAST("a" AS INT) AS "a", '
876876
'CAST("b" AS TEXT) AS "b" '
877-
"""FROM (VALUES ({'d': CAST('2020-01-01' AS DATE)}, 1, 'bla')) AS "t"("s", "a", "b")"""
877+
"""FROM (VALUES ({'d': CAST('2020-01-01' AS DATE)}, 1, 'bla')) AS "t"("s", "a", "b")""",
878+
False,
878879
)
879880

880881

@@ -1329,14 +1330,15 @@ def test_freeze_time(mocker: MockerFixture) -> None:
13291330

13301331
spy_execute.assert_has_calls(
13311332
[
1332-
call('CREATE SCHEMA IF NOT EXISTS "memory"."sqlmesh_test_jzngz56a"'),
1333+
call('CREATE SCHEMA IF NOT EXISTS "memory"."sqlmesh_test_jzngz56a"', False),
13331334
call(
13341335
"SELECT "
13351336
"""CAST('2023-01-01 12:05:03+00:00' AS DATE) AS "cur_date", """
13361337
"""CAST('2023-01-01 12:05:03+00:00' AS TIME) AS "cur_time", """
1337-
'''CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMP) AS "cur_timestamp"'''
1338+
'''CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMP) AS "cur_timestamp"''',
1339+
False,
13381340
),
1339-
call('DROP SCHEMA IF EXISTS "memory"."sqlmesh_test_jzngz56a" CASCADE'),
1341+
call('DROP SCHEMA IF EXISTS "memory"."sqlmesh_test_jzngz56a" CASCADE', False),
13401342
]
13411343
)
13421344

@@ -1361,7 +1363,12 @@ def test_freeze_time(mocker: MockerFixture) -> None:
13611363
_check_successful_or_raise(test.run())
13621364

13631365
spy_execute.assert_has_calls(
1364-
[call('''SELECT CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMPTZ) AS "cur_timestamp"''')]
1366+
[
1367+
call(
1368+
'''SELECT CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMPTZ) AS "cur_timestamp"''',
1369+
False,
1370+
)
1371+
]
13651372
)
13661373

13671374
@model("py_model", columns={"ts1": "timestamptz", "ts2": "timestamptz"})
@@ -1496,7 +1503,7 @@ def test_gateway(copy_to_temp_path: t.Callable, mocker: MockerFixture) -> None:
14961503
'AS "t"("id", "customer_id", "waiter_id", "start_ts", "end_ts", "event_date")'
14971504
)
14981505
test_adapter = t.cast(ModelTest, result.successes[0]).engine_adapter
1499-
assert call(test_adapter, expected_view_sql) in spy_execute.mock_calls
1506+
assert call(test_adapter, expected_view_sql, False) in spy_execute.mock_calls
15001507

15011508
_check_successful_or_raise(context.test())
15021509

@@ -1621,7 +1628,8 @@ def test_generate_input_data_using_sql(mocker: MockerFixture, tmp_path: Path) ->
16211628

16221629
spy_execute.assert_any_call(
16231630
'CREATE OR REPLACE VIEW "memory"."sqlmesh_test_jzngz56a"."foo" AS '
1624-
'''SELECT {'x': 1, 'n': {'y': 2}} AS "struct_value"'''
1631+
'''SELECT {'x': 1, 'n': {'y': 2}} AS "struct_value"''',
1632+
False,
16251633
)
16261634

16271635
with pytest.raises(
@@ -1817,9 +1825,9 @@ def test_custom_testing_schema(mocker: MockerFixture) -> None:
18171825

18181826
spy_execute.assert_has_calls(
18191827
[
1820-
call('CREATE SCHEMA IF NOT EXISTS "memory"."my_schema"'),
1821-
call('SELECT 1 AS "a"'),
1822-
call('DROP SCHEMA IF EXISTS "memory"."my_schema" CASCADE'),
1828+
call('CREATE SCHEMA IF NOT EXISTS "memory"."my_schema"', False),
1829+
call('SELECT 1 AS "a"', False),
1830+
call('DROP SCHEMA IF EXISTS "memory"."my_schema" CASCADE', False),
18231831
]
18241832
)
18251833

@@ -1845,9 +1853,9 @@ def test_pretty_query(mocker: MockerFixture) -> None:
18451853
_check_successful_or_raise(test.run())
18461854
spy_execute.assert_has_calls(
18471855
[
1848-
call('CREATE SCHEMA IF NOT EXISTS "memory"."my_schema"'),
1849-
call('SELECT\n 1 AS "a"'),
1850-
call('DROP SCHEMA IF EXISTS "memory"."my_schema" CASCADE'),
1856+
call('CREATE SCHEMA IF NOT EXISTS "memory"."my_schema"', False),
1857+
call('SELECT\n 1 AS "a"', False),
1858+
call('DROP SCHEMA IF EXISTS "memory"."my_schema" CASCADE', False),
18511859
]
18521860
)
18531861

@@ -2950,7 +2958,7 @@ def test_parameterized_name_sql_model() -> None:
29502958
outputs:
29512959
query:
29522960
- id: 1
2953-
name: foo
2961+
name: foo
29542962
""",
29552963
variables=variables,
29562964
),
@@ -2999,7 +3007,7 @@ def execute(
29993007
outputs:
30003008
query:
30013009
- id: 1
3002-
name: foo
3010+
name: foo
30033011
""",
30043012
variables=variables,
30053013
),
@@ -3049,7 +3057,7 @@ def test_parameterized_name_self_referential_model():
30493057
v: int
30503058
outputs:
30513059
query:
3052-
- v: 1
3060+
- v: 1
30533061
""",
30543062
variables=variables,
30553063
),
@@ -3171,7 +3179,7 @@ def execute(
31713179
- id: 5
31723180
outputs:
31733181
query:
3174-
- id: 8
3182+
- id: 8
31753183
""",
31763184
variables=variables,
31773185
),

0 commit comments

Comments
 (0)