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

snowflake/snowpark_connect/relation/write/map_write.py

@@ -2,8 +2,10 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import copy
 import os
 import shutil
+import uuid
 from pathlib import Path
 
 import pyspark.sql.connect.proto.base_pb2 as proto_base
@@ -29,6 +31,7 @@ from snowflake.snowpark.types import (
 )
 from snowflake.snowpark_connect.config import (
     auto_uppercase_column_identifiers,
+    get_success_file_generation_enabled,
     global_config,
     sessions_config,
     str_to_bool,
@@ -117,6 +120,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)
@@ -198,20 +252,9 @@ def map_write(request: proto_base.ExecutePlanRequest):
             )
 
             if overwrite:
-                try:
-                    path_after_stage = (
-                        write_path.split("/", 1)[1] if "/" in write_path else ""
-                    )
-                    if not path_after_stage or path_after_stage == "/":
-                        logger.warning(
-                            f"Skipping REMOVE for root path {write_path} - too broad scope"
-                        )
-                    else:
-                        remove_command = f"REMOVE '{write_path}/'"
-                        session.sql(remove_command).collect()
-                        logger.info(f"Successfully cleared directory: {write_path}")
-                except Exception as e:
-                    logger.warning(f"Could not clear directory {write_path}: {e}")
+                remove_command = f"REMOVE '{write_path}'"
+                session.sql(remove_command).collect()
+                logger.info(f"Successfully cleared directory: {write_path}")
 
             if should_write_to_single_file and partition_hint is None:
                 # Single file: generate complete filename with extension
@@ -240,7 +283,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.
@@ -300,15 +342,20 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 # Execute multiple COPY INTO operations, one per target file.
                 # Since we write per-partition with distinct prefixes, download from the base write path.
                 download_stage_path = write_path
+
+                # We need to create a new set of parameters with single=True
+                shared_uuid = str(uuid.uuid4())
+                part_params = copy.deepcopy(dict(parameters))
+                part_params["single"] = True
                 for part_idx in range(partition_hint):
-                    part_params = dict(parameters)
                     # Preserve Spark-like filename prefix per partition so downloaded basenames
                     # match the expected Spark pattern (with possible Snowflake counters appended).
                     per_part_prefix = generate_spark_compatible_filename(
                         task_id=part_idx,
                         attempt_number=0,
-                        compression=None,
-                        format_ext="",  # prefix only; Snowflake appends extension/counters
+                        compression=compression,
+                        format_ext=extension,
+                        shared_uuid=shared_uuid,
                     )
                     part_params["location"] = f"{write_path}/{per_part_prefix}"
                     (
@@ -318,6 +365,9 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     )
             else:
                 rewritten_df.write.copy_into_location(**parameters)
+
+            generate_success = get_success_file_generation_enabled()
+
             if not is_cloud_path(write_op.path):
                 store_files_locally(
                     download_stage_path,
@@ -325,6 +375,13 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     overwrite,
                     session,
                 )
+                if generate_success:
+                    _write_success_file_locally(write_op.path)
+            else:
+                if generate_success:
+                    _write_success_file_to_stage(
+                        download_stage_path, session, parameters
+                    )
         case "jdbc":
             from snowflake.snowpark_connect.relation.write.map_write_jdbc import (
                 map_write_jdbc,
@@ -347,14 +404,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 +425,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 +465,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 +529,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 +540,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 +615,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 +721,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 +773,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 +971,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 +989,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()
 
@@ -1133,6 +1051,47 @@ def handle_column_names(
     return df
 
 
+def _write_success_file_locally(directory_path: str) -> None:
+    """
+    Write a _SUCCESS marker file to a local directory.
+    """
+    try:
+        success_file = Path(directory_path) / "_SUCCESS"
+        success_file.touch()
+        logger.debug(f"Created _SUCCESS file at {directory_path}")
+    except Exception as e:
+        logger.warning(f"Failed to create _SUCCESS file at {directory_path}: {e}")
+
+
+def _write_success_file_to_stage(
+    stage_path: str,
+    session: snowpark.Session,
+    parameters: dict,
+) -> None:
+    """
+    Write a _SUCCESS marker file to a stage location.
+    """
+    try:
+        # Create a dummy dataframe with one row containing "SUCCESS"
+        success_df = session.create_dataframe([["SUCCESS"]]).to_df(["STATUS"])
+        success_params = copy.deepcopy(parameters)
+        success_params["location"] = f"{stage_path}/_SUCCESS"
+        success_params["single"] = True
+        success_params["header"] = True
+
+        # Set CSV format with explicit no compression for _SUCCESS file
+        success_params["file_format_type"] = "csv"
+        success_params["format_type_options"] = {
+            "COMPRESSION": "NONE",
+        }
+
+        success_df.write.copy_into_location(**success_params)
+
+        logger.debug(f"Created _SUCCESS file at {stage_path}")
+    except Exception as e:
+        logger.warning(f"Failed to create _SUCCESS file at {stage_path}: {e}")
+
+
 def store_files_locally(
     stage_path: str, target_path: str, overwrite: bool, session: snowpark.Session
 ) -> None: