Skip to content

Commit 7a02a31

Browse files
sfc-gh-alingsfc-gh-aling
authored
SNOW-3654850: escape identifier quotes in all dialect generate_select_query methods (#4272)
1 parent bf36bae commit 7a02a31

13 files changed

Lines changed: 497 additions & 55 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#### Bug Fixes
88

99
- Fixed a bug where stage paths and file format names that contain single quotes were not consistently escaped when generating SQL, which could produce malformed statements. This affects `INFER_SCHEMA` (used by `DataFrameReader.csv`/`json`/`parquet`/`orc`/`avro`) and `COPY FILES` (used by `FileOperation.copy_files`).
10+
- Fixed a bug where column names containing quote characters returned by an external database were not correctly escaped when generating the `SELECT` query for `DataFrameReader.dbapi`, which could produce malformed SQL. Embedded quote characters in identifiers are now doubled (backticks for Databricks/MySQL, double quotes for Oracle/PostgreSQL/SQL Server).
1011
- Fixed a bug where the destination passed to `DataFrameWriter.copy_into_location` (and `csv`/`json`/`parquet`/`save`) was embedded into the generated `COPY INTO` statement without quoting, which could produce malformed SQL for locations containing single quotes. The location is now consistently quoted and escaped, and a string that merely starts and ends with a single quote but contains unescaped interior quotes is no longer treated as an already-quoted literal; it is fully escaped so it stays a single SQL string literal.
1112
- Fixed a bug where UDF default argument values reconstructed from a source file in `register_from_file` were evaluated with `eval()`; they are now evaluated only against the documented set of supported default-value types, and unsupported expressions are ignored.
1213
- Fixed a bug where `object_name`, `object_domain`, or `object_version` values containing single quotes or backslashes in `session.lineage.trace()` caused incorrect SQL generation. These values are now properly escaped before being embedded in the `SYSTEM$DGQL` call.

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/base_dialect.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,11 @@
1010

1111

1212
class BaseDialect:
13+
@staticmethod
14+
def _quote_backtick(name: str) -> str:
15+
"""Escape a backtick-quoted identifier: double any embedded backticks."""
16+
return "`" + name.replace("`", "``") + "`"
17+
1318
@staticmethod
1419
def generate_select_query(
1520
table_or_query: str,
@@ -20,12 +25,11 @@ def generate_select_query(
2025
) -> str:
2126
cols = []
2227
for raw_field in raw_schema:
28+
quoted_name = BaseDialect._quote_backtick(raw_field[0])
2329
if is_query:
24-
cols.append(
25-
f"""{query_input_alias}.`{raw_field[0]}` AS {raw_field[0]}"""
26-
)
30+
cols.append(f"""{query_input_alias}.{quoted_name} AS {quoted_name}""")
2731
else:
28-
cols.append(raw_field[0])
32+
cols.append(quoted_name)
2933

3034
return QUERY_TEMPLATE.format(
3135
cols=", ".join(cols),

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/databricks_dialect.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,21 @@ def generate_select_query(
2121
) -> str:
2222
cols = []
2323
for field, raw_field in zip(schema.fields, raw_schema):
24+
quoted_name = self._quote_backtick(raw_field[0])
2425
field_name = (
25-
f"{query_input_alias}.`{raw_field[0]}`"
26-
if is_query
27-
else f"`{raw_field[0]}`"
26+
f"{query_input_alias}.{quoted_name}" if is_query else quoted_name
2827
)
28+
alias = quoted_name
2929
# databricks-sql-connector returns list of tuples for MapType
3030
# here we push down to-dict conversion to Databricks
3131
if isinstance(field.datatype, MapType):
32-
cols.append(f"""TO_JSON({field_name}) AS {raw_field[0]}""")
32+
cols.append(f"""TO_JSON({field_name}) AS {alias}""")
3333
elif isinstance(field.datatype, BinaryType):
34-
cols.append(f"""HEX({field_name}) AS {raw_field[0]}""")
34+
cols.append(f"""HEX({field_name}) AS {alias}""")
3535
else:
36-
cols.append(
37-
f"{field_name} AS {raw_field[0]}"
38-
) if is_query else cols.append(field_name)
36+
cols.append(f"{field_name} AS {alias}") if is_query else cols.append(
37+
field_name
38+
)
3939
return QUERY_TEMPLATE.format(
4040
cols=", ".join(cols),
4141
table_or_query=f"({table_or_query})" if is_query else table_or_query,

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/mysql_dialect.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,22 @@ def generate_select_query(
2424
) -> str:
2525
cols = []
2626
for field, raw_field in zip(schema.fields, raw_schema):
27+
quoted_name = self._quote_backtick(raw_field[0])
2728
field_name = (
28-
f"{query_input_alias}.`{raw_field[0]}`"
29-
if is_query
30-
else f"`{raw_field[0]}`"
29+
f"{query_input_alias}.{quoted_name}" if is_query else quoted_name
3130
)
31+
alias = quoted_name
3232
if isinstance(field.datatype, TimeType):
33-
cols.append(f"""CAST({field_name} AS CHAR) AS {raw_field[0]}""")
33+
cols.append(f"""CAST({field_name} AS CHAR) AS {alias}""")
3434
elif (
3535
isinstance(field.datatype, BinaryType)
3636
or raw_field[1] == PymysqlTypeCode.BIT
3737
):
38-
cols.append(f"""HEX({field_name}) AS {raw_field[0]}""")
38+
cols.append(f"""HEX({field_name}) AS {alias}""")
3939
else:
40-
cols.append(
41-
f"{field_name} AS {raw_field[0]}"
42-
) if is_query else cols.append(field_name)
40+
cols.append(f"{field_name} AS {alias}") if is_query else cols.append(
41+
field_name
42+
)
4343

4444
return QUERY_TEMPLATE.format(
4545
cols=", ".join(cols),

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/oracledb_dialect.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,29 +22,29 @@ def generate_select_query(
2222
) -> str:
2323
cols = []
2424
for field, raw_field in zip(schema.fields, raw_schema):
25+
quoted_name = quote_name(raw_field[0], keep_case=True)
2526
field_name = (
26-
f"{query_input_alias}.{quote_name(raw_field[0], keep_case=True)}"
27-
if is_query
28-
else f"{quote_name(raw_field[0], keep_case=True)}"
27+
f"{query_input_alias}.{quoted_name}" if is_query else f"{quoted_name}"
2928
)
29+
alias = quoted_name
3030
if (
3131
isinstance(field.datatype, TimestampType)
3232
and field.datatype.tz == TimestampTimeZone.TZ
3333
):
3434
cols.append(
35-
f"""TO_CHAR({field_name}, 'YYYY-MM-DD HH24:MI:SS.FF9 TZHTZM') AS {raw_field[0]}"""
35+
f"""TO_CHAR({field_name}, 'YYYY-MM-DD HH24:MI:SS.FF9 TZHTZM') AS {alias}"""
3636
)
3737
elif (
3838
isinstance(field.datatype, TimestampType)
3939
and field.datatype.tz == TimestampTimeZone.LTZ
4040
):
4141
cols.append(
42-
f"""TO_CHAR({field_name} AT TIME ZONE SESSIONTIMEZONE, 'YYYY-MM-DD HH24:MI:SS.FF9 TZHTZM') AS {raw_field[0]}"""
42+
f"""TO_CHAR({field_name} AT TIME ZONE SESSIONTIMEZONE, 'YYYY-MM-DD HH24:MI:SS.FF9 TZHTZM') AS {alias}"""
4343
)
4444
else:
45-
cols.append(
46-
f"{field_name} AS {raw_field[0]}"
47-
) if is_query else cols.append(field_name)
45+
cols.append(f"{field_name} AS {alias}") if is_query else cols.append(
46+
field_name
47+
)
4848
return QUERY_TEMPLATE.format(
4949
cols=", ".join(cols),
5050
table_or_query=f"({table_or_query})" if is_query else table_or_query,

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/postgresql_dialect.py‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,11 @@
1414

1515

1616
class PostgresDialect(BaseDialect):
17+
@staticmethod
18+
def _quote_ident(name: str) -> str:
19+
"""Escape a PostgreSQL identifier: double any embedded double-quotes."""
20+
return '"' + name.replace('"', '""') + '"'
21+
1722
@staticmethod
1823
def generate_select_query(
1924
table_or_query: str,
@@ -27,30 +32,30 @@ def generate_select_query(
2732
# databricks-sql-connector returns list of tuples for MapType
2833
# here we push down to-dict conversion to Databricks
2934
type_code = raw_field[1]
35+
quoted_name = PostgresDialect._quote_ident(raw_field[0])
3036
field_name = (
31-
f'{query_input_alias}."{raw_field[0]}"'
32-
if is_query
33-
else f'"{raw_field[0]}"'
37+
f"{query_input_alias}.{quoted_name}" if is_query else quoted_name
3438
)
39+
alias = quoted_name
3540
if type_code in (
3641
Psycopg2TypeCode.JSONB.value,
3742
Psycopg2TypeCode.JSON.value,
3843
):
39-
cols.append(f"""TO_JSON({field_name})::TEXT AS {raw_field[0]}""")
44+
cols.append(f"""TO_JSON({field_name})::TEXT AS {alias}""")
4045
elif type_code == Psycopg2TypeCode.CASHOID.value:
4146
cols.append(
42-
f"""CASE WHEN {field_name} IS NULL THEN NULL ELSE FORMAT('"%s"', "{raw_field[0]}"::TEXT) END AS {raw_field[0]}"""
47+
f"""CASE WHEN {field_name} IS NULL THEN NULL ELSE FORMAT('"%s"', {quoted_name}::TEXT) END AS {alias}"""
4348
)
4449
elif type_code == Psycopg2TypeCode.BYTEAOID.value:
45-
cols.append(f"""ENCODE({field_name}, 'HEX') AS {raw_field[0]}""")
50+
cols.append(f"""ENCODE({field_name}, 'HEX') AS {alias}""")
4651
elif type_code == Psycopg2TypeCode.TIMETZOID.value:
47-
cols.append(f"""{field_name}::TIME AS {raw_field[0]}""")
52+
cols.append(f"""{field_name}::TIME AS {alias}""")
4853
elif type_code == Psycopg2TypeCode.INTERVALOID.value:
49-
cols.append(f"""{field_name}::TEXT AS {raw_field[0]}""")
54+
cols.append(f"""{field_name}::TEXT AS {alias}""")
5055
else:
51-
cols.append(
52-
f"{field_name} AS {raw_field[0]}"
53-
) if is_query else cols.append(field_name)
56+
cols.append(f"{field_name} AS {alias}") if is_query else cols.append(
57+
field_name
58+
)
5459

5560
return QUERY_TEMPLATE.format(
5661
cols=", ".join(cols),

‎src/snowflake/snowpark/_internal/data_source/dbms_dialects/sqlserver_dialect.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,12 @@ def generate_select_query(
2222
) -> str:
2323
cols = []
2424
for _field, raw_field in zip(schema.fields, raw_schema):
25+
quoted_name = quote_name(raw_field[0], keep_case=True)
2526
field_name = (
26-
f"{query_input_alias}.{quote_name(raw_field[0], keep_case=True)}"
27-
if is_query
28-
else f"{quote_name(raw_field[0], keep_case=True)}"
27+
f"{query_input_alias}.{quoted_name}" if is_query else f"{quoted_name}"
2928
)
30-
cols.append(f"{field_name} AS {raw_field[0]}") if is_query else cols.append(
29+
alias = quoted_name
30+
cols.append(f"{field_name} AS {alias}") if is_query else cols.append(
3131
field_name
3232
)
3333
return QUERY_TEMPLATE.format(

‎tests/integ/datasource/test_databricks.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,30 @@ def test_double_quoted_column_databricks(session, custom_schema):
172172
assert df.schema == databricks_double_quoted_schema
173173

174174

175+
def test_embedded_quote_alias_has_no_side_effects_databricks(session):
176+
conn = create_databricks_connection()
177+
try:
178+
with conn.cursor() as cur:
179+
cur.execute(f"SELECT COUNT(*) FROM {TEST_TABLE_NAME}")
180+
before_count = cur.fetchone()[0]
181+
182+
embedded_quote_alias = (
183+
f"y` FROM {TEST_TABLE_NAME};DELETE FROM {TEST_TABLE_NAME};--"
184+
)
185+
escaped_alias = embedded_quote_alias.replace("`", "``")
186+
crafted_query = f"SELECT 1 AS `{escaped_alias}`"
187+
188+
df = session.read.dbapi(create_databricks_connection, query=crafted_query)
189+
assert len(df.collect()) == 1
190+
191+
with conn.cursor() as cur:
192+
cur.execute(f"SELECT COUNT(*) FROM {TEST_TABLE_NAME}")
193+
after_count = cur.fetchone()[0]
194+
assert after_count == before_count
195+
finally:
196+
conn.close()
197+
198+
175199
@pytest.mark.parametrize(
176200
"input_type, input_value",
177201
[("table", TEST_TABLE_NAME), ("query", f"(SELECT * FROM {TEST_TABLE_NAME})")],

‎tests/integ/datasource/test_mysql.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,30 @@ def test_double_quoted_column_name_mysql(session, custom_schema):
193193
assert df.schema == mysql_double_quoted_schema
194194

195195

196+
def test_embedded_quote_alias_has_no_side_effects_mysql(session):
197+
conn = create_connection_mysql()
198+
try:
199+
with conn.cursor() as cur:
200+
cur.execute(f"SELECT COUNT(*) FROM {TEST_TABLE_NAME}")
201+
before_count = cur.fetchone()[0]
202+
203+
embedded_quote_alias = (
204+
f"y` FROM {TEST_TABLE_NAME};DELETE FROM {TEST_TABLE_NAME};--"
205+
)
206+
escaped_alias = embedded_quote_alias.replace("`", "``")
207+
crafted_query = f"SELECT 1 AS `{escaped_alias}`"
208+
209+
df = session.read.dbapi(create_connection_mysql, query=crafted_query)
210+
assert len(df.collect()) == 1
211+
212+
with conn.cursor() as cur:
213+
cur.execute(f"SELECT COUNT(*) FROM {TEST_TABLE_NAME}")
214+
after_count = cur.fetchone()[0]
215+
assert after_count == before_count
216+
finally:
217+
conn.close()
218+
219+
196220
@pytest.mark.parametrize(
197221
"data, number_of_columns, expected_result",
198222
[

‎tests/integ/datasource/test_oracledb.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from snowflake.snowpark._internal.data_source.drivers.oracledb_driver import (
1414
output_type_handler,
1515
)
16+
from snowflake.snowpark._internal.data_source import DataSourcePartitioner
1617
from snowflake.snowpark._internal.data_source.drivers import (
1718
OracledbDriver,
1819
)
@@ -240,6 +241,36 @@ def test_double_quoted_column_name_oracledb(session, custom_schema):
240241
assert df.schema == oracledb_double_quoted_schema
241242

242243

244+
def test_embedded_quote_alias_has_no_side_effects_oracledb(session):
245+
conn = create_connection_oracledb()
246+
try:
247+
with conn.cursor() as cur:
248+
cur.execute(f"SELECT COUNT(*) FROM {ORACLEDB_TABLE_NAME}")
249+
before_count = cur.fetchone()[0]
250+
251+
# Use a reserved keyword as a quoted identifier to exercise the
252+
# is_query=True alias quoting path without Oracle's ORA-03001 limitation
253+
# for embedded double-quotes in nested subqueries.
254+
crafted_query = 'SELECT 1 AS "select" FROM dual'
255+
partitioner = DataSourcePartitioner(
256+
create_connection_oracledb, crafted_query, is_query=True
257+
)
258+
partition_query = partitioner.partitions[0]
259+
assert ' AS "select"' in partition_query
260+
261+
with conn.cursor() as cur:
262+
cur.execute(partition_query)
263+
rows = cur.fetchall()
264+
assert len(rows) == 1
265+
266+
with conn.cursor() as cur:
267+
cur.execute(f"SELECT COUNT(*) FROM {ORACLEDB_TABLE_NAME}")
268+
after_count = cur.fetchone()[0]
269+
assert after_count == before_count
270+
finally:
271+
conn.close()
272+
273+
243274
def test_unsupported_type():
244275
invalid_type = OracleDBType("ID", "UNKNOWN", None, None, False)
245276
MockDescription = namedtuple(

0 commit comments

Comments
 (0)