ygg 0.1.33__py3-none-any.whl → 0.1.35__py3-none-any.whl

yggdrasil/databricks/sql/engine.py

@@ -1,4 +1,15 @@
-"""Databricks SQL engine utilities and helpers."""
+"""Databricks SQL engine utilities and helpers.
+
+This module provides a thin “do the right thing” layer over:
+- Databricks SQL Statement Execution API (warehouse)
+- Spark SQL / Delta Lake (when running inside a Spark-enabled context)
+
+It includes helpers to:
+- Build fully-qualified table names
+- Execute SQL via Spark or Databricks SQL API
+- Insert Arrow/Spark data into Delta tables (append/overwrite/merge)
+- Generate DDL from Arrow schemas
+"""
 
 import dataclasses
 import logging
@@ -8,7 +19,6 @@ import time
 from typing import Optional, Union, Any, Dict, List, Literal
 
 import pyarrow as pa
-import pyarrow.parquet as pq
 
 from .statement_result import StatementResult
 from .types import column_info_to_arrow_field
@@ -28,7 +38,6 @@ except ImportError:
         @classmethod
         def forName(cls, *args, **kwargs):
             from delta.tables import DeltaTable
-
             return DeltaTable.forName(*args, **kwargs)
 
 
@@ -37,23 +46,18 @@ if databricks_sdk is not None:
         StatementResponse, Disposition, Format,
         ExecuteStatementRequestOnWaitTimeout, StatementParameterListItem
     )
-
     StatementResponse = StatementResponse
 else:
-    class StatementResponse:
+    class StatementResponse:  # pragma: no cover
         pass
 
 
 logger = logging.getLogger(__name__)
 
-
 if pyspark is not None:
     import pyspark.sql.functions as F
 
-__all__ = [
-    "SQLEngine",
-    "StatementResult"
-]
+__all__ = ["SQLEngine", "StatementResult"]
 
 
 class SqlExecutionError(RuntimeError):
@@ -62,7 +66,7 @@ class SqlExecutionError(RuntimeError):
 
 @dataclasses.dataclass
 class SQLEngine(WorkspaceService):
-    """Execute SQL statements and manage tables via Databricks."""
+    """Execute SQL statements and manage tables via Databricks SQL / Spark."""
     warehouse_id: Optional[str] = None
     catalog_name: Optional[str] = None
     schema_name: Optional[str] = None
@@ -72,18 +76,18 @@ class SQLEngine(WorkspaceService):
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
-        safe_chars: bool = True
-    ):
-        """Build a fully qualified table name for the current catalog/schema.
+        safe_chars: bool = True,
+    ) -> str:
+        """Build a fully qualified table name (catalog.schema.table).
 
         Args:
-            catalog_name: Optional catalog override.
-            schema_name: Optional schema override.
+            catalog_name: Optional catalog override (defaults to engine.catalog_name).
+            schema_name: Optional schema override (defaults to engine.schema_name).
             table_name: Table name to qualify.
-            safe_chars: Whether to wrap identifiers in backticks.
+            safe_chars: Whether to wrap each identifier in backticks.
 
         Returns:
-            The fully qualified table name.
+            Fully qualified table name string.
         """
         catalog_name = catalog_name or self.catalog_name
         schema_name = schema_name or self.schema_name
@@ -96,21 +100,23 @@ class SQLEngine(WorkspaceService):
             return f"`{catalog_name}`.`{schema_name}`.`{table_name}`"
         return f"{catalog_name}.{schema_name}.{table_name}"
 
-    def _catalog_schema_table_names(
-        self,
-        full_name: str,
-    ):
+    def _catalog_schema_table_names(self, full_name: str) -> tuple[Optional[str], Optional[str], Optional[str]]:
         """Parse a catalog.schema.table string into components.
 
+        Supports partial names:
+        - table
+        - schema.table
+        - catalog.schema.table
+
+        Backticks are stripped.
+
         Args:
-            full_name: A fully qualified name or partial name.
+            full_name: Fully qualified or partial table name.
 
         Returns:
-            A tuple of (catalog_name, schema_name, table_name).
+            Tuple of (catalog_name, schema_name, table_name).
         """
-        parts = [
-            _.strip("`") for _ in full_name.split(".")
-        ]
+        parts = [_.strip("`") for _ in full_name.split(".")]
 
         if len(parts) == 0:
             return self.catalog_name, self.schema_name, None
@@ -122,20 +128,20 @@ class SQLEngine(WorkspaceService):
         catalog_name, schema_name, table_name = parts[-3], parts[-2], parts[-1]
         catalog_name = catalog_name or self.catalog_name
         schema_name = schema_name or self.schema_name
-
         return catalog_name, schema_name, table_name
 
-    def _default_warehouse(
-        self,
-        cluster_size: str = "Small"
-    ):
-        """Return a default SQL warehouse matching the desired size.
+    def _default_warehouse(self, cluster_size: str = "Small"):
+        """Pick a default SQL warehouse (best-effort) matching the desired size.
 
         Args:
-            cluster_size: Desired warehouse size filter.
+            cluster_size: Desired warehouse size (Databricks "cluster_size"), e.g. "Small".
+                If empty/None, returns the first warehouse encountered.
 
         Returns:
-            The matched warehouse object.
+            Warehouse object.
+
+        Raises:
+            ValueError: If no warehouses exist in the workspace.
         """
         wk = self.workspace.sdk()
         existing = list(wk.warehouses.list())
@@ -146,48 +152,54 @@ class SQLEngine(WorkspaceService):
                 first = warehouse
 
             if cluster_size:
-                if warehouse.cluster_size == cluster_size:
+                if getattr(warehouse, "cluster_size", None) == cluster_size:
+                    logger.debug("Default warehouse match found: id=%s cluster_size=%s", warehouse.id, warehouse.cluster_size)
                     return warehouse
             else:
+                logger.debug("Default warehouse selected (first): id=%s", warehouse.id)
                 return warehouse
 
         if first is not None:
+            logger.info(
+                "No warehouse matched cluster_size=%s; falling back to first warehouse id=%s cluster_size=%s",
+                cluster_size,
+                getattr(first, "id", None),
+                getattr(first, "cluster_size", None),
+            )
             return first
 
         raise ValueError(f"No default warehouse found in {wk.config.host}")
 
-    def _get_or_default_warehouse_id(
-        self,
-        cluster_size = "Small"
-    ):
-        """Return the configured warehouse id or a default one.
+    def _get_or_default_warehouse_id(self, cluster_size: str = "Small") -> str:
+        """Return configured warehouse_id or resolve a default one.
 
         Args:
-            cluster_size: Desired warehouse size filter.
+            cluster_size: Desired warehouse size filter used when resolving defaults.
 
         Returns:
-            The warehouse id string.
+            Warehouse id string.
         """
         if not self.warehouse_id:
             dft = self._default_warehouse(cluster_size=cluster_size)
-
             self.warehouse_id = dft.id
+            logger.info("Resolved default warehouse_id=%s (cluster_size=%s)", self.warehouse_id, cluster_size)
+
         return self.warehouse_id
 
     @staticmethod
     def _random_suffix(prefix: str = "") -> str:
-        """Generate a unique suffix for temporary resources.
-
-        Args:
-            prefix: Optional prefix to prepend.
-
-        Returns:
-            A unique suffix string.
-        """
-        unique = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
+        """Generate a unique suffix for temporary resources."""
+        unique = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
         timestamp = int(time.time() * 1000)
         return f"{prefix}{timestamp}_{unique}"
 
+    @staticmethod
+    def _sql_preview(sql: str, limit: int = 220) -> str:
+        """Short, single-line preview for logs (avoids spewing giant SQL)."""
+        if not sql:
+            return ""
+        return sql[:limit] + ("…" if len(sql) > limit else "")
+
     def execute(
         self,
         statement: Optional[str] = None,
@@ -205,56 +217,67 @@ class SQLEngine(WorkspaceService):
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         wait_result: bool = True,
-        **kwargs,
     ) -> "StatementResult":
-        """
-        Execute a SQL statement on a SQL warehouse.
+        """Execute a SQL statement via Spark or Databricks SQL Statement Execution API.
 
-        - If wait=True (default): poll until terminal state.
-            - On SUCCEEDED: return final statement object
-            - On FAILED / CANCELED: raise SqlExecutionError
-        - If wait=False: return initial execution handle without polling.
+        Engine resolution:
+        - If `engine` is not provided and a Spark session is active -> uses Spark.
+        - Otherwise uses Databricks SQL API (warehouse).
+
+        Waiting behavior (`wait_result`):
+        - If True (default): returns a StatementResult in terminal state (SUCCEEDED/FAILED/CANCELED).
+        - If False: returns immediately with the initial handle (caller can `.wait()` later).
 
         Args:
-            statement: SQL statement to execute. If omitted, selects from the table.
-            engine: Execution engine ("spark" or "api").
-            warehouse_id: Optional warehouse id override.
+            statement: SQL statement to execute. If None, a `SELECT *` is generated from the table params.
+            engine: "spark" or "api".
+            warehouse_id: Warehouse override (for API engine).
             byte_limit: Optional byte limit for results.
-            disposition: Result disposition mode.
-            format: Result format for Databricks SQL API.
-            on_wait_timeout: Timeout behavior for waiting.
-            parameters: Optional statement parameters.
-            row_limit: Optional row limit.
-            wait_timeout: Optional API wait timeout.
-            catalog_name: Optional catalog override.
-            schema_name: Optional schema override.
-            table_name: Optional table name override.
-            wait_result: Whether to block until completion.
-            **kwargs: Additional API parameters.
+            disposition: Result disposition mode (API engine).
+            format: Result format (API engine).
+            on_wait_timeout: Timeout behavior for waiting (API engine).
+            parameters: Optional statement parameters (API engine).
+            row_limit: Optional row limit for results (API engine).
+            wait_timeout: API wait timeout value.
+            catalog_name: Optional catalog override for API engine.
+            schema_name: Optional schema override for API engine.
+            table_name: Optional table override used when `statement` is None.
+            wait_result: Whether to block until completion (API engine).
 
         Returns:
-            A StatementResult wrapper for the execution.
+            StatementResult.
         """
+        # --- Engine auto-detection ---
         if not engine:
             if pyspark is not None:
                 spark_session = SparkSession.getActiveSession()
-
                 if spark_session is not None:
                     engine = "spark"
 
+        # --- Spark path ---
         if engine == "spark":
             spark_session = SparkSession.getActiveSession()
-
             if spark_session is None:
                 raise ValueError("No spark session found to run sql query")
 
+            df: SparkDataFrame = spark_session.sql(statement)
+
+            if row_limit:
+                df = df.limit(row_limit)
+
+            logger.info("Spark SQL executed: %s", self._sql_preview(statement))
+
+            # Avoid Disposition dependency if SDK imports are absent
+            spark_disp = disposition if disposition is not None else getattr(globals().get("Disposition", object), "EXTERNAL_LINKS", None)
+
             return StatementResult(
                 engine=self,
                 statement_id="sparksql",
-                disposition=Disposition.EXTERNAL_LINKS,
-                _spark_df=spark_session.sql(statement)
+                disposition=spark_disp,
+                _spark_df=df,
             )
 
+        # --- API path defaults ---
         if format is None:
             format = Format.ARROW_STREAM
 
@@ -280,7 +303,6 @@ class SQLEngine(WorkspaceService):
             wait_timeout=wait_timeout,
             catalog=catalog_name or self.catalog_name,
             schema=schema_name or self.schema_name,
-            **kwargs,
         )
 
         execution = StatementResult(
@@ -288,10 +310,15 @@ class SQLEngine(WorkspaceService):
             statement_id=response.statement_id,
             _response=response,
             _response_refresh_time=time.time(),
-            disposition=disposition
+            disposition=disposition,
+        )
+
+        logger.info(
+            "API SQL executed: %s",
+            self._sql_preview(statement)
         )
 
-        return execution.wait() if wait_result else wait_result
+        return execution.wait() if wait_result else execution
 
     def spark_table(
         self,
@@ -300,35 +327,21 @@ class SQLEngine(WorkspaceService):
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
     ):
-        """Return a DeltaTable handle for a given table name.
-
-        Args:
-            full_name: Fully qualified table name.
-            catalog_name: Optional catalog override.
-            schema_name: Optional schema override.
-            table_name: Optional table name override.
-
-        Returns:
-            A Spark DeltaTable handle.
-        """
+        """Return a DeltaTable handle for a given table name (Spark context required)."""
         if not full_name:
             full_name = self.table_full_name(
                 catalog_name=catalog_name,
                 schema_name=schema_name,
-                table_name=table_name
+                table_name=table_name,
             )
-
         return SparkDeltaTable.forName(
             sparkSession=SparkSession.getActiveSession(),
-            tableOrViewName=full_name
+            tableOrViewName=full_name,
         )
 
     def insert_into(
         self,
-        data: Union[
-            pa.Table, pa.RecordBatch, pa.RecordBatchReader,
-            SparkDataFrame
-        ],
+        data: Union[pa.Table, pa.RecordBatch, pa.RecordBatchReader, SparkDataFrame],
         location: Optional[str] = None,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
@@ -336,14 +349,18 @@ class SQLEngine(WorkspaceService):
         mode: str = "auto",
         cast_options: Optional[CastOptions] = None,
         overwrite_schema: bool | None = None,
-        match_by: list[str] = None,
-        zorder_by: list[str] = None,
+        match_by: Optional[list[str]] = None,
+        zorder_by: Optional[list[str]] = None,
         optimize_after_merge: bool = False,
-        vacuum_hours: int | None = None,  # e.g., 168 for 7 days
+        vacuum_hours: int | None = None,
         spark_session: Optional[SparkSession] = None,
-        spark_options: Optional[Dict[str, Any]] = None
+        spark_options: Optional[Dict[str, Any]] = None,
     ):
-        """Insert data into a table using Spark or Arrow paths.
+        """Insert data into a Delta table using Spark when available; otherwise stage Arrow.
+
+        Strategy:
+        - If Spark is available and we have an active session (or Spark DF input) -> use `spark_insert_into`.
+        - Otherwise -> use `arrow_insert_into` (stages Parquet to a temp volume + runs SQL INSERT/MERGE).
 
         Args:
             data: Arrow or Spark data to insert.
@@ -353,18 +370,18 @@ class SQLEngine(WorkspaceService):
             table_name: Optional table name override.
             mode: Insert mode ("auto", "append", "overwrite").
             cast_options: Optional casting options.
-            overwrite_schema: Whether to overwrite schema (Spark).
-            match_by: Optional merge keys for upserts.
-            zorder_by: Optional Z-ORDER columns.
-            optimize_after_merge: Whether to run OPTIMIZE after merge.
+            overwrite_schema: Whether to overwrite schema (Spark path).
+            match_by: Merge keys for upserts (MERGE semantics). When set, mode affects behavior.
+            zorder_by: Z-ORDER columns (SQL path uses OPTIMIZE ZORDER; Spark path uses Delta optimize API).
+            optimize_after_merge: Whether to run OPTIMIZE after a merge (SQL path) / after merge+zorder (Spark path).
             vacuum_hours: Optional VACUUM retention window.
             spark_session: Optional SparkSession override.
             spark_options: Optional Spark write options.
 
         Returns:
