snowpark-connect 0.32.0__py3-none-any.whl → 0.33.0__py3-none-any.whl

snowflake/snowpark_connect/relation/write/map_write.py

@@ -117,6 +117,57 @@ def _spark_to_snowflake(multipart_id: str) -> str:
     )
 
 
+def _validate_table_exist_and_of_type(
+    snowpark_table_name: str,
+    session: snowpark.Session,
+    table_type: str,
+    table_schema_or_error: DataType | SnowparkSQLException,
+) -> None:
+    if not isinstance(table_schema_or_error, DataType):
+        exception = AnalysisException(
+            f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{snowpark_table_name}` cannot be found."
+        )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
+    _validate_table_type(snowpark_table_name, session, table_type)
+
+
+def _validate_table_type(
+    snowpark_table_name: str,
+    session: snowpark.Session,
+    table_type: str,
+) -> None:
+    actual_type = get_table_type(snowpark_table_name, session)
+    if table_type == "iceberg":
+        if actual_type not in ("ICEBERG", "TABLE"):
+            exception = AnalysisException(
+                f"Table {snowpark_table_name} is not an iceberg table"
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
+    elif table_type == "fdn":
+        if actual_type not in ("NORMAL", "TABLE"):
+            exception = AnalysisException(
+                f"Table {snowpark_table_name} is not a FDN table"
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
+    else:
+        raise ValueError(
+            f"Invalid table_type: {table_type}. Must be 'iceberg' or 'fdn'"
+        )
+
+
+def _validate_table_does_not_exist(
+    snowpark_table_name: str,
+    table_schema_or_error: DataType | SnowparkSQLException,
+) -> None:
+    if isinstance(table_schema_or_error, DataType):
+        exception = AnalysisException(f"Table {snowpark_table_name} already exists")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
+
+
 def map_write(request: proto_base.ExecutePlanRequest):
     write_op = request.plan.command.write_operation
     telemetry.report_io_write(write_op.source)
@@ -240,7 +291,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 "format_type_options": {
                     "COMPRESSION": compression,
                 },
-                "overwrite": overwrite,
             }
             # Download from the base write path to ensure we fetch whatever Snowflake produced.
             # Using the base avoids coupling to exact filenames/prefixes.
@@ -347,14 +397,9 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     table_schema_or_error = _get_table_schema_or_error(
                         snowpark_table_name, session
                     )
-                    if isinstance(table_schema_or_error, DataType):  # Table exists
-                        exception = AnalysisException(
-                            f"Table {snowpark_table_name} already exists"
-                        )
-                        attach_custom_error_code(
-                            exception, ErrorCodes.INVALID_OPERATION
-                        )
-                        raise exception
+                    _validate_table_does_not_exist(
+                        snowpark_table_name, table_schema_or_error
+                    )
                     create_iceberg_table(
                         snowpark_table_name=snowpark_table_name,
                         location=write_op.options.get("location", None),
@@ -373,17 +418,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         snowpark_table_name, session
                     )
                     if isinstance(table_schema_or_error, DataType):  # Table exists
-                        if get_table_type(snowpark_table_name, session) not in (
-                            "ICEBERG",
-                            "TABLE",
-                        ):
-                            exception = AnalysisException(
-                                f"Table {snowpark_table_name} is not an iceberg table"
-                            )
-                            attach_custom_error_code(
-                                exception, ErrorCodes.INVALID_OPERATION
-                            )
-                            raise exception
+                        _validate_table_type(snowpark_table_name, session, "iceberg")
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -423,29 +458,25 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         snowpark_table_name, session
                     )
                     if isinstance(table_schema_or_error, DataType):  # Table exists
-                        if get_table_type(snowpark_table_name, session) not in (
-                            "ICEBERG",
-                            "TABLE",
-                        ):
-                            exception = AnalysisException(
-                                f"Table {snowpark_table_name} is not an iceberg table"
-                            )
-                            attach_custom_error_code(
-                                exception, ErrorCodes.INVALID_OPERATION
-                            )
-                            raise exception
+                        _validate_table_type(snowpark_table_name, session, "iceberg")
+                        create_iceberg_table(
+                            snowpark_table_name=snowpark_table_name,
+                            location=write_op.options.get("location", None),
+                            schema=input_df.schema,
+                            snowpark_session=session,
+                            mode="replace",
+                        )
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
                             location=write_op.options.get("location", None),
                             schema=input_df.schema,
                             snowpark_session=session,
+                            mode="create",
                         )
-                    _validate_schema_and_get_writer(
-                        input_df, "truncate", snowpark_table_name, table_schema_or_error
-                    ).saveAsTable(
+                    _get_writer_for_table_creation(input_df).saveAsTable(
                         table_name=snowpark_table_name,
-                        mode="truncate",
+                        mode="append",
                         column_order=_column_order_for_write,
                     )
                 case _:
@@ -491,18 +522,9 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             snowpark_table_name, session
                         )
                         if isinstance(table_schema_or_error, DataType):  # Table exists
-                            if get_table_type(snowpark_table_name, session) not in (
-                                "NORMAL",
-                                "TABLE",
-                            ):
-                                exception = AnalysisException(
-                                    f"Table {snowpark_table_name} is not a FDN table"
-                                )
-                                attach_custom_error_code(
-                                    exception, ErrorCodes.INVALID_OPERATION
-                                )
-                                raise exception
-                            write_mode = "truncate"
+                            _validate_table_type(snowpark_table_name, session, "fdn")
+
+                        write_mode = "overwrite"
                         _validate_schema_and_get_writer(
                             input_df,
                             write_mode,
@@ -511,27 +533,15 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         ).saveAsTable(
                             table_name=snowpark_table_name,
                             mode=write_mode,
+                            copy_grants=True,
                             column_order=_column_order_for_write,
                         )
                     case "append":
                         table_schema_or_error = _get_table_schema_or_error(
                             snowpark_table_name, session
                         )
-                        if isinstance(
-                            table_schema_or_error, DataType
-                        ) and get_table_type(  # Table exists
-                            snowpark_table_name, session
-                        ) not in (
-                            "NORMAL",
-                            "TABLE",
-                        ):
-                            exception = AnalysisException(
-                                f"Table {snowpark_table_name} is not a FDN table"
-                            )
-                            attach_custom_error_code(
-                                exception, ErrorCodes.INVALID_OPERATION
-                            )
-                            raise exception
+                        if isinstance(table_schema_or_error, DataType):  # Table exists
+                            _validate_table_type(snowpark_table_name, session, "fdn")
 
                         _validate_schema_and_get_writer(
                             input_df,
@@ -598,113 +608,105 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
         attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
         raise exception
 
-    if write_op.provider.lower() == "iceberg":
-        match write_op.mode:
-            case commands_proto.WriteOperationV2.MODE_CREATE:
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if isinstance(table_schema_or_error, DataType):  # Table exists
-                    exception = AnalysisException(
-                        f"Table {snowpark_table_name} already exists"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
+    is_iceberg = write_op.provider.lower() == "iceberg"
+    table_type = "iceberg" if is_iceberg else "fdn"
+
+    match write_op.mode:
+        case commands_proto.WriteOperationV2.MODE_CREATE:
+            table_schema_or_error = _get_table_schema_or_error(
+                snowpark_table_name, session
+            )
+            _validate_table_does_not_exist(snowpark_table_name, table_schema_or_error)
+
+            if is_iceberg:
                 create_iceberg_table(
                     snowpark_table_name=snowpark_table_name,
                     location=write_op.table_properties.get("location"),
                     schema=input_df.schema,
                     snowpark_session=session,
                 )
-                _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="append",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_APPEND:
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    exception = AnalysisException(
-                        f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
-                    raise exception
-                if get_table_type(snowpark_table_name, session) not in (
-                    "ICEBERG",
-                    "TABLE",
-                ):
-                    exception = AnalysisException(
-                        f"Table {snowpark_table_name} is not an iceberg table"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="append",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
-                # TODO: handle the filter condition for MODE_OVERWRITE
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if isinstance(table_schema_or_error, DataType):  # Table exists
-                    if get_table_type(snowpark_table_name, session) not in (
-                        "ICEBERG",
-                        "TABLE",
-                    ):
-                        exception = AnalysisException(
-                            f"Table {snowpark_table_name} is not an iceberg table"
-                        )
-                        attach_custom_error_code(
-                            exception, ErrorCodes.INVALID_OPERATION
-                        )
-                        raise exception
-                else:
-                    exception = AnalysisException(
-                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "truncate", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="truncate",
-                    column_order=_column_order_for_write,
+            _get_writer_for_table_creation(input_df).saveAsTable(
+                table_name=snowpark_table_name,
+                mode="append" if is_iceberg else "errorifexists",
+                column_order=_column_order_for_write,
+            )
+
+        case commands_proto.WriteOperationV2.MODE_APPEND:
+            table_schema_or_error = _get_table_schema_or_error(
+                snowpark_table_name, session
+            )
+            _validate_table_exist_and_of_type(
+                snowpark_table_name, session, table_type, table_schema_or_error
+            )
+            _validate_schema_and_get_writer(
+                input_df, "append", snowpark_table_name, table_schema_or_error
+            ).saveAsTable(
+                table_name=snowpark_table_name,
+                mode="append",
+                column_order=_column_order_for_write,
+            )
+
+        case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
+            # TODO: handle the filter condition for MODE_OVERWRITE
+            table_schema_or_error = _get_table_schema_or_error(
+                snowpark_table_name, session
+            )
+            _validate_table_exist_and_of_type(
+                snowpark_table_name, session, table_type, table_schema_or_error
+            )
+
+            if is_iceberg:
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.options.get("location", None),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                    mode="replace",
                 )
-            case commands_proto.WriteOperationV2.MODE_REPLACE:
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
+                writer = _get_writer_for_table_creation(input_df)
+                save_mode = "append"
+            else:
+                writer = _validate_schema_and_get_writer(
+                    input_df, "overwrite", snowpark_table_name, table_schema_or_error
                 )
-                if isinstance(table_schema_or_error, DataType):  # Table exists
-                    create_iceberg_table(
-                        snowpark_table_name=snowpark_table_name,
-                        location=write_op.table_properties.get("location"),
-                        schema=input_df.schema,
-                        snowpark_session=session,
-                        mode="replace",
-                    )
-                else:
-                    exception = AnalysisException(
-                        f"Table {snowpark_table_name} does not exist"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="append",
-                    column_order=_column_order_for_write,
+                save_mode = "overwrite"
+
+            writer.saveAsTable(
+                table_name=snowpark_table_name,
+                mode=save_mode,
+                column_order=_column_order_for_write,
+            )
+
+        case commands_proto.WriteOperationV2.MODE_REPLACE:
+            table_schema_or_error = _get_table_schema_or_error(
+                snowpark_table_name, session
+            )
+            _validate_table_exist_and_of_type(
+                snowpark_table_name, session, table_type, table_schema_or_error
+            )
+
+            if is_iceberg:
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.table_properties.get("location"),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                    mode="replace",
                 )
-            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+                save_mode = "append"
+            else:
+                save_mode = "overwrite"
+
+            _validate_schema_and_get_writer(
+                input_df, "replace", snowpark_table_name, table_schema_or_error
+            ).saveAsTable(
+                table_name=snowpark_table_name,
+                mode=save_mode,
+                column_order=_column_order_for_write,
+            )
+
+        case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+            if is_iceberg:
                 create_iceberg_table(
                     snowpark_table_name=snowpark_table_name,
                     location=write_op.table_properties.get("location"),
@@ -712,116 +714,24 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_session=session,
                     mode="create_or_replace",
                 )
-                _validate_schema_and_get_writer(
-                    input_df, "create_or_replace", snowpark_table_name
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="append",
-                    column_order=_column_order_for_write,
-                )
-            case _:
-                exception = SnowparkConnectNotImplementedError(
-                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
-                )
-                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
-                raise exception
-    else:
-        match write_op.mode:
-            case commands_proto.WriteOperationV2.MODE_CREATE:
-                _validate_schema_and_get_writer(
-                    input_df, "errorifexists", snowpark_table_name
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="errorifexists",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_APPEND:
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    exception = AnalysisException(
-                        f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
-                if get_table_type(snowpark_table_name, session) not in (
-                    "NORMAL",
-                    "TABLE",
-                ):
-                    exception = AnalysisException(
-                        f"Table {snowpark_table_name} is not a FDN table"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="append",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
-                # TODO: handle the filter condition for MODE_OVERWRITE
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if isinstance(table_schema_or_error, DataType):  # Table exists
-                    if get_table_type(snowpark_table_name, session) not in (
-                        "NORMAL",
-                        "TABLE",
-                    ):
-                        exception = AnalysisException(
-                            f"Table {snowpark_table_name} is not a FDN table"
-                        )
-                        attach_custom_error_code(
-                            exception, ErrorCodes.INVALID_OPERATION
-                        )
-                        raise exception
-                else:
-                    exception = AnalysisException(
-                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.TABLE_NOT_FOUND)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "truncate", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="truncate",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_REPLACE:
-                table_schema_or_error = _get_table_schema_or_error(
-                    snowpark_table_name, session
-                )
-                if not isinstance(table_schema_or_error, DataType):  # Table not exists
-                    exception = AnalysisException(
-                        f"Table {snowpark_table_name} does not exist"
-                    )
-                    attach_custom_error_code(exception, ErrorCodes.TABLE_NOT_FOUND)
-                    raise exception
-                _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name, table_schema_or_error
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="overwrite",
-                    column_order=_column_order_for_write,
-                )
-            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
-                _validate_schema_and_get_writer(
-                    input_df, "create_or_replace", snowpark_table_name
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode="overwrite",
-                    column_order=_column_order_for_write,
-                )
-            case _:
-                exception = SnowparkConnectNotImplementedError(
-                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
-                )
-                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
-                raise exception
+                save_mode = "append"
+            else:
+                save_mode = "overwrite"
+
+            _validate_schema_and_get_writer(
+                input_df, "create_or_replace", snowpark_table_name
+            ).saveAsTable(
+                table_name=snowpark_table_name,
+                mode=save_mode,
+                column_order=_column_order_for_write,
+            )
+
+        case _:
+            exception = SnowparkConnectNotImplementedError(
+                f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def _get_table_schema_or_error(
@@ -856,6 +766,7 @@ def _validate_schema_and_get_writer(
     if write_mode is not None and write_mode.lower() in (
         "replace",
         "create_or_replace",
+        "overwrite",
     ):
         return _get_writer_for_table_creation(input_df)
 
@@ -1053,14 +964,13 @@ def create_iceberg_table(
         if config_external_volume is None or config_external_volume == ""
         else f"EXTERNAL_VOLUME = '{config_external_volume}'"
     )
-
+    copy_grants = ""
     match mode:
         case "create":
             create_sql = "CREATE"
-        case "replace":
+        case "replace" | "create_or_replace":
             # There's no replace for iceberg table, so we use create or replace
-            create_sql = "CREATE OR REPLACE"
-        case "create_or_replace":
+            copy_grants = "COPY GRANTS"
             create_sql = "CREATE OR REPLACE"
         case _:
             exception = SnowparkConnectNotImplementedError(
@@ -1072,7 +982,8 @@ def create_iceberg_table(
         {create_sql} ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
         CATALOG = 'SNOWFLAKE'
         {external_volume}
-        {base_location};
+        {base_location}
+        {copy_grants};
         """
     snowpark_session.sql(sql).collect()
 

snowflake/snowpark_connect/resources_initializer.py

@@ -1,7 +1,6 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-import pathlib
 import threading
 import time
 
@@ -51,11 +50,9 @@ def initialize_resources() -> None:
         """Upload Spark jar files required for creating Scala UDFs."""
         stage = session.get_session_stage()
         resource_path = stage + RESOURCE_PATH
-        import snowflake
+        import snowpark_connect_deps_1
+        import snowpark_connect_deps_2
 
-        pyspark_jars = (
-            pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes/jars"
-        )
         jar_files = [
             f"spark-sql_2.12-{SPARK_VERSION}.jar",
             f"spark-connect-client-jvm_2.12-{SPARK_VERSION}.jar",
@@ -64,14 +61,29 @@ def initialize_resources() -> None:
             "json4s-ast_2.12-3.7.0-M11.jar",
         ]
 
-        for jar in jar_files:
-            session.file.put(
-                str(pyspark_jars) + "/" + jar,
-                resource_path,
-                auto_compress=False,
-                overwrite=False,
-                source_compression="NONE",
-            )
+        for jar_name in jar_files:
+            # Try to find the JAR in package 1 first, then package 2
+            jar_path = None
+            try:
+                jar_path = snowpark_connect_deps_1.get_jar_path(jar_name)
+            except FileNotFoundError:
+                try:
+                    jar_path = snowpark_connect_deps_2.get_jar_path(jar_name)
+                except FileNotFoundError:
+                    raise FileNotFoundError(
+                        f"JAR {jar_name} not found in either package"
+                    )
+
+            try:
+                session.file.put(
+                    str(jar_path),
+                    resource_path,
+                    auto_compress=False,
+                    overwrite=False,
+                    source_compression="NONE",
+                )
+            except Exception as e:
+                raise RuntimeError(f"Failed to upload JAR {jar_name}: {e}")
 
     start_time = time.time()