fastflowtransform.executors.snowflake_snowpark

SnowflakeSnowparkExecutor

Bases: SqlIdentifierMixin, SnapshotSqlMixin, BaseExecutor[SNDF]

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

class SnowflakeSnowparkExecutor(SqlIdentifierMixin, SnapshotSqlMixin, BaseExecutor[SNDF]):
    ENGINE_NAME = "snowflake_snowpark"
    """Snowflake executor operating on Snowpark DataFrames (no pandas)."""
    _BUDGET_GUARD = BudgetGuard(
        env_var="FF_SF_MAX_BYTES",
        estimator_attr="_estimate_query_bytes",
        engine_label="Snowflake",
        what="query",
    )

    def __init__(self, cfg: dict):
        # cfg: {account, user, password, warehouse, database, schema, role?}
        self.session = Session.builder.configs(cfg).create()
        self.database = cfg["database"]
        self.schema = cfg["schema"]

        self.allow_create_schema: bool = bool(cfg["allow_create_schema"])
        self._ensure_schema()

        # Provide a tiny testing shim so tests can call executor.con.execute("SQL")
        self.con = _SFCursorShim(self.session)

    # ---------- Cost estimation & central execution ----------

    def _estimate_query_bytes(self, sql: str) -> int | None:
        """
        Best-effort Snowflake bytes estimation.

        Uses `EXPLAIN USING TEXT` and tries to extract a "bytes=<n>"-style
        metric from the textual plan. If parsing fails or Snowflake doesn't
        expose such info, returns None and the guard is effectively disabled.
        """
        try:
            body = self._selectable_body(sql)
        except Exception:
            body = sql

        try:
            rows = self.session.sql(f"EXPLAIN USING JSON {body}").collect()
            if not rows:
                return None

            parts: list[str] = []
            for r in rows:
                try:
                    parts.append(str(r[0]))
                except Exception:
                    as_dict: dict[str, Any] = getattr(r, "asDict", lambda: {})()
                    if as_dict:
                        parts.extend(str(v) for v in as_dict.values())

            plan_text = "\n".join(parts).strip()
            if not plan_text:
                return None

            try:
                plan_data = json.loads(plan_text)
            except Exception:
                return None

            bytes_val = self._extract_bytes_from_plan(plan_data)
            if bytes_val is None or bytes_val <= 0:
                return None
            return bytes_val
        except Exception:
            # Any parsing / EXPLAIN issues → no estimate, guard skipped
            return None

    def _extract_bytes_from_plan(self, plan_data: Any) -> int | None:
        def _to_int(value: Any) -> int | None:
            if value is None:
                return None
            try:
                return int(value)
            except Exception:
                return None

        if isinstance(plan_data, dict):
            global_stats = plan_data.get("GlobalStats") or plan_data.get("globalStats")
            if isinstance(global_stats, dict):
                candidate = _to_int(
                    global_stats.get("bytesAssigned") or global_stats.get("bytes_assigned")
                )
                if candidate:
                    return candidate
            for val in plan_data.values():
                bytes_val = self._extract_bytes_from_plan(val)
                if bytes_val:
                    return bytes_val
        elif isinstance(plan_data, list):
            for item in plan_data:
                bytes_val = self._extract_bytes_from_plan(item)
                if bytes_val:
                    return bytes_val
        return None

    def _execute_sql(self, sql: str) -> SNDF:
        """
        Central Snowflake SQL runner.

        - Returns a Snowpark DataFrame (same as session.sql).
        - Records best-effort query stats for run_results.json.
        """

        def _exec() -> SNDF:
            return self.session.sql(sql)

        return run_sql_with_budget(
            self,
            sql,
            guard=self._BUDGET_GUARD,
            exec_fn=_exec,
            estimate_fn=self._estimate_query_bytes,
        )

    def _exec_many(self, sql: str) -> None:
        """
        Execute multiple SQL statements separated by ';' on the same connection.
        Snowflake normally accepts one statement per execute(), so we split here.
        """
        for stmt in (part.strip() for part in sql.split(";")):
            if not stmt:
                continue
            self._execute_sql(stmt).collect()

    # ---------- Helpers ----------
    def _q(self, s: str) -> str:
        return '"' + s.replace('"', '""') + '"'

    def _quote_identifier(self, ident: str) -> str:
        # Keep identifiers unquoted to match legacy Snowflake behaviour.
        return ident

    def _default_schema(self) -> str | None:
        return self.schema

    def _default_catalog(self) -> str | None:
        return self.database

    def _should_include_catalog(
        self, catalog: str | None, schema: str | None, *, explicit: bool
    ) -> bool:
        # Always include database when present; Snowflake expects DB.SCHEMA.TABLE.
        return bool(catalog)

    def _qualified(self, rel: str) -> str:
        # DATABASE.SCHEMA.TABLE  (no quotes)
        return self._qualify_identifier(rel, quote=False)

    def _ensure_schema(self) -> None:
        """
        Best-effort schema creation when allow_create_schema=True.

        Mirrors BigQuery's `_ensure_dataset` behaviour:
        - If the flag is false → do nothing.
        - If true → `CREATE SCHEMA IF NOT EXISTS "DB"."SCHEMA"`.
        """
        if not getattr(self, "allow_create_schema", False):
            return
        if not self.database or not self.schema:
            # Misconfigured; let downstream errors surface naturally.
            return

        db = self._q(self.database)
        sch = self._q(self.schema)
        with suppress(Exception):
            # Fully qualified CREATE SCHEMA is allowed in Snowflake.
            self.session.sql(f"CREATE SCHEMA IF NOT EXISTS {db}.{sch}").collect()
            # Best-effort; permission issues or race conditions shouldn't crash the executor.
            # If the schema truly doesn't exist and we can't create it, later queries will fail
            # with a clearer engine error.

    # ---------- Frame-Hooks ----------
    def _read_relation(self, relation: str, node: Node, deps: Iterable[str]) -> SNDF:
        df = self.session.table(self._qualified(relation))
        # Present a *logical* lowercase schema to Python models:
        lowered = [c.lower() for c in df.schema.names]
        return df.toDF(*lowered)

    def _materialize_relation(self, relation: str, df: SNDF, node: Node) -> None:
        if not self._is_frame(df):
            raise TypeError("Snowpark model must return a Snowpark DataFrame")

        # Normalize to uppercase for storage in Snowflake
        cols = list(df.schema.names)
        upper_cols = [c.upper() for c in cols]
        if cols != upper_cols:
            df = df.toDF(*upper_cols)

        start = perf_counter()
        df.write.save_as_table(self._qualified(relation), mode="overwrite")
        duration_ms = int((perf_counter() - start) * 1000)
        bytes_est = self._estimate_frame_bytes(df)
        self._record_query_stats(
            QueryStats(
                bytes_processed=bytes_est,
                rows=None,
                duration_ms=duration_ms,
            )
        )

    def _estimate_frame_bytes(self, df: SNDF) -> int | None:
        """
        Best-effort bytes estimate for a Snowpark DataFrame.

        Strategy:
        1) Use DataFrame.queries["queries"] (public Snowpark API) to get SQL.
        2) Optionally fall back to df._plan.sql() if queries is missing/empty.
        3) Run our existing _estimate_query_bytes(sql_text).
        """
        try:
            sql_text = self._snowpark_df_sql(df)
            if not isinstance(sql_text, str) or not sql_text.strip():
                return None
            return self._estimate_query_bytes(sql_text)
        except Exception:
            return None

    def _snowpark_df_sql(self, df: Any) -> str | None:
        """
        Extract the main SQL statement for a Snowpark DataFrame.

        Uses the documented public APIs:
        - DataFrame.queries -> {"queries": [sql1, sql2, ...], "post_actions": [...]}
        - Optionally falls back to df._plan.sql() if needed.
        """
        # 1) Primary source: DataFrame.queries
        queries_dict = getattr(df, "queries", None)

        if isinstance(queries_dict, dict):
            queries = queries_dict.get("queries")
            if isinstance(queries, list) and queries:
                # Pick the most likely "main" query.
                # Snowflake examples use queries['queries'][0],
                # but we can be a bit safer and pick the longest non-empty SQL.
                candidates = [q.strip() for q in queries if isinstance(q, str) and q.strip()]
                if candidates:
                    # Heuristic: longest SQL string is usually the main SELECT/CTE.
                    return max(candidates, key=len)

        # 2) Fallback: internal plan (undocumented but widely used)
        plan = getattr(df, "_plan", None)
        if plan is not None:
            # Prefer simplified plan if available
            with suppress(Exception):
                simplify = getattr(plan, "simplify", None)
                if callable(simplify):
                    simplified = simplify()
                    to_sql = getattr(simplified, "sql", None)
                    if callable(to_sql):
                        sql = to_sql()
                        if isinstance(sql, str) and sql.strip():
                            return sql.strip()

            # Raw plan.sql()
            with suppress(Exception):
                to_sql = getattr(plan, "sql", None)
                if callable(to_sql):
                    sql = to_sql()
                    if isinstance(sql, str) and sql.strip():
                        return sql.strip()

        return None

    def _create_view_over_table(self, view_name: str, backing_table: str, node: Node) -> None:
        qv = self._qualified(view_name)
        qb = self._qualified(backing_table)
        self._execute_sql(f"CREATE OR REPLACE VIEW {qv} AS SELECT * FROM {qb}").collect()

    def _validate_required(
        self, node_name: str, inputs: Any, requires: dict[str, set[str]]
    ) -> None:
        if not requires:
            return

        def cols(df: SNDF) -> set[str]:
            # Compare in lowercase to be case-insensitive for Snowflake
            return {c.lower() for c in df.schema.names}

        # Normalize the required sets too
        normalized_requires = {rel: {c.lower() for c in needed} for rel, needed in requires.items()}

        errors: list[str] = []

        if isinstance(inputs, SNDF):
            need = next(iter(normalized_requires.values()), set())
            missing = need - cols(inputs)
            if missing:
                errors.append(f"- missing columns: {sorted(missing)} | have={sorted(cols(inputs))}")
        else:
            for rel, need in normalized_requires.items():
                if rel not in inputs:
                    errors.append(f"- missing dependency key '{rel}'")
                    continue
                missing = need - cols(inputs[rel])
                if missing:
                    errors.append(
                        f"- [{rel}] missing: {sorted(missing)} | have={sorted(cols(inputs[rel]))}"
                    )

        if errors:
            raise ValueError(
                "Required columns check failed for Snowpark model "
                f"'{node_name}'.\n" + "\n".join(errors)
            )

    def _columns_of(self, frame: SNDF) -> list[str]:
        return list(frame.schema.names)

    def _is_frame(self, obj: Any) -> bool:
        # Accept real Snowpark DataFrames and test doubles with a compatible surface.
        schema = getattr(obj, "schema", None)
        return isinstance(obj, SNDF) or (
            schema is not None
            and hasattr(schema, "names")
            and callable(getattr(obj, "collect", None))
        )

    def _frame_name(self) -> str:
        return "Snowpark"

    # ---- SQL hooks ----
    def _this_identifier(self, node: Node) -> str:
        """
        Identifier for {{ this }} in SQL models.
        Use fully-qualified DB.SCHEMA.TABLE so all build/read/test paths agree.
        """
        return self._qualify_identifier(relation_for(node.name), quote=False)

    def _format_source_reference(
        self, cfg: dict[str, Any], source_name: str, table_name: str
    ) -> str:
        if cfg.get("location"):
            raise NotImplementedError("Snowflake executor does not support path-based sources.")

        ident = cfg.get("identifier")
        if not ident:
            raise KeyError(f"Source {source_name}.{table_name} missing identifier")

        sch = self._pick_schema(cfg)
        db = self._pick_catalog(cfg, sch)
        if not db or not sch:
            raise KeyError(
                f"Source {source_name}.{table_name} missing database/schema for Snowflake"
            )
        return self._qualify_identifier(ident, schema=sch, catalog=db, quote=False)

    def _create_or_replace_view(self, target_sql: str, select_body: str, node: Node) -> None:
        self._execute_sql(f"CREATE OR REPLACE VIEW {target_sql} AS {select_body}").collect()

    def _create_or_replace_table(self, target_sql: str, select_body: str, node: Node) -> None:
        self._execute_sql(f"CREATE OR REPLACE TABLE {target_sql} AS {select_body}").collect()

    def _create_or_replace_view_from_table(
        self, view_name: str, backing_table: str, node: Node
    ) -> None:
        view_id = self._qualified(view_name)
        back_id = self._qualified(backing_table)
        self._execute_sql(f"CREATE OR REPLACE VIEW {view_id} AS SELECT * FROM {back_id}").collect()

    def _format_test_table(self, table: str | None) -> str | None:
        formatted = super()._format_test_table(table)
        if formatted is None:
            return None

        # If it's already qualified (DB.SCHEMA.TABLE) or quoted, leave it alone.
        if "." in formatted or '"' in formatted:
            return formatted

        # Otherwise, treat it as a logical relation name and fully-qualify it
        # with the executor's configured database/schema.
        return self._qualified(formatted)

    # ---- Meta hook ----
    def on_node_built(self, node: Node, relation: str, fingerprint: str) -> None:
        """After successful materialization, upsert _ff_meta (best-effort)."""
        ensure_meta_table(self)
        upsert_meta(self, node.name, relation, fingerprint, "snowflake_snowpark")

    # ── Incremental API (parity with DuckDB/PG) ──────────────────────────
    def exists_relation(self, relation: str) -> bool:
        """Check existence via information_schema.tables."""
        db = self._q(self.database)
        schema_lit = f"'{self.schema.upper()}'"
        rel_lit = f"'{relation.upper()}'"
        q = f"""
        select 1
        from {db}.information_schema.tables
        where upper(table_schema) = {schema_lit}
            and upper(table_name) = {rel_lit}
        limit 1
        """
        try:
            return bool(self._execute_sql(q).collect())
        except Exception:
            return False

    def create_table_as(self, relation: str, select_sql: str) -> None:
        body = self._selectable_body(select_sql).strip().rstrip(";\n\t ")
        self._execute_sql(
            f"CREATE OR REPLACE TABLE {self._qualified(relation)} AS {body}"
        ).collect()

    def full_refresh_table(self, relation: str, select_sql: str) -> None:
        """
        Engine-specific full refresh for incremental fallbacks.
        """
        body = self._selectable_body(select_sql).strip().rstrip(";\n\t ")
        self._execute_sql(
            f"CREATE OR REPLACE TABLE {self._qualified(relation)} AS {body}"
        ).collect()

    def incremental_insert(self, relation: str, select_sql: str) -> None:
        body = self._selectable_body(select_sql).strip().rstrip(";\n\t ")
        self._execute_sql(f"INSERT INTO {self._qualified(relation)} {body}").collect()

    def incremental_merge(self, relation: str, select_sql: str, unique_key: list[str]) -> None:
        body = self._selectable_body(select_sql).strip().rstrip(";\n\t ")
        pred = " AND ".join([f"t.{k}=s.{k}" for k in unique_key]) or "FALSE"
        qrel = self._qualified(relation)

        # 1) Delete matching keys
        delete_sql = f"""
        DELETE FROM {qrel} AS t
        USING ({body}) AS s
        WHERE {pred}
        """
        self._execute_sql(delete_sql).collect()

        # 2) Insert all rows from the delta
        insert_sql = f"INSERT INTO {qrel} SELECT * FROM ({body})"
        self._execute_sql(insert_sql).collect()

    def alter_table_sync_schema(
        self, relation: str, select_sql: str, *, mode: str = "append_new_columns"
    ) -> None:
        """
        Best-effort additive schema sync:
        - infer SELECT schema via LIMIT 0
        - add missing columns as STRING
        """
        if mode not in {"append_new_columns", "sync_all_columns"}:
            return

        qrel = self._qualified(relation)

        # Use identifiers in FROM, but *string literals* in WHERE
        db_ident = self._q(self.database)
        schema_lit = self.schema.replace("'", "''")
        rel_lit = relation.replace("'", "''")

        try:
            existing = {
                r[0]
                for r in self._execute_sql(
                    f"""
                    select column_name
                    from {db_ident}.information_schema.columns
                    where upper(table_schema) = upper('{schema_lit}')
                    and upper(table_name)   = upper('{rel_lit}')
                    """
                ).collect()
            }
        except Exception:
            existing = set()

        # Probe SELECT columns
        body = self._first_select_body(select_sql).strip().rstrip(";\n\t ")
        probe = self.session.sql(f"SELECT * FROM ({body}) q WHERE 1=0")
        probe_cols = list(probe.schema.names)

        to_add = [c for c in probe_cols if c not in existing]
        if not to_add:
            return

        # Column names are identifiers → _q is correct here
        cols_sql = ", ".join(f"{self._q(c)} STRING" for c in to_add)
        self._execute_sql(f"ALTER TABLE {qrel} ADD COLUMN {cols_sql}").collect()

    # ---- Snapshot API (mixin hooks) --------------------------------------
    def _snapshot_target_identifier(self, rel_name: str) -> str:
        return self._qualified(rel_name)

    def _snapshot_current_timestamp(self) -> str:
        return "CURRENT_TIMESTAMP()"

    def _snapshot_create_keyword(self) -> str:
        return "CREATE OR REPLACE TABLE"

    def _snapshot_null_timestamp(self) -> str:
        return "CAST(NULL AS TIMESTAMP)"

    def _snapshot_null_hash(self) -> str:
        return "CAST(NULL AS VARCHAR)"

    def _snapshot_hash_expr(self, check_cols: list[str], src_alias: str) -> str:
        concat_expr = self._snapshot_concat_expr(check_cols, src_alias)
        return f"CAST(MD5({concat_expr}) AS VARCHAR)"

    def _snapshot_cast_as_string(self, expr: str) -> str:
        return f"CAST({expr} AS VARCHAR)"

    def _snapshot_source_ref(
        self, rel_name: str, select_body: str
    ) -> tuple[str, Callable[[], None]]:
        src_name = f"__ff_snapshot_src_{rel_name}".replace(".", "_")
        src_quoted = self._q(src_name)
        self._execute_sql(
            f"CREATE OR REPLACE TEMPORARY VIEW {src_quoted} AS {select_body}"
        ).collect()

        def _cleanup() -> None:
            self._execute_sql(f"DROP VIEW IF EXISTS {src_quoted}").collect()

        return src_quoted, _cleanup

    def execute_hook_sql(self, sql: str) -> None:
        """
        Execute one SQL statement for pre/post/on_run hooks.
        """
        self._exec_many(sql)

    # ---- Unit-test helpers -----------------------------------------------

    def utest_read_relation(self, relation: str) -> pd.DataFrame:
        """
        Read a relation into a pandas DataFrame for unit-test assertions.

        We use Snowpark to read the table and convert to pandas,
        normalizing column names to lowercase to match _read_relation.
        """
        df = self.session.table(self._qualified(relation))
        # Mirror _read_relation: present lowercase schema to the test layer
        lowered = [c.lower() for c in df.schema.names]
        df = df.toDF(*lowered)

        to_pandas = getattr(df, "to_pandas", None)

        pdf: pd.DataFrame
        if callable(to_pandas):
            pdf = cast(pd.DataFrame, to_pandas())
        else:
            rows = df.collect()
            records = [r.asDict() for r in rows]
            pdf = pd.DataFrame.from_records(records)

        # Return a new DF with lowercase columns (no attribute assignment)
        return pdf.rename(columns=lambda c: str(c).lower())

    def utest_load_relation_from_rows(self, relation: str, rows: list[dict]) -> None:
        """
        Load rows into a Snowflake table for unit tests (replace if exists).

        We build a Snowpark DataFrame from the Python rows and overwrite the
        target table using save_as_table().
        """
        # Best-effort: if rows are empty, create an empty table with no rows.
        # We assume at least one row in normal test usage so we can infer schema.
        if not rows:
            # Without any rows we don't know the schema; create a trivial
            # single-column table to surface the situation clearly.
            tmp_df = self.session.create_dataframe([[None]], schema=["__empty__"])
            tmp_df.write.save_as_table(self._qualified(relation), mode="overwrite")
            return

        # Infer column order from the first row
        first = rows[0]
        columns = list(first.keys())

        # Normalize data to a list of lists in a fixed column order
        data = [[row.get(col) for col in columns] for row in rows]

        df = self.session.create_dataframe(data, schema=columns)

        # Store with uppercase column names in Snowflake (conventional)
        upper_cols = [c.upper() for c in columns]
        if columns != upper_cols:
            df = df.toDF(*upper_cols)

        # Overwrite the target table
        df.write.save_as_table(self._qualified(relation), mode="overwrite")

    def utest_clean_target(self, relation: str) -> None:
        """
        For unit tests: drop any table or view with this name in the configured
        database/schema.

        We:
          - try DROP VIEW IF EXISTS DB.SCHEMA.REL
          - try DROP TABLE IF EXISTS DB.SCHEMA.REL

        and ignore "not a view/table" style errors so it doesn't matter what
        kind of object is currently there - after this, nothing with that name
        should remain (best-effort).
        """
        qualified = self._qualified(relation)

        # Drop view first; ignore errors if it's actually a table or doesn't exist.
        with suppress(Exception):
            self.session.sql(f"DROP VIEW IF EXISTS {qualified}").collect()

        # Then drop table; ignore errors if it's actually a view or doesn't exist.
        with suppress(Exception):
            self.session.sql(f"DROP TABLE IF EXISTS {qualified}").collect()

ENGINE_NAME class-attribute instance-attribute

ENGINE_NAME = 'snowflake_snowpark'

Snowflake executor operating on Snowpark DataFrames (no pandas).

on_node_built

on_node_built(node, relation, fingerprint)

After successful materialization, upsert _ff_meta (best-effort).

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def on_node_built(self, node: Node, relation: str, fingerprint: str) -> None:
    """After successful materialization, upsert _ff_meta (best-effort)."""
    ensure_meta_table(self)
    upsert_meta(self, node.name, relation, fingerprint, "snowflake_snowpark")

exists_relation

exists_relation(relation)

Check existence via information_schema.tables.

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def exists_relation(self, relation: str) -> bool:
    """Check existence via information_schema.tables."""
    db = self._q(self.database)
    schema_lit = f"'{self.schema.upper()}'"
    rel_lit = f"'{relation.upper()}'"
    q = f"""
    select 1
    from {db}.information_schema.tables
    where upper(table_schema) = {schema_lit}
        and upper(table_name) = {rel_lit}
    limit 1
    """
    try:
        return bool(self._execute_sql(q).collect())
    except Exception:
        return False

full_refresh_table

full_refresh_table(relation, select_sql)

Engine-specific full refresh for incremental fallbacks.

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def full_refresh_table(self, relation: str, select_sql: str) -> None:
    """
    Engine-specific full refresh for incremental fallbacks.
    """
    body = self._selectable_body(select_sql).strip().rstrip(";\n\t ")
    self._execute_sql(
        f"CREATE OR REPLACE TABLE {self._qualified(relation)} AS {body}"
    ).collect()

