diff --git a/CHANGELOG.md b/CHANGELOG.md index 9397856a16..af7c7b2f6e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ #### New Features - Added interval type support for Python UDFs and stored procedures. Use `datetime.timedelta` as the type annotation for day-time interval (`DayTimeIntervalType`) parameters and return values, and `YearMonthInterval` (a type annotation sentinel from `snowflake.snowpark.types`) for year-month interval (`YearMonthIntervalType`) parameters and return values. +- Added support for a `table_properties` key in the `iceberg_config` dictionary of `DataFrameWriter.save_as_table`, which emits a `TABLE_PROPERTIES = ('k'='v', ...)` clause on Iceberg table creation (CREATE / CTAS). #### Bug Fixes diff --git a/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py b/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py index cf8f701570..73ac803a0f 100644 --- a/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py +++ b/src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py @@ -150,6 +150,7 @@ TARGET_FILE_SIZE = " TARGET_FILE_SIZE " CATALOG_SYNC = " CATALOG_SYNC " STORAGE_SERIALIZATION_POLICY = " STORAGE_SERIALIZATION_POLICY " +TABLE_PROPERTIES = " TABLE_PROPERTIES" REG_EXP = " REGEXP " COLLATE = " COLLATE " RESULT_SCAN = " RESULT_SCAN" @@ -334,6 +335,30 @@ def iceberg_partition_clause(partition_exprs: List[str]) -> str: ) +def iceberg_table_properties_clause(iceberg_config: Optional[dict]) -> str: + """Emit ``TABLE_PROPERTIES = ('k'='v', ...)`` from an ``iceberg_config``'s + ``table_properties`` map (keys kept verbatim, keys/values escaped); empty + when absent.""" + if not iceberg_config: + return EMPTY_STRING + normalized = {k.lower(): v for k, v in iceberg_config.items()} + table_properties = normalized.get("table_properties") + if not table_properties: + return EMPTY_STRING + + # Escape via ``escape_quotes_and_backslashes`` (not ``single_quote``): it + # escapes backslashes as well as quotes and never skips escaping, so a + # crafted key/value cannot break out of its literal and inject extra SQL + # (e.g. "'x', 'injected'='surprise'" or "x\\') COPY GRANTS --"). + def _quote(value: object) -> str: + return SINGLE_QUOTE + escape_quotes_and_backslashes(str(value)) + SINGLE_QUOTE + + pairs = COMMA.join(f"{_quote(k)}={_quote(v)}" for k, v in table_properties.items()) + return ( + SPACE + TABLE_PROPERTIES + EQUALS + LEFT_PARENTHESIS + pairs + RIGHT_PARENTHESIS + ) + + def order_by_spec(col_exprs: List[str]) -> str: if not col_exprs: return EMPTY_STRING @@ -1168,13 +1193,14 @@ def create_table_statement( options_statement = get_options_statement(options) partition_by_clause = iceberg_partition_clause(partition_exprs) + table_properties_clause = iceberg_table_properties_clause(iceberg_config) return ( f"{CREATE}{(OR + REPLACE) if replace else EMPTY_STRING}" f" {(get_temp_type_for_object(use_scoped_temp_objects, is_generated) if table_type.lower() in TEMPORARY_STRING_SET else table_type).upper()} " f"{ICEBERG if iceberg_options else EMPTY_STRING}{TABLE}{table_name}{(IF + NOT + EXISTS) if not replace and not error else EMPTY_STRING}" f"{LEFT_PARENTHESIS}{schema}{RIGHT_PARENTHESIS}{partition_by_clause}{cluster_by_clause}" - f"{options_statement}{COPY_GRANTS if copy_grants else EMPTY_STRING}{comment_sql}" + f"{options_statement}{table_properties_clause}{COPY_GRANTS if copy_grants else EMPTY_STRING}{comment_sql}" ) @@ -1259,13 +1285,14 @@ def create_table_as_select_statement( options_statement = get_options_statement(options) partition_by_clause = iceberg_partition_clause(partition_exprs) + table_properties_clause = iceberg_table_properties_clause(iceberg_config) return ( f"{CREATE}{OR + REPLACE if replace else EMPTY_STRING}" f" {(get_temp_type_for_object(use_scoped_temp_objects, is_generated) if table_type.lower() in TEMPORARY_STRING_SET else table_type).upper()} " f"{ICEBERG if iceberg_options else EMPTY_STRING}{TABLE}" f"{IF + NOT + EXISTS if not replace and not error else EMPTY_STRING} " - f"{table_name}{column_definition_sql}{partition_by_clause}{cluster_by_clause}{options_statement}" + f"{table_name}{column_definition_sql}{partition_by_clause}{cluster_by_clause}{options_statement}{table_properties_clause}" f"{COPY_GRANTS if copy_grants else EMPTY_STRING}{comment_sql} {AS}{project_statement([], child)}" ) diff --git a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py index c4bd0a72ca..82a47801d9 100644 --- a/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py +++ b/src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py @@ -1318,6 +1318,8 @@ def save_as_table( catalog_sync: optionally sets the catalog integration configured for Polaris Catalog storage_serialization_policy: specifies the storage serialization policy for the table iceberg_version: Overrides the version of iceberg to use. Defaults to 2 when unset. + table_properties: optional mapping of Iceberg property names to values, emitted + verbatim as a ``TABLE_PROPERTIES = ('k'='v', ...)`` clause on CREATE / CTAS. table_exists: whether the table already exists in the database. Only used for APPEND and TRUNCATE mode. """ diff --git a/src/snowflake/snowpark/dataframe_writer.py b/src/snowflake/snowpark/dataframe_writer.py index bce50edf91..b0e75ef072 100644 --- a/src/snowflake/snowpark/dataframe_writer.py +++ b/src/snowflake/snowpark/dataframe_writer.py @@ -323,6 +323,8 @@ def save_as_table( * storage_serialization_policy: specifies the storage serialization policy for the table * iceberg_version: Overrides the version of iceberg to use. Defaults to 2 when unset. + * table_properties: an optional mapping of Iceberg table property names to values, + emitted verbatim as a ``TABLE_PROPERTIES = ('k'='v', ...)`` clause on table creation. table_exists: Optional parameter to specify if the table is known to exist or not. Set to ``True`` if table exists, ``False`` if it doesn't, or ``None`` (default) for automatic detection. Primarily useful for "append", "truncate", and "overwrite" with overwrite_condition modes to avoid running query for automatic detection. diff --git a/tests/unit/test_analyzer_util_suite.py b/tests/unit/test_analyzer_util_suite.py index a8ef99f404..bbf584ddda 100644 --- a/tests/unit/test_analyzer_util_suite.py +++ b/tests/unit/test_analyzer_util_suite.py @@ -25,6 +25,7 @@ create_table_as_select_statement, create_table_statement, file_operation_statement, + iceberg_table_properties_clause, join_statement, project_statement, table_function_statement, @@ -495,6 +496,89 @@ def test_create_iceberg_table_as_select_statement(): ) +def test_iceberg_table_properties_clause(): + # Absent / empty / no table_properties key -> nothing emitted. + assert iceberg_table_properties_clause(None) == "" + assert iceberg_table_properties_clause({}) == "" + assert iceberg_table_properties_clause({"catalog": "SNOWFLAKE"}) == "" + assert iceberg_table_properties_clause({"table_properties": {}}) == "" + + # Keys are preserved verbatim (not lowercased); the top-level config key + # is matched case-insensitively; keys/values are escaped single-quoted + # literals. + assert ( + iceberg_table_properties_clause( + { + "table_properties": { + "write.metadata.compression-codec": "gzip", + "custom": "it's ok", + } + } + ) + == " TABLE_PROPERTIES = ('write.metadata.compression-codec'='gzip', 'custom'='it''s ok')" + ) + assert ( + iceberg_table_properties_clause({"TABLE_PROPERTIES": {"a": "b"}}) + == " TABLE_PROPERTIES = ('a'='b')" + ) + + # Values/keys cannot inject additional SQL: a crafted value that starts and + # ends with a quote is still fully escaped into a single literal (it must + # not add a second property to the clause). + assert ( + iceberg_table_properties_clause( + {"table_properties": {"k": "'x', 'injected'='surprise'"}} + ) + == " TABLE_PROPERTIES = ('k'='''x'', ''injected''=''surprise''')" + ) + assert ( + iceberg_table_properties_clause({"table_properties": {"k')=(": "v"}}) + == " TABLE_PROPERTIES = ('k'')=('='v')" + ) + # Backslash must also be escaped: a bare backslash can escape the quote that + # follows and break out of the literal, so it is doubled too. + assert ( + iceberg_table_properties_clause( + {"table_properties": {"k": r"x\') COPY GRANTS --"}} + ) + == r" TABLE_PROPERTIES = ('k'='x\\'') COPY GRANTS --')" + ) + + +def test_create_iceberg_table_statement_with_table_properties(): + assert create_table_statement( + table_name="test_table", + schema="test_col varchar", + iceberg_config={ + "catalog": "SNOWFLAKE", + "base_location": "/root", + "table_properties": { + "read.split.target-size": "134217728", + "comment": "hi", + }, + }, + ) == ( + " CREATE ICEBERG TABLE test_table(test_col varchar) CATALOG = 'SNOWFLAKE' " + " BASE_LOCATION = '/root' TABLE_PROPERTIES = ('read.split.target-size'='134217728', 'comment'='hi')" + ) + + +def test_create_iceberg_table_as_select_statement_with_table_properties(): + assert create_table_as_select_statement( + table_name="test_table", + child="select * from foo", + column_definition=None, + iceberg_config={ + "catalog": "SNOWFLAKE", + "table_properties": {"read.split.target-size": "134217728"}, + }, + ) == ( + " CREATE ICEBERG TABLE test_table CATALOG = 'SNOWFLAKE' " + " TABLE_PROPERTIES = ('read.split.target-size'='134217728') AS SELECT * \n" + " FROM (\nselect * from foo\n)" + ) + + def test_create_dynamic_iceberg_table(): dt_name = "my_dt" warehouse = "my_warehouse"