-            None for Arrow inserts, or the Spark insert result.
+            None (mutates the destination table).
         """
-        # -------- existing logic you provided (kept intact) ----------
+
         if pyspark is not None:
             spark_session = SparkSession.getActiveSession() if spark_session is None else spark_session
 
@@ -382,7 +399,7 @@ class SQLEngine(WorkspaceService):
                     zorder_by=zorder_by,
                     optimize_after_merge=optimize_after_merge,
                     vacuum_hours=vacuum_hours,
-                    spark_options=spark_options
+                    spark_options=spark_options,
                 )
 
         return self.arrow_insert_into(
@@ -402,9 +419,7 @@ class SQLEngine(WorkspaceService):
 
     def arrow_insert_into(
         self,
-        data: Union[
-            pa.Table, pa.RecordBatch, pa.RecordBatchReader,
-        ],
+        data: Union[pa.Table, pa.RecordBatch, pa.RecordBatchReader],
         location: Optional[str] = None,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
@@ -412,14 +427,19 @@ class SQLEngine(WorkspaceService):
         mode: str = "auto",
         cast_options: Optional[CastOptions] = None,
         overwrite_schema: bool | None = None,
-        match_by: list[str] = None,
-        zorder_by: list[str] = None,
+        match_by: Optional[list[str]] = None,
+        zorder_by: Optional[list[str]] = None,
         optimize_after_merge: bool = False,
-        vacuum_hours: int | None = None,  # e.g., 168 for 7 days
+        vacuum_hours: int | None = None,
         existing_schema: pa.Schema | None = None,
-        temp_volume_path: Optional[Union[str, DatabricksPath]] = None
+        temp_volume_path: Optional[Union[str, DatabricksPath]] = None,
     ):
-        """Insert Arrow data by staging to a temp volume and running SQL.
+        """Insert Arrow data by staging Parquet to a temp volume and running Databricks SQL.
+
+        Notes:
+        - If the table does not exist, it is created from the input Arrow schema (best-effort).
+        - If `match_by` is provided, uses MERGE INTO (upsert).
+        - Otherwise uses INSERT INTO / INSERT OVERWRITE depending on mode.
 
         Args:
             data: Arrow table/batch data to insert.
@@ -427,14 +447,14 @@ class SQLEngine(WorkspaceService):
             catalog_name: Optional catalog override.
             schema_name: Optional schema override.
             table_name: Optional table name override.
-            mode: Insert mode ("auto", "append", "overwrite").
+            mode: Insert mode ("auto", "append", "overwrite"). ("auto" behaves like append here.)
             cast_options: Optional casting options.
-            overwrite_schema: Whether to overwrite schema.
-            match_by: Optional merge keys for upserts.
-            zorder_by: Optional Z-ORDER columns.
-            optimize_after_merge: Whether to run OPTIMIZE after merge.
-            vacuum_hours: Optional VACUUM retention window.
-            existing_schema: Optional pre-fetched schema.
+            overwrite_schema: Reserved for parity with Spark path (unused here).
+            match_by: Merge keys for MERGE INTO upserts.
+            zorder_by: Columns for OPTIMIZE ZORDER BY.
+            optimize_after_merge: Run OPTIMIZE after MERGE (in addition to ZORDER optimization).
+            vacuum_hours: Optional VACUUM retention window in hours.
+            existing_schema: Optional pre-fetched destination schema (Arrow).
             temp_volume_path: Optional temp volume path override.
 
         Returns:
@@ -445,26 +465,26 @@ class SQLEngine(WorkspaceService):
             catalog_name=catalog_name,
             schema_name=schema_name,
             table_name=table_name,
-            safe_chars=True
+            safe_chars=True,
         )
 
-        with self as connected:
+        with self.connect() as connected:
             if existing_schema is None:
                 try:
                     existing_schema = connected.get_table_schema(
                         catalog_name=catalog_name,
                         schema_name=schema_name,
                         table_name=table_name,
-                        to_arrow_schema=True
+                        to_arrow_schema=True,
                     )
                 except ValueError as exc:
-                    data = convert(data, pa.Table)
-                    existing_schema = data.schema
+                    data_tbl = convert(data, pa.Table)
+                    existing_schema = data_tbl.schema
                     logger.warning(
-                        "Table %s not found, %s, creating it based on input data %s",
+                        "Table %s not found (%s). Creating it from input schema (columns=%s)",
                         location,
                         exc,
-                        existing_schema.names
+                        existing_schema.names,
                     )
 
                     connected.create_table(
@@ -472,12 +492,12 @@ class SQLEngine(WorkspaceService):
                         catalog_name=catalog_name,
                         schema_name=schema_name,
                         table_name=table_name,
-                        if_not_exists=True
+                        if_not_exists=True,
                     )
 
                     try:
                         return connected.arrow_insert_into(
-                            data=data,
+                            data=data_tbl,
                             location=location,
                             catalog_name=catalog_name,
                             schema_name=schema_name,
@@ -489,54 +509,62 @@ class SQLEngine(WorkspaceService):
                             zorder_by=zorder_by,
                             optimize_after_merge=optimize_after_merge,
                             vacuum_hours=vacuum_hours,
-                            existing_schema=existing_schema
+                            existing_schema=existing_schema,
                         )
-                    except:
+                    except Exception:
+                        logger.exception("Arrow insert failed after auto-creating %s; attempting cleanup (DROP TABLE)", location)
                         try:
                             connected.drop_table(location=location)
-                        except Exception as e:
-                            logger.warning("Failed to drop table %s after auto creation on error: %s", location, e)
+                        except Exception:
+                            logger.exception("Failed to drop table %s after auto creation error", location)
                         raise
 
             transaction_id = self._random_suffix()
 
-            data = convert(data, pa.Table, options=cast_options, target_field=existing_schema)
+            data_tbl = convert(
+                data, pa.Table,
+                options=cast_options, target_field=existing_schema
+            )
+            num_rows = data_tbl.num_rows
+
+            logger.debug(
+                "Arrow inserting %s rows into %s (mode=%s, match_by=%s, zorder_by=%s)",
+                num_rows,
+                location,
+                mode,
+                match_by,
+                zorder_by,
+            )
 
             # Write in temp volume
             temp_volume_path = connected.dbfs_path(
                 kind=DatabricksPathKind.VOLUME,
-                parts=[catalog_name, schema_name, "tmp", "sql", transaction_id]
+                parts=[catalog_name, schema_name, "tmp", "sql", transaction_id],
             ) if temp_volume_path is None else DatabricksPath.parse(obj=temp_volume_path, workspace=connected.workspace)
 
+            logger.debug("Staging Parquet to temp volume: %s", temp_volume_path)
             temp_volume_path.mkdir()
+            temp_volume_path.write_arrow_table(data_tbl)
 
-            temp_volume_path.write_arrow_table(data)
-
-            # get column list from arrow schema
-            columns = [c for c in existing_schema.names]
+            columns = list(existing_schema.names)
             cols_quoted = ", ".join([f"`{c}`" for c in columns])
 
-            statements = []
+            statements: list[str] = []
 
-            # Decide how to ingest
-            # If merge keys provided -> use MERGE
             if match_by:
-                # build ON condition using match_by
-                on_clauses = []
-                for k in match_by:
-                    on_clauses.append(f"T.`{k}` = S.`{k}`")
-                on_condition = " AND ".join(on_clauses)
+                on_condition = " AND ".join([f"T.`{k}` = S.`{k}`" for k in match_by])
 
-                # build UPDATE set (all columns except match_by)
                 update_cols = [c for c in columns if c not in match_by]
                 if update_cols:
                     update_set = ", ".join([f"T.`{c}` = S.`{c}`" for c in update_cols])
                     update_clause = f"WHEN MATCHED THEN UPDATE SET {update_set}"
                 else:
-                    update_clause = ""  # nothing to update
+                    update_clause = ""
 
-                # build INSERT clause
-                insert_clause = f"WHEN NOT MATCHED THEN INSERT ({cols_quoted}) VALUES ({', '.join([f'S.`{c}`' for c in columns])})"
+                insert_clause = (
+                    f"WHEN NOT MATCHED THEN INSERT ({cols_quoted}) "
+                    f"VALUES ({', '.join([f'S.`{c}`' for c in columns])})"
+                )
 
                 merge_sql = f"""MERGE INTO {location} AS T
 USING (
@@ -546,41 +574,47 @@ ON {on_condition}
 {update_clause}
 {insert_clause}"""
                 statements.append(merge_sql)
-
             else:
-                # No match_by -> plain insert
                 if mode.lower() in ("overwrite",):
                     insert_sql = f"""INSERT OVERWRITE {location}
 SELECT {cols_quoted}
 FROM parquet.`{temp_volume_path}`"""
                 else:
-                    # default: append
                     insert_sql = f"""INSERT INTO {location} ({cols_quoted})
 SELECT {cols_quoted}
 FROM parquet.`{temp_volume_path}`"""
                 statements.append(insert_sql)
 
-            # Execute statements (use your existing execute helper)
             try:
                 for stmt in statements:
-                    # trim and run
                     connected.execute(stmt.strip())
             finally:
                 try:
                     temp_volume_path.rmdir(recursive=True)
-                except Exception as e:
-                    logger.warning(e)
+                except Exception:
+                    logger.exception("Failed cleaning temp volume: %s", temp_volume_path)
+
+            logger.info(
+                "Arrow inserted %s rows into %s (mode=%s, match_by=%s, zorder_by=%s)",
+                num_rows,
+                location,
+                mode,
+                match_by,
+                zorder_by,
+            )
 
-            # Optionally run OPTIMIZE / ZORDER / VACUUM if requested (Databricks SQL)
             if zorder_by:
                 zcols = ", ".join([f"`{c}`" for c in zorder_by])
                 optimize_sql = f"OPTIMIZE {location} ZORDER BY ({zcols})"
+                logger.info("Running OPTIMIZE ZORDER BY: %s", zorder_by)
                 connected.execute(optimize_sql)
 
             if optimize_after_merge and match_by:
+                logger.info("Running OPTIMIZE after MERGE")
                 connected.execute(f"OPTIMIZE {location}")
 
             if vacuum_hours is not None:
+                logger.info("Running VACUUM retain=%s hours", vacuum_hours)
                 connected.execute(f"VACUUM {location} RETAIN {vacuum_hours} HOURS")
 
         return None
@@ -596,13 +630,20 @@ FROM parquet.`{temp_volume_path}`"""
         mode: str = "auto",
         cast_options: Optional[CastOptions] = None,
         overwrite_schema: bool | None = None,
-        match_by: list[str] = None,
-        zorder_by: list[str] = None,
+        match_by: Optional[list[str]] = None,
+        zorder_by: Optional[list[str]] = None,
         optimize_after_merge: bool = False,
-        vacuum_hours: int | None = None,  # e.g., 168 for 7 days
+        vacuum_hours: int | None = None,
         spark_options: Optional[Dict[str, Any]] = None,
     ):
-        """Insert a Spark DataFrame into a Delta table with optional merge semantics.
+        """Insert a Spark DataFrame into a Delta table (append/overwrite/merge).
+
+        Behavior:
+        - If the table does not exist: creates it via `saveAsTable(location)` (overwrite).
+        - If `match_by` is provided: uses Delta MERGE for upserts.
+          - If mode == "overwrite": deletes matching keys first, then appends the batch (fast-ish overwrite-by-key).
+          - Else: updates matching rows + inserts new ones.
+        - Else: uses `DataFrameWriter.saveAsTable` with mode.
 
         Args:
             data: Spark DataFrame to insert.
@@ -611,12 +652,12 @@ FROM parquet.`{temp_volume_path}`"""
             schema_name: Optional schema override.
             table_name: Optional table name override.
             mode: Insert mode ("auto", "append", "overwrite").
-            cast_options: Optional casting options.
-            overwrite_schema: Whether to overwrite schema.
-            match_by: Optional merge keys for upserts.
-            zorder_by: Optional Z-ORDER columns.
-            optimize_after_merge: Whether to run OPTIMIZE after merge.
-            vacuum_hours: Optional VACUUM retention window.
+            cast_options: Optional casting options (align to destination schema).
+            overwrite_schema: Whether to overwrite schema on write (when supported).
+            match_by: Merge keys for upserts.
+            zorder_by: Z-ORDER columns (used only if `optimize_after_merge` is True).
+            optimize_after_merge: Whether to run Delta optimize (and z-order) after merge.
+            vacuum_hours: Optional VACUUM retention window in hours.
             spark_options: Optional Spark write options.
 
         Returns:
@@ -627,7 +668,15 @@ FROM parquet.`{temp_volume_path}`"""
             catalog_name=catalog_name,
             schema_name=schema_name,
             table_name=table_name,
-            safe_chars=True
+            safe_chars=True,
+        )
+
+        logger.info(
+            "Spark insert into %s (mode=%s, match_by=%s, overwrite_schema=%s)",
+            location,
+            mode,
+            match_by,
+            overwrite_schema,
         )
 
         spark_options = spark_options if spark_options else {}
@@ -636,11 +685,13 @@ FROM parquet.`{temp_volume_path}`"""
 
         try:
             existing_schema = self.get_table_schema(
-                catalog_name=catalog_name, schema_name=schema_name,
+                catalog_name=catalog_name,
+                schema_name=schema_name,
                 table_name=table_name,
-                to_arrow_schema=False
+                to_arrow_schema=False,
             )
         except ValueError:
+            logger.warning("Destination table missing; creating table %s via overwrite write", location)
             data = convert(data, pyspark.sql.DataFrame)
             data.write.mode("overwrite").options(**spark_options).saveAsTable(location)
             return
@@ -651,29 +702,25 @@ FROM parquet.`{temp_volume_path}`"""
             cast_options = CastOptions.check_arg(options=cast_options, target_field=existing_schema)
             data = cast_spark_dataframe(data, options=cast_options)
 
-        # --- Sanity checks & pre-cleaning (avoid nulls in keys) ---
-        if match_by:
-            notnull: pyspark.sql.Column = None
+        logger.debug("Incoming Spark columns: %s", data.columns)
 
+        if match_by:
+            notnull = None
             for k in match_by:
                 if k not in data.columns:
                     raise ValueError(f"Missing match key '{k}' in DataFrame columns: {data.columns}")
-
-                notnull = data[k].isNotNull() if notnull is None else notnull & (data[k].isNotNull())
+                notnull = data[k].isNotNull() if notnull is None else notnull & data[k].isNotNull()
 
             data = data.filter(notnull)
+            logger.debug("Filtered null keys for match_by=%s", match_by)
 
-        # --- Merge (upsert) ---
         target = self.spark_table(full_name=location)
 
         if match_by:
-            # Build merge condition on the composite key
             cond = " AND ".join([f"t.`{k}` <=> s.`{k}`" for k in match_by])
 
             if mode.casefold() == "overwrite":
                 data = data.cache()
-
-                # Step 1: get unique key combos from source
                 distinct_keys = data.select([f"`{k}`" for k in match_by]).distinct()
 
                 (
@@ -683,35 +730,30 @@ FROM parquet.`{temp_volume_path}`"""
                     .execute()
                 )
 
-                # Step 3: append the clean batch
-                data.write.format("delta").mode("append").saveAsTable(location)
+                data.write.format("delta").mode("append").options(**spark_options).saveAsTable(location)
             else:
                 update_cols = [c for c in data.columns if c not in match_by]
-                set_expr = {
-                    c: F.expr(f"s.`{c}`") for c in update_cols
-                }
+                set_expr = {c: F.expr(f"s.`{c}`") for c in update_cols}
 
-                # Execute MERGE - update matching records first, then insert new ones
                 (
                     target.alias("t")
                     .merge(data.alias("s"), cond)
-                    .whenMatchedUpdate(set=set_expr)  # update matched rows
-                    .whenNotMatchedInsertAll()  # insert new rows
+                    .whenMatchedUpdate(set=set_expr)
+                    .whenNotMatchedInsertAll()
                     .execute()
                 )
         else:
             if mode == "auto":
                 mode = "append"
+            logger.info("Spark write saveAsTable mode=%s", mode)
             data.write.mode(mode).options(**spark_options).saveAsTable(location)
 
-        # --- Optimize: Z-ORDER for faster lookups by composite key (Databricks) ---
         if optimize_after_merge and zorder_by:
-            # pass columns as varargs
+            logger.info("Delta optimize + zorder (%s)", zorder_by)
             target.optimize().executeZOrderBy(*zorder_by)
 
-        # --- Optional VACUUM ---
         if vacuum_hours is not None:
-            # Beware data retention policies; set to a safe value or use default 7 days
+            logger.info("Delta vacuum retain=%s hours", vacuum_hours)
             target.vacuum(vacuum_hours)
 
     def get_table_schema(
@@ -719,24 +761,24 @@ FROM parquet.`{temp_volume_path}`"""
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
-        to_arrow_schema: bool = True
+        to_arrow_schema: bool = True,
     ) -> Union[pa.Field, pa.Schema]:
-        """Fetch a table schema from Unity Catalog as Arrow types.
+        """Fetch a table schema from Unity Catalog and convert it to Arrow types.
 
         Args:
             catalog_name: Optional catalog override.
             schema_name: Optional schema override.
             table_name: Optional table name override.
-            to_arrow_schema: Whether to return an Arrow schema or field.
+            to_arrow_schema: If True returns pa.Schema; else returns a pa.Field(STRUCT<...>).
 
         Returns:
-            Arrow Schema or Field representing the table.
+            Arrow Schema or a STRUCT Field representing the table.
         """
         full_name = self.table_full_name(
             catalog_name=catalog_name,
             schema_name=schema_name,
             table_name=table_name,
-            safe_chars=False
+            safe_chars=False,
         )
 
         wk = self.workspace.sdk()
@@ -746,10 +788,7 @@ FROM parquet.`{temp_volume_path}`"""
         except Exception as e:
             raise ValueError(f"Table %s not found, {type(e)} {e}" % full_name)
 
-        fields = [
-            column_info_to_arrow_field(_)
-            for _ in table.columns
-        ]
+        fields = [column_info_to_arrow_field(_) for _ in table.columns]
 
         if to_arrow_schema:
             return pa.schema(fields, metadata={b"name": table_name})
@@ -762,25 +801,15 @@ FROM parquet.`{temp_volume_path}`"""
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
     ):
-        """Drop a table if it exists.
-
-        Args:
-            location: Fully qualified table name override.
-            catalog_name: Optional catalog override.
-            schema_name: Optional schema override.
-            table_name: Optional table name override.
-
-        Returns:
-            The StatementResult from executing the drop statement.
-        """
+        """Drop a table if it exists."""
         location, _, _, _ = self._check_location_params(
             location=location,
             catalog_name=catalog_name,
             schema_name=schema_name,
             table_name=table_name,
-            safe_chars=True
+            safe_chars=True,
         )
-
+        logger.info("Dropping table if exists: %s", location)
         return self.execute(f"DROP TABLE IF EXISTS {location}")
 
     def create_table(
@@ -797,23 +826,29 @@ FROM parquet.`{temp_volume_path}`"""
         if_not_exists: bool = True,
         optimize_write: bool = True,
         auto_compact: bool = True,