alter_table_sync_schema

alter_table_sync_schema(relation, select_sql, *, mode='append_new_columns')

Best-effort additive schema sync: - infer SELECT schema via LIMIT 0 - add missing columns as STRING

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def alter_table_sync_schema(
    self, relation: str, select_sql: str, *, mode: str = "append_new_columns"
) -> None:
    """
    Best-effort additive schema sync:
    - infer SELECT schema via LIMIT 0
    - add missing columns as STRING
    """
    if mode not in {"append_new_columns", "sync_all_columns"}:
        return

    qrel = self._qualified(relation)

    # Use identifiers in FROM, but *string literals* in WHERE
    db_ident = self._q(self.database)
    schema_lit = self.schema.replace("'", "''")
    rel_lit = relation.replace("'", "''")

    try:
        existing = {
            r[0]
            for r in self._execute_sql(
                f"""
                select column_name
                from {db_ident}.information_schema.columns
                where upper(table_schema) = upper('{schema_lit}')
                and upper(table_name)   = upper('{rel_lit}')
                """
            ).collect()
        }
    except Exception:
        existing = set()

    # Probe SELECT columns
    body = self._first_select_body(select_sql).strip().rstrip(";\n\t ")
    probe = self.session.sql(f"SELECT * FROM ({body}) q WHERE 1=0")
    probe_cols = list(probe.schema.names)

    to_add = [c for c in probe_cols if c not in existing]
    if not to_add:
        return

    # Column names are identifiers → _q is correct here
    cols_sql = ", ".join(f"{self._q(c)} STRING" for c in to_add)
    self._execute_sql(f"ALTER TABLE {qrel} ADD COLUMN {cols_sql}").collect()

execute_hook_sql

execute_hook_sql(sql)

Execute one SQL statement for pre/post/on_run hooks.

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def execute_hook_sql(self, sql: str) -> None:
    """
    Execute one SQL statement for pre/post/on_run hooks.
    """
    self._exec_many(sql)

utest_read_relation

utest_read_relation(relation)

Read a relation into a pandas DataFrame for unit-test assertions.

We use Snowpark to read the table and convert to pandas, normalizing column names to lowercase to match _read_relation.

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def utest_read_relation(self, relation: str) -> pd.DataFrame:
    """
    Read a relation into a pandas DataFrame for unit-test assertions.

    We use Snowpark to read the table and convert to pandas,
    normalizing column names to lowercase to match _read_relation.
    """
    df = self.session.table(self._qualified(relation))
    # Mirror _read_relation: present lowercase schema to the test layer
    lowered = [c.lower() for c in df.schema.names]
    df = df.toDF(*lowered)

    to_pandas = getattr(df, "to_pandas", None)

    pdf: pd.DataFrame
    if callable(to_pandas):
        pdf = cast(pd.DataFrame, to_pandas())
    else:
        rows = df.collect()
        records = [r.asDict() for r in rows]
        pdf = pd.DataFrame.from_records(records)

    # Return a new DF with lowercase columns (no attribute assignment)
    return pdf.rename(columns=lambda c: str(c).lower())

utest_load_relation_from_rows

utest_load_relation_from_rows(relation, rows)

Load rows into a Snowflake table for unit tests (replace if exists).

We build a Snowpark DataFrame from the Python rows and overwrite the target table using save_as_table().

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def utest_load_relation_from_rows(self, relation: str, rows: list[dict]) -> None:
    """
    Load rows into a Snowflake table for unit tests (replace if exists).

    We build a Snowpark DataFrame from the Python rows and overwrite the
    target table using save_as_table().
    """
    # Best-effort: if rows are empty, create an empty table with no rows.
    # We assume at least one row in normal test usage so we can infer schema.
    if not rows:
        # Without any rows we don't know the schema; create a trivial
        # single-column table to surface the situation clearly.
        tmp_df = self.session.create_dataframe([[None]], schema=["__empty__"])
        tmp_df.write.save_as_table(self._qualified(relation), mode="overwrite")
        return

    # Infer column order from the first row
    first = rows[0]
    columns = list(first.keys())

    # Normalize data to a list of lists in a fixed column order
    data = [[row.get(col) for col in columns] for row in rows]

    df = self.session.create_dataframe(data, schema=columns)

    # Store with uppercase column names in Snowflake (conventional)
    upper_cols = [c.upper() for c in columns]
    if columns != upper_cols:
        df = df.toDF(*upper_cols)

    # Overwrite the target table
    df.write.save_as_table(self._qualified(relation), mode="overwrite")

utest_clean_target

utest_clean_target(relation)

For unit tests: drop any table or view with this name in the configured database/schema.

We

• try DROP VIEW IF EXISTS DB.SCHEMA.REL
• try DROP TABLE IF EXISTS DB.SCHEMA.REL

and ignore "not a view/table" style errors so it doesn't matter what kind of object is currently there - after this, nothing with that name should remain (best-effort).

Source code in src/fastflowtransform/executors/snowflake_snowpark.py

def utest_clean_target(self, relation: str) -> None:
    """
    For unit tests: drop any table or view with this name in the configured
    database/schema.

    We:
      - try DROP VIEW IF EXISTS DB.SCHEMA.REL
      - try DROP TABLE IF EXISTS DB.SCHEMA.REL

    and ignore "not a view/table" style errors so it doesn't matter what
    kind of object is currently there - after this, nothing with that name
    should remain (best-effort).
    """
    qualified = self._qualified(relation)

    # Drop view first; ignore errors if it's actually a table or doesn't exist.
    with suppress(Exception):
        self.session.sql(f"DROP VIEW IF EXISTS {qualified}").collect()

    # Then drop table; ignore errors if it's actually a view or doesn't exist.
    with suppress(Exception):
        self.session.sql(f"DROP TABLE IF EXISTS {qualified}").collect()

run_sql

run_sql(node, env)

Orchestrate SQL models

1) Render Jinja (ref/source/this) and strip leading {{ config(...) }}. 2) If the SQL is full DDL (CREATE …), execute it verbatim (passthrough). 3) Otherwise, normalize to CREATE OR REPLACE {TABLE|VIEW} AS . The body is CTE-aware (keeps WITH … SELECT … intact).

On failure, raise ModelExecutionError with a helpful snippet.

Source code in src/fastflowtransform/executors/base.py

def run_sql(self, node: Node, env: Environment) -> None:
    """
    Orchestrate SQL models:
      1) Render Jinja (ref/source/this) and strip leading {{ config(...) }}.
      2) If the SQL is full DDL (CREATE …), execute it verbatim (passthrough).
      3) Otherwise, normalize to CREATE OR REPLACE {TABLE|VIEW} AS <body>.
         The body is CTE-aware (keeps WITH … SELECT … intact).
    On failure, raise ModelExecutionError with a helpful snippet.
    """
    meta = getattr(node, "meta", {}) or {}
    if self._meta_is_incremental(meta):
        # Delegates to incremental engine: render, schema sync, merge/insert, etc.
        return _ff_incremental.run_or_dispatch(self, node, env)

    if self._meta_is_snapshot(meta):
        # Snapshots are executed via the dedicated CLI: `fft snapshot run`.
        raise ModelExecutionError(
            node_name=node.name,
            relation=relation_for(node.name),
            message=(
                "Snapshot models cannot be executed via 'fft run'. "
                "Use 'fft snapshot run' instead."
            ),
            sql_snippet="",
        )

    sql_rendered = self.render_sql(
        node,
        env,
        ref_resolver=lambda name: self._resolve_ref(name, env),
        source_resolver=self._resolve_source,
    )
    sql = self._strip_leading_config(sql_rendered).strip()

    materialization = (node.meta or {}).get("materialized", "table")
    if materialization == "ephemeral":
        return

    # 1) Direct DDL passthrough (CREATE [OR REPLACE] {TABLE|VIEW} …)
    if self._looks_like_direct_ddl(sql):
        try:
            self._execute_sql_direct(sql, node)
            return
        except NotImplementedError:
            # Engine doesn't implement direct DDL → fall back to normalized materialization.
            pass
        except Exception as e:
            raise ModelExecutionError(
                node_name=node.name,
                relation=relation_for(node.name),
                message=str(e),
                sql_snippet=sql,
            ) from e

    # 2) Normalized materialization path (CTE-safe body)
    body = self._selectable_body(sql).rstrip(" ;\n\t")
    target_sql = self._format_relation_for_ref(node.name)

    # Centralized SQL preview logging (applies to ALL engines)
    preview = (
        f"=== MATERIALIZE ===\n"
        f"-- model: {node.name}\n"
        f"-- materialized: {materialization}\n"
        f"-- target: {target_sql}\n"
        f"{body}\n"
    )
    echo_debug(preview)

    try:
        self._apply_sql_materialization(node, target_sql, body, materialization)
    except Exception as e:
        preview = f"-- materialized={materialization}\n-- target={target_sql}\n{body}"
        raise ModelExecutionError(
            node_name=node.name,
            relation=relation_for(node.name),
            message=str(e),
            sql_snippet=preview,
        ) from e

configure_query_budget_limit

configure_query_budget_limit(limit)

Inject a configured per-query byte limit (e.g. from budgets.yml).

Source code in src/fastflowtransform/executors/base.py

def configure_query_budget_limit(self, limit: int | None) -> None:
    """
    Inject a configured per-query byte limit (e.g. from budgets.yml).
    """
    if limit is None:
        self._ff_configured_query_limit = None
        return
    try:
        iv = int(limit)
    except Exception:
        self._ff_configured_query_limit = None
        return
    self._ff_configured_query_limit = iv if iv > 0 else None

reset_node_stats

reset_node_stats()

Reset per-node statistics buffer.

The run engine calls this before executing a model so that all stats recorded via _record_query_stats(...) belong to that node.

Source code in src/fastflowtransform/executors/base.py

def reset_node_stats(self) -> None:
    """
    Reset per-node statistics buffer.

    The run engine calls this before executing a model so that all
    stats recorded via `_record_query_stats(...)` belong to that node.
    """
    # just clear the buffer; next recording will re-create it
    self._ff_query_stats_buffer = []

get_node_stats

get_node_stats()

Aggregate buffered QueryStats into a simple dict:

{
  "bytes_scanned": <sum>,
  "rows": <sum>,
  "query_duration_ms": <sum>,
}

Called by the run engine after a node finishes.

Source code in src/fastflowtransform/executors/base.py

def get_node_stats(self) -> dict[str, int]:
    """
    Aggregate buffered QueryStats into a simple dict:

        {
          "bytes_scanned": <sum>,
          "rows": <sum>,
          "query_duration_ms": <sum>,
        }

    Called by the run engine after a node finishes.
    """
    stats_list = self._drain_query_stats()
    if not stats_list:
        return {}

    total_bytes = 0
    total_rows = 0
    total_duration = 0

    for s in stats_list:
        if s.bytes_processed is not None:
            total_bytes += int(s.bytes_processed)
        if s.rows is not None:
            total_rows += int(s.rows)
        if s.duration_ms is not None:
            total_duration += int(s.duration_ms)

    return {
        "bytes_scanned": total_bytes,
        "rows": total_rows,
        "query_duration_ms": total_duration,
    }

run_python

run_python(node)

Execute the Python model for a given node and materialize its result.

Source code in src/fastflowtransform/executors/base.py

def run_python(self, node: Node) -> None:
    """Execute the Python model for a given node and materialize its result."""
    func = REGISTRY.py_funcs[node.name]
    deps = REGISTRY.nodes[node.name].deps or []

    self._reset_http_ctx(node)

    # arg = self._build_python_args(node, deps)
    args, argmap = self._build_python_inputs(node, deps)
    requires = REGISTRY.py_requires.get(node.name, {})
    # if deps:
    #     self._validate_required(node.name, arg, requires)
    if deps:
        # Required-columns check works against the mapping
        self._validate_required(node.name, argmap, requires)

    # out = self._execute_python_func(func, arg, node)
    out = self._execute_python_func(func, args, node)

    target = relation_for(node.name)
    meta = getattr(node, "meta", {}) or {}
    mat = self._resolve_materialization_strategy(meta)

    if mat == "incremental":
        self._materialize_incremental(target, out, node, meta)
    elif mat == "view":
        self._materialize_view(target, out, node)
    else:
        self._materialize_relation(target, out, node)

    self._snapshot_http_ctx(node)

snapshot_prune

snapshot_prune(relation, unique_key, keep_last, *, dry_run=False)

Delete older snapshot versions while keeping the most recent keep_last rows per business key (including the current row).

Source code in src/fastflowtransform/executors/_snapshot_sql_mixin.py

    def snapshot_prune(
        self,
        relation: str,
        unique_key: list[str],
        keep_last: int,
        *,
        dry_run: bool = False,
    ) -> None:
        """
        Delete older snapshot versions while keeping the most recent `keep_last`
        rows per business key (including the current row).
        """
        ex = cast("BaseExecutor[Any]", self)

        if keep_last <= 0:
            return

        keys = [k for k in unique_key if k]
        if not keys:
            return

        target = self._snapshot_target_identifier(relation)
        vf = self.SNAPSHOT_VALID_FROM_COL  # type: ignore[attr-defined]

        key_select = ", ".join(keys)
        part_by = ", ".join(keys)

        ranked_sql = f"""
SELECT
  {key_select},
  {vf},
  ROW_NUMBER() OVER (
    PARTITION BY {part_by}
    ORDER BY {vf} DESC
  ) AS rn
FROM {target}
"""

        if dry_run:
            sql = f"""
WITH ranked AS (
  {ranked_sql}
)
SELECT COUNT(*) AS rows_to_delete
FROM ranked
WHERE rn > {int(keep_last)}
"""
            res = ex._execute_sql(sql)
            count = self._snapshot_fetch_count(res)
            echo(
                f"[DRY-RUN] snapshot_prune({relation}): would delete {count} row(s) "
                f"(keep_last={keep_last})"
            )
            return

        join_pred = " AND ".join([f"t.{k} = r.{k}" for k in keys])
        delete_sql = f"""
DELETE FROM {target} t
USING (
  {ranked_sql}
) r
WHERE
  r.rn > {int(keep_last)}
  AND {join_pred}
  AND t.{vf} = r.{vf}
"""
        ex._execute_sql(delete_sql)