-        execute: bool = True
-    ) -> str:
-        """
-        Generate DDL (Data Definition Language) SQL for creating a table from a PyField schema.
+        execute: bool = True,
+        wait_result: bool = True
+    ) -> Union[str, "StatementResult"]:
+        """Generate (and optionally execute) CREATE TABLE DDL from an Arrow schema/field.
 
         Args:
-            field: PyField schema that defines the table structure
-            table_name: Name of the table to create (defaults to schema.name)
-            catalog_name: Optional catalog name (defaults to "hive_metastore")
-            schema_name: Optional schema name (defaults to "default")
-            partition_by: Optional list of column names to partition the table by
-            comment: Optional table comment
-            options: Optional table properties
-            if_not_exists: Whether to add IF NOT EXISTS clause
+            field: Arrow Field or Schema describing the table. If `field` is a schema, it's converted.
+            location: Fully qualified table name override.
+            table_name: Table name override (used if location not provided).
+            catalog_name: Catalog override.
+            schema_name: Schema override.
+            partition_by: Optional partition columns.
+            cluster_by: If True -> CLUSTER BY AUTO. If list[str] -> CLUSTER BY (..). If False -> no clustering.
+            comment: Optional table comment (falls back to field metadata b"comment" when present).
+            options: Extra table properties.
+            if_not_exists: Add IF NOT EXISTS clause.
+            optimize_write: Sets delta.autoOptimize.optimizeWrite table property.
+            auto_compact: Sets delta.autoOptimize.autoCompact table property.
+            execute: If True, executes DDL and returns StatementResult; otherwise returns SQL string.
+            wait_result: Waits execution to complete
 
         Returns:
-            A SQL string for creating the table
+            StatementResult if execute=True, else the DDL SQL string.
         """
         if not isinstance(field, pa.Field):
             field = convert(field, pa.Field)
@@ -823,7 +858,7 @@ FROM parquet.`{temp_volume_path}`"""
             catalog_name=catalog_name,
             schema_name=schema_name,
             table_name=table_name,
-            safe_chars=True
+            safe_chars=True,
         )
 
         if pa.types.is_struct(field.type):
@@ -831,28 +866,22 @@ FROM parquet.`{temp_volume_path}`"""
         else:
             children = [field]
 
-        # Create the DDL statement
-        column_definitions = [
-            self._field_to_ddl(child)
-            for child in children
-        ]
+        column_definitions = [self._field_to_ddl(child) for child in children]
 
         sql = [
             f"CREATE TABLE {'IF NOT EXISTS ' if if_not_exists else ''}{location} (",
             ",\n  ".join(column_definitions),
-            ")"
+            ")",
         ]
 
-        # Add partition by clause if provided
-        if partition_by and len(partition_by) > 0:
+        if partition_by:
             sql.append(f"\nPARTITIONED BY ({', '.join(partition_by)})")
         elif cluster_by:
             if isinstance(cluster_by, bool):
-                sql.append(f"\nCLUSTER BY AUTO")
+                sql.append("\nCLUSTER BY AUTO")
             else:
                 sql.append(f"\nCLUSTER BY ({', '.join(cluster_by)})")
 
-        # Add comment if provided
         if not comment and field.metadata:
             comment = field.metadata.get(b"comment")
 
@@ -862,32 +891,33 @@ FROM parquet.`{temp_volume_path}`"""
         if comment:
             sql.append(f"\nCOMMENT '{comment}'")
 
-        # Add options if provided
         options = {} if options is None else options
         options.update({
             "delta.autoOptimize.optimizeWrite": optimize_write,
-            "delta.autoOptimize.autoCompact": auto_compact
+            "delta.autoOptimize.autoCompact": auto_compact,
         })
 
         option_strs = []
-
-        if options:
-            for key, value in options.items():
-                if isinstance(value, str):
-                    option_strs.append(f"'{key}' = '{value}'")
-                elif isinstance(value, bool):
-                    b_value = "true" if value else "false"
-                    option_strs.append(f"'{key}' = '{b_value}'")
-                else:
-                    option_strs.append(f"'{key}' = {value}")
+        for key, value in (options or {}).items():
+            if isinstance(value, str):
+                option_strs.append(f"'{key}' = '{value}'")
+            elif isinstance(value, bool):
+                option_strs.append(f"'{key}' = '{'true' if value else 'false'}'")
+            else:
+                option_strs.append(f"'{key}' = {value}")
 
         if option_strs:
             sql.append(f"\nTBLPROPERTIES ({', '.join(option_strs)})")
 
         statement = "\n".join(sql)
 
+        logger.debug(
+            "Generated CREATE TABLE DDL for %s:\n%s",
+            location, statement
+        )
+
         if execute:
-            return self.execute(statement)
+            return self.execute(statement, wait_result=wait_result)
         return statement
 
     def _check_location_params(
@@ -896,28 +926,18 @@ FROM parquet.`{temp_volume_path}`"""
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
-        safe_chars: bool = True
-    ):
-        """Resolve location/catalog/schema/table parameters to a full name.
-
-        Args:
-            location: Fully qualified table name override.
-            catalog_name: Optional catalog override.
-            schema_name: Optional schema override.
-            table_name: Optional table name override.
-            safe_chars: Whether to wrap identifiers in backticks.
-
-        Returns:
-            A tuple of (location, catalog_name, schema_name, table_name).
-        """
+        safe_chars: bool = True,
+    ) -> tuple[str, Optional[str], Optional[str], Optional[str]]:
+        """Resolve (location OR catalog/schema/table) into a fully-qualified name."""
         if location:
             c, s, t = self._catalog_schema_table_names(location)
             catalog_name, schema_name, table_name = catalog_name or c, schema_name or s, table_name or t
 
         location = self.table_full_name(
-            catalog_name=catalog_name, schema_name=schema_name,
+            catalog_name=catalog_name,
+            schema_name=schema_name,
             table_name=table_name,
-            safe_chars=safe_chars
+            safe_chars=safe_chars,
         )
 
         return location, catalog_name or self.catalog_name, schema_name or self.schema_name, table_name
@@ -927,98 +947,68 @@ FROM parquet.`{temp_volume_path}`"""
         field: pa.Field,
         put_name: bool = True,
         put_not_null: bool = True,
-        put_comment: bool = True
+        put_comment: bool = True,
     ) -> str:
-        """
-        Convert a PyField to a DDL column definition.
-
-        Args:
-            field: The PyField to convert
-
-        Returns:
-            A string containing the column definition in DDL format
-        """
+        """Convert an Arrow Field to a Databricks SQL column DDL fragment."""
         name = field.name
         nullable_str = " NOT NULL" if put_not_null and not field.nullable else ""
         name_str = f"{name} " if put_name else ""
 
-        # Get comment if available
         comment_str = ""
         if put_comment and field.metadata and b"comment" in field.metadata:
             comment = field.metadata[b"comment"].decode("utf-8")
             comment_str = f" COMMENT '{comment}'"
 
-        # Handle primitive types
         if not pa.types.is_nested(field.type):
             sql_type = SQLEngine._arrow_to_sql_type(field.type)
             return f"{name_str}{sql_type}{nullable_str}{comment_str}"
 
-        # Handle struct type
         if pa.types.is_struct(field.type):
             child_defs = [SQLEngine._field_to_ddl(child) for child in field.type]
             struct_body = ", ".join(child_defs)
             return f"{name_str}STRUCT<{struct_body}>{nullable_str}{comment_str}"
 
-        # Handle map type
         if pa.types.is_map(field.type):
             map_type: pa.MapType = field.type
             key_type = SQLEngine._field_to_ddl(map_type.key_field, put_name=False, put_comment=False, put_not_null=False)
             val_type = SQLEngine._field_to_ddl(map_type.item_field, put_name=False, put_comment=False, put_not_null=False)
             return f"{name_str}MAP<{key_type}, {val_type}>{nullable_str}{comment_str}"
 
-        # Handle list type after map
         if pa.types.is_list(field.type) or pa.types.is_large_list(field.type):
             list_type: pa.ListType = field.type
             elem_type = SQLEngine._field_to_ddl(list_type.value_field, put_name=False, put_comment=False, put_not_null=False)
             return f"{name_str}ARRAY<{elem_type}>{nullable_str}{comment_str}"
 
-        # Default fallback to string for unknown types
         raise TypeError(f"Cannot make ddl field from {field}")
 
     @staticmethod
-    def _arrow_to_sql_type(
-        arrow_type: Union[pa.DataType, pa.Decimal128Type]
-    ) -> str:
-        """
-        Convert an Arrow data type to SQL data type.
-
-        Args:
-            arrow_type: The Arrow data type
-
-        Returns:
-            A string containing the SQL data type
-        """
+    def _arrow_to_sql_type(arrow_type: Union[pa.DataType, pa.Decimal128Type]) -> str:
+        """Convert an Arrow data type to a Databricks SQL type string."""
         if pa.types.is_boolean(arrow_type):
             return "BOOLEAN"
-        elif pa.types.is_int8(arrow_type):
+        if pa.types.is_int8(arrow_type):
             return "TINYINT"
-        elif pa.types.is_int16(arrow_type):
+        if pa.types.is_int16(arrow_type):
             return "SMALLINT"
-        elif pa.types.is_int32(arrow_type):
+        if pa.types.is_int32(arrow_type):
             return "INT"
-        elif pa.types.is_int64(arrow_type):
+        if pa.types.is_int64(arrow_type):
             return "BIGINT"
-        elif pa.types.is_float32(arrow_type):
+        if pa.types.is_float32(arrow_type):
             return "FLOAT"
-        elif pa.types.is_float64(arrow_type):
+        if pa.types.is_float64(arrow_type):
             return "DOUBLE"
-        elif is_arrow_type_string_like(arrow_type):
+        if is_arrow_type_string_like(arrow_type):
             return "STRING"
-        elif is_arrow_type_binary_like(arrow_type):
+        if is_arrow_type_binary_like(arrow_type):
             return "BINARY"
-        elif pa.types.is_timestamp(arrow_type):
+        if pa.types.is_timestamp(arrow_type):
             tz = getattr(arrow_type, "tz", None)
-
-            if tz:
-                return "TIMESTAMP"
-            return "TIMESTAMP_NTZ"
-        elif pa.types.is_date(arrow_type):
+            return "TIMESTAMP" if tz else "TIMESTAMP_NTZ"
+        if pa.types.is_date(arrow_type):
             return "DATE"
-        elif pa.types.is_decimal(arrow_type):
-            precision = arrow_type.precision
-            scale = arrow_type.scale
-            return f"DECIMAL({precision}, {scale})"
-        elif pa.types.is_null(arrow_type):
+        if pa.types.is_decimal(arrow_type):
+            return f"DECIMAL({arrow_type.precision}, {arrow_type.scale})"
+        if pa.types.is_null(arrow_type):
             return "STRING"
-        else:
-            raise ValueError(f"Cannot make ddl type for {arrow_type}")
+        raise ValueError(f"Cannot make ddl type for {arrow_type}")