fakesnow 0.9.23__py3-none-any.whl → 0.9.25__py3-none-any.whl

fakesnow/__init__.py

@@ -90,3 +90,4 @@ def patch(
         yield None
     finally:
         stack.close()
+        fs.duck_conn.close()

fakesnow/arrow.py

@@ -1,33 +1,45 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import cast
 
 import pyarrow as pa
+import pyarrow.compute as pc
+
+from fakesnow.types import ColumnInfo
+
+
+def to_sf_schema(schema: pa.Schema, rowtype: list[ColumnInfo]) -> pa.Schema:
+    # expected by the snowflake connector
+    # uses rowtype to populate metadata, rather than the arrow schema type, for consistency with
+    # rowtype returned in the response
 
+    assert len(schema) == len(rowtype), f"schema and rowtype must be same length but f{len(schema)=} f{len(rowtype)=}"
 
-def with_sf_metadata(schema: pa.Schema) -> pa.Schema:
     # see https://github.com/snowflakedb/snowflake-connector-python/blob/e9393a6/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.cpp#L32
     # and https://github.com/snowflakedb/snowflake-connector-python/blob/e9393a6/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.cpp#L10
-    fms = []
-    for i, t in enumerate(schema.types):
-        f = schema.field(i)
-
-        # TODO: precision, scale, charLength etc. for all types
-
-        if t == pa.bool_():
-            fm = f.with_metadata({"logicalType": "BOOLEAN"})
-        elif t == pa.int64():
-            # scale and precision required, see here
-            # https://github.com/snowflakedb/snowflake-connector-python/blob/416ff57/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.cpp#L147
-            fm = f.with_metadata({"logicalType": "FIXED", "precision": "38", "scale": "0"})
-        elif t == pa.float64():
-            fm = f.with_metadata({"logicalType": "REAL"})
-        elif isinstance(t, pa.Decimal128Type):
-            fm = f.with_metadata({"logicalType": "FIXED", "precision": str(t.precision), "scale": str(t.scale)})
-        elif t == pa.string():
-            # TODO: set charLength to size of column
-            fm = f.with_metadata({"logicalType": "TEXT", "charLength": "16777216"})
-        else:
-            raise NotImplementedError(f"Unsupported Arrow type: {t}")
-        fms.append(fm)
+
+    def sf_field(field: pa.Field, c: ColumnInfo) -> pa.Field:
+        if isinstance(field.type, pa.TimestampType):
+            # snowflake uses a struct to represent timestamps, see timestamp_to_sf_struct
+            fields = [pa.field("epoch", pa.int64(), nullable=False), pa.field("fraction", pa.int32(), nullable=False)]
+            if field.type.tz:
+                fields.append(pa.field("timezone", nullable=False, type=pa.int32()))
+            field = field.with_type(pa.struct(fields))
+        elif isinstance(field.type, pa.Time64Type):
+            field = field.with_type(pa.int64())
+
+        return field.with_metadata(
+            {
+                "logicalType": c["type"].upper(),
+                # required for FIXED type see
+                # https://github.com/snowflakedb/snowflake-connector-python/blob/416ff57/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.cpp#L147
+                "precision": str(c["precision"] or 38),
+                "scale": str(c["scale"] or 0),
+                "charLength": str(c["length"] or 0),
+            }
+        )
+
+    fms = [sf_field(schema.field(i), c) for i, c in enumerate(rowtype)]
     return pa.schema(fms)
 
 
@@ -39,29 +51,56 @@ def to_ipc(table: pa.Table) -> pa.Buffer:
 
     sink = pa.BufferOutputStream()
 
-    with pa.ipc.new_stream(sink, with_sf_metadata(table.schema)) as writer:
+    with pa.ipc.new_stream(sink, table.schema) as writer:
         writer.write_batch(batch)
 
     return sink.getvalue()
 
 
-# TODO: should this be derived before with_schema?
-def to_rowtype(schema: pa.Schema) -> list[dict[str, Any]]:
-    return [
-        {
-            "name": f.name,
-            # TODO
-            # "database": "",
-            # "schema": "",
-            # "table": "",
-            "nullable": f.nullable,
-            "type": f.metadata.get(b"logicalType").decode("utf-8").lower(),  # type: ignore
-            # TODO
-            # "byteLength": 20,
-            "length": int(f.metadata.get(b"charLength")) if f.metadata.get(b"charLength") else None,  # type: ignore
-            "scale": int(f.metadata.get(b"scale")) if f.metadata.get(b"scale") else None,  # type: ignore
-            "precision": int(f.metadata.get(b"precision")) if f.metadata.get(b"precision") else None,  # type: ignore
-            "collation": None,
-        }
-        for f in schema
-    ]
+def to_sf(table: pa.Table, rowtype: list[ColumnInfo]) -> pa.Table:
+    def to_sf_col(col: pa.Array) -> pa.Array:
+        if pa.types.is_timestamp(col.type):
+            return timestamp_to_sf_struct(col)
+        elif pa.types.is_time(col.type):
+            # as nanoseconds
+            return pc.multiply(col.cast(pa.int64()), 1000)  # type: ignore https://github.com/zen-xu/pyarrow-stubs/issues/44
+        return col
+
+    return pa.Table.from_arrays([to_sf_col(c) for c in table.columns], schema=to_sf_schema(table.schema, rowtype))
+
+
+def timestamp_to_sf_struct(ts: pa.Array | pa.ChunkedArray) -> pa.Array:
+    if isinstance(ts, pa.ChunkedArray):
+        # combine because pa.StructArray.from_arrays doesn't support ChunkedArray
+        ts = cast(pa.Array, ts.combine_chunks())  # see https://github.com/zen-xu/pyarrow-stubs/issues/46
+
+    if not isinstance(ts.type, pa.TimestampType):
+        raise ValueError(f"Expected TimestampArray, got {type(ts)}")
+
+    # Round to seconds, ie: strip subseconds
+    tsa_without_us = pc.floor_temporal(ts, unit="second")  # type: ignore https://github.com/zen-xu/pyarrow-stubs/issues/45
+    epoch = pc.divide(tsa_without_us.cast(pa.int64()), 1_000_000)  # type: ignore https://github.com/zen-xu/pyarrow-stubs/issues/44
+
+    # Calculate fractional part as nanoseconds
+    fraction = pc.multiply(pc.subsecond(ts), 1_000_000_000).cast(pa.int32())  # type: ignore
+
+    if ts.type.tz:
+        assert ts.type.tz == "UTC", f"Timezone {ts.type.tz} not yet supported"
+        timezone = pa.array([1440] * len(ts), type=pa.int32())
+
+        return pa.StructArray.from_arrays(
+            arrays=[epoch, fraction, timezone],  # type: ignore https://github.com/zen-xu/pyarrow-stubs/issues/42
+            fields=[
+                pa.field("epoch", nullable=False, type=pa.int64()),
+                pa.field("fraction", nullable=False, type=pa.int32()),
+                pa.field("timezone", nullable=False, type=pa.int32()),
+            ],
+        )
+    else:
+        return pa.StructArray.from_arrays(
+            arrays=[epoch, fraction],  # type: ignore https://github.com/zen-xu/pyarrow-stubs/issues/42
+            fields=[
+                pa.field("epoch", nullable=False, type=pa.int64()),
+                pa.field("fraction", nullable=False, type=pa.int32()),
+            ],
+        )

fakesnow/checks.py

@@ -68,3 +68,11 @@ def is_unqualified_table_expression(expression: exp.Expression) -> tuple[bool, b
         no_schema = not node.args.get("db")
 
     return no_database, no_schema
+
+
+def equal(left: exp.Identifier, right: exp.Identifier) -> bool:
+    # as per https://docs.snowflake.com/en/sql-reference/identifiers-syntax#label-identifier-casing
+    lid = left.this if left.quoted else left.this.upper()
+    rid = right.this if right.quoted else right.this.upper()
+
+    return lid == rid

fakesnow/conn.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-import json
 import os
 from collections.abc import Iterable
 from pathlib import Path
 from types import TracebackType
 from typing import Any
 
-import pandas as pd
 import snowflake.connector.converter
 import snowflake.connector.errors
 import sqlglot
@@ -147,29 +145,3 @@ class FakeSnowflakeConnection:
 
     def rollback(self) -> None:
         self.cursor().execute("ROLLBACK")
-
-    def _insert_df(self, df: pd.DataFrame, table_name: str) -> int:
-        # Objects in dataframes are written as parquet structs, and snowflake loads parquet structs as json strings.
-        # Whereas duckdb analyses a dataframe see https://duckdb.org/docs/api/python/data_ingestion.html#pandas-dataframes--object-columns
-        # and converts a object to the most specific type possible, eg: dict -> STRUCT, MAP or varchar, and list -> LIST
-        # For dicts see https://github.com/duckdb/duckdb/pull/3985 and https://github.com/duckdb/duckdb/issues/9510
-        #
-        # When the rows have dicts with different keys there isn't a single STRUCT that can cover them, so the type is
-        # varchar and value a string containing a struct representation. In order to support dicts with different keys
-        # we first convert the dicts to json strings. A pity we can't do something inside duckdb and avoid the dataframe
-        # copy and transform in python.
-
-        df = df.copy()
-
-        # Identify columns of type object
-        object_cols = df.select_dtypes(include=["object"]).columns
-
-        # Apply json.dumps to these columns
-        for col in object_cols:
-            # don't jsonify string
-            df[col] = df[col].apply(lambda x: json.dumps(x) if isinstance(x, (dict, list)) else x)
-
-        escaped_cols = ",".join(f'"{col}"' for col in df.columns.to_list())
-        self._duck_conn.execute(f"INSERT INTO {table_name}({escaped_cols}) SELECT * FROM df")
-
-        return self._duck_conn.fetchall()[0][0]

fakesnow/cursor.py

@@ -26,9 +26,11 @@ import fakesnow.transforms as transforms
 from fakesnow.types import describe_as_result_metadata
 
 if TYPE_CHECKING:
+    # don't require pandas at import time
     import pandas as pd
     import pyarrow.lib
 
+    # avoid circular import
     from fakesnow.conn import FakeSnowflakeConnection
 
 
@@ -110,13 +112,15 @@ class FakeSnowflakeCursor:
 
     @property
     def description(self) -> list[ResultMetadata]:
+        return describe_as_result_metadata(self._describe_last_sql())
+
+    def _describe_last_sql(self) -> list:
         # use a separate cursor to avoid consuming the result set on this cursor
         with self._conn.cursor() as cur:
+            # TODO: can we replace with self._duck_conn.description?
             expression = sqlglot.parse_one(f"DESCRIBE {self._last_sql}", read="duckdb")
             cur._execute(expression, self._last_params)  # noqa: SLF001
-            meta = describe_as_result_metadata(cur.fetchall())
-
-        return meta
+            return cur.fetchall()
 
     def execute(
         self,
@@ -135,10 +139,15 @@ class FakeSnowflakeCursor:
             command, params = self._rewrite_with_params(command, params)
             if self._conn.nop_regexes and any(re.match(p, command, re.IGNORECASE) for p in self._conn.nop_regexes):
                 transformed = transforms.SUCCESS_NOP
-            else:
-                expression = parse_one(command, read="snowflake")
-                transformed = self._transform(expression)
-            return self._execute(transformed, params)
+                self._execute(transformed, params)
+                return self
+
+            expression = parse_one(command, read="snowflake")
+            for exp in self._transform_explode(expression):
+                transformed = self._transform(exp)
+                self._execute(transformed, params)
+
+            return self
         except snowflake.connector.errors.ProgrammingError as e:
             self._sqlstate = e.sqlstate
             raise e
@@ -203,9 +212,12 @@ class FakeSnowflakeCursor:
             .transform(transforms.alter_table_strip_cluster_by)
         )
 
-    def _execute(
-        self, transformed: exp.Expression, params: Sequence[Any] | dict[Any, Any] | None = None
-    ) -> FakeSnowflakeCursor:
+    def _transform_explode(self, expression: exp.Expression) -> list[exp.Expression]:
+        # Applies transformations that require splitting the expression into multiple expressions
+        # Split transforms have limited support at the moment.
+        return transforms.merge(expression)
+
+    def _execute(self, transformed: exp.Expression, params: Sequence[Any] | dict[Any, Any] | None = None) -> None:
         self._arrow_table = None
         self._arrow_table_fetch_index = None
         self._rowcount = None
@@ -341,8 +353,6 @@ class FakeSnowflakeCursor:
         self._last_sql = result_sql or sql
         self._last_params = params
 
-        return self
-
     def _log_sql(self, sql: str, params: Sequence[Any] | dict[Any, Any] | None = None) -> None:
         if (fs_debug := os.environ.get("FAKESNOW_DEBUG")) and fs_debug != "snowflake":
             print(f"{sql};{params=}" if params else f"{sql};", file=sys.stderr)

fakesnow/pandas_tools.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
+import json
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import numpy as np
+from duckdb import DuckDBPyConnection
+
+from fakesnow.conn import FakeSnowflakeConnection
 
 if TYPE_CHECKING:
+    # don't require pandas at import time
     import pandas as pd
 
-    from fakesnow.conn import FakeSnowflakeConnection
 
 CopyResult = tuple[
     str,
@@ -68,10 +72,37 @@ def write_pandas(
 
         conn.cursor().execute(f"CREATE TABLE IF NOT EXISTS {name} ({','.join(cols)})")
 
-    count = conn._insert_df(df, name)  # noqa: SLF001
+    count = _insert_df(conn._duck_conn, df, name)  # noqa: SLF001
 
     # mocks https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#output
     mock_copy_results = [("fakesnow/file0.txt", "LOADED", count, count, 1, 0, None, None, None, None)]
 
     # return success
     return (True, len(mock_copy_results), count, mock_copy_results)
+
+
+def _insert_df(duck_conn: DuckDBPyConnection, df: pd.DataFrame, table_name: str) -> int:
+    # Objects in dataframes are written as parquet structs, and snowflake loads parquet structs as json strings.
+    # Whereas duckdb analyses a dataframe see https://duckdb.org/docs/api/python/data_ingestion.html#pandas-dataframes--object-columns
+    # and converts a object to the most specific type possible, eg: dict -> STRUCT, MAP or varchar, and list -> LIST
+    # For dicts see https://github.com/duckdb/duckdb/pull/3985 and https://github.com/duckdb/duckdb/issues/9510
+    #
+    # When the rows have dicts with different keys there isn't a single STRUCT that can cover them, so the type is
+    # varchar and value a string containing a struct representation. In order to support dicts with different keys
+    # we first convert the dicts to json strings. A pity we can't do something inside duckdb and avoid the dataframe
+    # copy and transform in python.
+
+    df = df.copy()
+
+    # Identify columns of type object
+    object_cols = df.select_dtypes(include=["object"]).columns
+
+    # Apply json.dumps to these columns
+    for col in object_cols:
+        # don't jsonify string
+        df[col] = df[col].apply(lambda x: json.dumps(x) if isinstance(x, (dict, list)) else x)
+
+    escaped_cols = ",".join(f'"{col}"' for col in df.columns.to_list())
+    duck_conn.execute(f"INSERT INTO {table_name}({escaped_cols}) SELECT * FROM df")
+
+    return duck_conn.fetchall()[0][0]

fakesnow/server.py

@@ -5,19 +5,22 @@ import json
 import secrets
 from base64 import b64encode
 from dataclasses import dataclass
+from typing import Any
 
+import snowflake.connector.errors
 from starlette.applications import Starlette
 from starlette.concurrency import run_in_threadpool
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.routing import Route
 
-from fakesnow.arrow import to_ipc, to_rowtype, with_sf_metadata
+from fakesnow.arrow import to_ipc, to_sf
 from fakesnow.fakes import FakeSnowflakeConnection
 from fakesnow.instance import FakeSnow
+from fakesnow.types import describe_as_rowtype
 
-fs = FakeSnow()
-sessions = {}
+shared_fs = FakeSnow()
+sessions: dict[str, FakeSnowflakeConnection] = {}
 
 
 @dataclass
@@ -27,9 +30,19 @@ class ServerError(Exception):
     message: str
 
 
-def login_request(request: Request) -> JSONResponse:
+async def login_request(request: Request) -> JSONResponse:
     database = request.query_params.get("databaseName")
     schema = request.query_params.get("schemaName")
+    body = await request.body()
+    body_json = json.loads(gzip.decompress(body))
+    session_params: dict[str, Any] = body_json["data"]["SESSION_PARAMETERS"]
+    if db_path := session_params.get("FAKESNOW_DB_PATH"):
+        # isolated creates a new in-memory database, rather than using the shared in-memory database
+        # so this connection won't share any tables with other connections
+        fs = FakeSnow() if db_path == ":isolated:" else FakeSnow(db_path=db_path)
+    else:
+        # share the in-memory database across connections
+        fs = shared_fs
     token = secrets.token_urlsafe(32)
     sessions[token] = fs.connect(database, schema)
     return JSONResponse({"data": {"token": token}, "success": True})
@@ -44,16 +57,30 @@ async def query_request(request: Request) -> JSONResponse:
 
         sql_text = body_json["sqlText"]
 
-        # only a single sql statement is sent at a time by the python snowflake connector
-        cur = await run_in_threadpool(conn.cursor().execute, sql_text)
-
-        assert cur._arrow_table, "No result set"  # noqa: SLF001
-
-        batch_bytes = to_ipc(cur._arrow_table)  # noqa: SLF001
-        rowset_b64 = b64encode(batch_bytes).decode("utf-8")
-
-        # TODO: avoid calling with_sf_metadata twice
-        rowtype = to_rowtype(with_sf_metadata(cur._arrow_table.schema))  # noqa: SLF001
+        try:
+            # only a single sql statement is sent at a time by the python snowflake connector
+            cur = await run_in_threadpool(conn.cursor().execute, sql_text)
+        except snowflake.connector.errors.ProgrammingError as e:
+            code = f"{e.errno:06d}"
+            return JSONResponse(
+                {
+                    "data": {
+                        "errorCode": code,
+                        "sqlState": e.sqlstate,
+                    },
+                    "code": code,
+                    "message": e.msg,
+                    "success": False,
+                }
+            )
+
+        rowtype = describe_as_rowtype(cur._describe_last_sql())  # noqa: SLF001
+
+        if cur._arrow_table:  # noqa: SLF001
+            batch_bytes = to_ipc(to_sf(cur._arrow_table, rowtype))  # noqa: SLF001
+            rowset_b64 = b64encode(batch_bytes).decode("utf-8")
+        else:
+            rowset_b64 = ""
 
         return JSONResponse(
             {

fakesnow/transforms.py

@@ -7,6 +7,7 @@ from typing import ClassVar, Literal, cast
 import sqlglot
 from sqlglot import exp
 
+from fakesnow import transforms_merge
 from fakesnow.instance import USERS_TABLE_FQ_NAME
 from fakesnow.variables import Variables
 
@@ -691,6 +692,10 @@ def json_extract_precedence(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def merge(expression: exp.Expression) -> list[exp.Expression]:
+    return transforms_merge.merge(expression)
+
+
 def random(expression: exp.Expression) -> exp.Expression:
     """Convert random() and random(seed).
 

fakesnow/transforms_merge.py

@@ -0,0 +1,203 @@
+import sqlglot
+from sqlglot import exp
+
+from fakesnow import checks
+
+# Implements snowflake's MERGE INTO functionality in duckdb (https://docs.snowflake.com/en/sql-reference/sql/merge).
+
+
+def merge(merge_expr: exp.Expression) -> list[exp.Expression]:
+    if not isinstance(merge_expr, exp.Merge):
+        return [merge_expr]
+
+    return [_create_merge_candidates(merge_expr), *_mutations(merge_expr), _counts(merge_expr)]
+
+
+def _create_merge_candidates(merge_expr: exp.Merge) -> exp.Expression:
+    """
+    Given a merge statement, produce a temporary table that joins together the target and source tables.
+    The merge_op column identifies which merge clause applies to the row.
+    """
+    target_tbl = merge_expr.this
+
+    source = merge_expr.args.get("using")
+    assert isinstance(source, exp.Expression)
+    source_id = (alias := source.args.get("alias")) and alias.this if isinstance(source, exp.Subquery) else source.this
+    assert isinstance(source_id, exp.Identifier)
+
+    join_expr = merge_expr.args.get("on")
+    assert isinstance(join_expr, exp.Binary)
+
+    case_when_clauses: list[str] = []
+    values: set[str] = set()
+
+    # extract keys that reference the source table from the join expression
+    # so they can be used by the mutation statements for joining
+    # will include the source table identifier
+    values.update(
+        map(
+            str,
+            {
+                c
+                for c in join_expr.find_all(exp.Column)
+                if (table := c.args.get("table"))
+                and isinstance(table, exp.Identifier)
+                and checks.equal(table, source_id)
+            },
+        )
+    )
+
+    # Iterate through the WHEN clauses to build up the CASE WHEN clauses
+    for w_idx, w in enumerate(merge_expr.expressions):
+        assert isinstance(w, exp.When), f"Expected When expression, got {w}"
+
+        predicate = join_expr.copy()
+        matched = w.args.get("matched")
+        then = w.args.get("then")
+        condition = w.args.get("condition")
+
+        if matched:
+            # matchedClause see https://docs.snowflake.com/en/sql-reference/sql/merge#matchedclause-for-updates-or-deletes
+            if condition:
+                # Combine the top level ON expression with the AND condition
+                # from this specific WHEN into a subquery, we use to target rows.
+                # Eg. MERGE INTO t1 USING t2 ON t1.t1Key = t2.t2Key
+                #       WHEN MATCHED AND t2.marked = 1 THEN DELETE
+                predicate = exp.And(this=predicate, expression=condition)
+
+            if isinstance(then, exp.Update):
+                case_when_clauses.append(f"WHEN {predicate} THEN {w_idx}")
+                values.update([str(c.expression) for c in then.expressions if isinstance(c.expression, exp.Column)])
+            elif isinstance(then, exp.Var) and then.args.get("this") == "DELETE":
+                case_when_clauses.append(f"WHEN {predicate} THEN {w_idx}")
+            else:
+                raise AssertionError(f"Expected 'Update' or 'Delete', got {then}")
+        else:
+            # notMatchedClause see https://docs.snowflake.com/en/sql-reference/sql/merge#notmatchedclause-for-inserts
+            assert isinstance(then, exp.Insert), f"Expected 'Insert', got {then}"
+            insert_values = then.expression.expressions
+            values.update([str(c) for c in insert_values if isinstance(c, exp.Column)])
+            predicate = f"AND {condition}" if condition else ""
+            case_when_clauses.append(f"WHEN {target_tbl}.rowid is NULL {predicate} THEN {w_idx}")
+
+    sql = f"""
+    CREATE OR REPLACE TEMPORARY TABLE merge_candidates AS
+    SELECT
+        {', '.join(sorted(values))},
+        CASE
+            {' '.join(case_when_clauses)}
+            ELSE NULL
+        END AS MERGE_OP
+    FROM {target_tbl}
+    FULL OUTER JOIN {source} ON {join_expr.sql()}
+    WHERE MERGE_OP IS NOT NULL
+    """
+
+    return sqlglot.parse_one(sql)
+
+
+def _mutations(merge_expr: exp.Merge) -> list[exp.Expression]:
+    """
+    Given a merge statement, produce a list of delete, update and insert statements that use the
+    merge_candidates and source table to update the target target.
+    """
+    target_tbl = merge_expr.this
+    source = merge_expr.args.get("using")
+    source_tbl = source.alias if isinstance(source, exp.Subquery) else source
+    join_expr = merge_expr.args.get("on")
+
+    statements: list[exp.Expression] = []
+
+    # Iterate through the WHEN clauses to generate delete/update/insert statements
+    for w_idx, w in enumerate(merge_expr.expressions):
+        assert isinstance(w, exp.When), f"Expected When expression, got {w}"
+
+        matched = w.args.get("matched")
+        then = w.args.get("then")
+
+        if matched:
+            if isinstance(then, exp.Var) and then.args.get("this") == "DELETE":
+                delete_sql = f"""
+                    DELETE FROM {target_tbl}
+                    USING merge_candidates AS {source_tbl}
+                    WHERE {join_expr}
+                    AND {source_tbl}.merge_op = {w_idx}
+                """
+                statements.append(sqlglot.parse_one(delete_sql))
+            elif isinstance(then, exp.Update):
+                # when the update statement has a table alias, duckdb doesn't support the alias in the set
+                # column name, so we use e.this.this to get just the column name without its table prefix
+                set_clauses = ", ".join(
+                    [f"{e.this.this} = {e.expression.sql()}" for e in then.args.get("expressions", [])]
+                )
+                update_sql = f"""
+                    UPDATE {target_tbl}
+                    SET {set_clauses}
+                    FROM merge_candidates AS {source_tbl}
+                    WHERE {join_expr}
+                    AND {source_tbl}.merge_op = {w_idx}
+                """
+                statements.append(sqlglot.parse_one(update_sql))
+            else:
+                raise AssertionError(f"Expected 'Update' or 'Delete', got {then}")
+        else:
+            assert isinstance(then, exp.Insert), f"Expected 'Insert', got {then}"
+            cols = [str(c) for c in then.this.expressions] if then.this else []
+            columns = f"({', '.join(cols)})" if cols else ""
+            values = ", ".join(map(str, then.expression.expressions))
+            insert_sql = f"""
+                INSERT INTO {target_tbl} {columns}
+                SELECT {values}
+                FROM merge_candidates AS {source_tbl}
+                WHERE {source_tbl}.merge_op = {w_idx}
+            """
+            statements.append(sqlglot.parse_one(insert_sql))
+
+    return statements
+
+
+def _counts(merge_expr: exp.Merge) -> exp.Expression:
+    """
+    Given a merge statement, derive the a SQL statement which produces the following columns using the merge_candidates
+    table:
+
+    - "number of rows inserted"
+    - "number of rows updated"
+    - "number of rows deleted"
+
+    Only columns relevant to the merge operation are included, eg: if no rows are deleted, the "number of rows deleted"
+    column is not included.
+    """
+
+    # Initialize dictionaries to store operation types and their corresponding indices
+    operations = {"inserted": [], "updated": [], "deleted": []}
+
+    # Iterate through the WHEN clauses to categorize operations
+    for w_idx, w in enumerate(merge_expr.expressions):
+        assert isinstance(w, exp.When), f"Expected When expression, got {w}"
+
+        matched = w.args.get("matched")
+        then = w.args.get("then")
+
+        if matched:
+            if isinstance(then, exp.Update):
+                operations["updated"].append(w_idx)
+            elif isinstance(then, exp.Var) and then.args.get("this") == "DELETE":
+                operations["deleted"].append(w_idx)
+            else:
+                raise AssertionError(f"Expected 'Update' or 'Delete', got {then}")
+        else:
+            assert isinstance(then, exp.Insert), f"Expected 'Insert', got {then}"
+            operations["inserted"].append(w_idx)
+
+    count_statements = [
+        f"""COUNT_IF(merge_op in ({','.join(map(str, indices))})) as \"number of rows {op}\""""
+        for op, indices in operations.items()
+        if indices
+    ]
+    sql = f"""
+    SELECT {', '.join(count_statements)}
+    FROM merge_candidates
+    """
+
+    return sqlglot.parse_one(sql)

{fakesnow-0.9.23.dist-info → fakesnow-0.9.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fakesnow
-Version: 0.9.23
+Version: 0.9.25
 Summary: Fake Snowflake Connector for Python. Run, mock and test Snowflake DB locally.
 License: Apache License
                                    Version 2.0, January 2004
@@ -220,10 +220,10 @@ Requires-Dist: dirty-equals; extra == "dev"
 Requires-Dist: pandas-stubs; extra == "dev"
 Requires-Dist: snowflake-connector-python[pandas,secure-local-storage]; extra == "dev"
 Requires-Dist: pre-commit~=3.4; extra == "dev"
-Requires-Dist: pyarrow-stubs; extra == "dev"
+Requires-Dist: pyarrow-stubs==10.0.1.9; extra == "dev"
 Requires-Dist: pytest~=8.0; extra == "dev"
 Requires-Dist: pytest-asyncio; extra == "dev"
-Requires-Dist: ruff~=0.5.1; extra == "dev"
+Requires-Dist: ruff~=0.6.3; extra == "dev"
 Requires-Dist: twine~=5.0; extra == "dev"
 Requires-Dist: snowflake-sqlalchemy~=1.5.0; extra == "dev"
 Provides-Extra: notebook

fakesnow-0.9.25.dist-info/RECORD

@@ -0,0 +1,26 @@
+fakesnow/__init__.py,sha256=qUfgucQYPdELrJaxczalhJgWAWQ6cfTCUAHx6nUqRaI,3528
+fakesnow/__main__.py,sha256=GDrGyNTvBFuqn_UfDjKs7b3LPtU6gDv1KwosVDrukIM,76
+fakesnow/arrow.py,sha256=EGAYeuCnRuvmWBEGqw2YOcgQR4zcCsZBu85kSRl70dQ,4698
+fakesnow/checks.py,sha256=N8sXldhS3u1gG32qvZ4VFlsKgavRKrQrxLiQU8am1lw,2691
+fakesnow/cli.py,sha256=9qfI-Ssr6mo8UmIlXkUAOz2z2YPBgDsrEVaZv9FjGFs,2201
+fakesnow/conn.py,sha256=Gy_Z7BZRm5yMjV3x6hR4iegDQFdG9aJBjqWdc3iWYFU,5353
+fakesnow/cursor.py,sha256=JbLSTzIN5Hu6ECn1kQ8-hC8V-ENeEWTID4JxHbGpEIo,19948
+fakesnow/expr.py,sha256=CAxuYIUkwI339DQIBzvFF0F-m1tcVGKEPA5rDTzmH9A,892
+fakesnow/fakes.py,sha256=JQTiUkkwPeQrJ8FDWhPFPK6pGwd_aR2oiOrNzCWznlM,187
+fakesnow/fixtures.py,sha256=G-NkVeruSQAJ7fvSS2fR2oysUn0Yra1pohHlOvacKEk,455
+fakesnow/info_schema.py,sha256=DObVOrhzppAFHsdtj4YI9oRISn9SkJUG6ONjVleQQ_Y,6303
+fakesnow/instance.py,sha256=3cJvPRuFy19dMKXbtBLl6imzO48pEw8uTYhZyFDuwhk,3133
+fakesnow/macros.py,sha256=pX1YJDnQOkFJSHYUjQ6ErEkYIKvFI6Ncz_au0vv1csA,265
+fakesnow/pandas_tools.py,sha256=WjyjTV8QUCQQaCGboaEOvx2uo4BkknpWYjtLwkeCY6U,3468
+fakesnow/py.typed,sha256=B-DLSjYBi7pkKjwxCSdpVj2J02wgfJr-E7B1wOUyxYU,80
+fakesnow/server.py,sha256=SO5xKZ4rvySsuKDsoSPSCZcFuIX_K7d1XJYhRRJ-7Bk,4150
+fakesnow/transforms.py,sha256=hPQ9L1TvsDjFiGbs8mGUnIFyPUR7JxFU8FZRKFj5ZD0,54568
+fakesnow/transforms_merge.py,sha256=7rq-UPjfFNRrFsqR8xx3otwP6-k4eslLVLhfuqSXq1A,8314
+fakesnow/types.py,sha256=9Tt83Z7ctc9_v6SYyayXYz4MEI4RZo4zq_uqdj4g3Dk,2681
+fakesnow/variables.py,sha256=WXyPnkeNwD08gy52yF66CVe2twiYC50tztNfgXV4q1k,3032
+fakesnow-0.9.25.dist-info/LICENSE,sha256=kW-7NWIyaRMQiDpryfSmF2DObDZHGR1cJZ39s6B1Svg,11344
+fakesnow-0.9.25.dist-info/METADATA,sha256=1RktYqC8KfU4ekh8IGApjHSLb5oe3KMgBLjyKJKXHRc,18074
+fakesnow-0.9.25.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+fakesnow-0.9.25.dist-info/entry_points.txt,sha256=2riAUgu928ZIHawtO8EsfrMEJhi-EH-z_Vq7Q44xKPM,47
+fakesnow-0.9.25.dist-info/top_level.txt,sha256=500evXI1IFX9so82cizGIEMHAb_dJNPaZvd2H9dcKTA,24
+fakesnow-0.9.25.dist-info/RECORD,,

fakesnow-0.9.23.dist-info/RECORD

@@ -1,25 +0,0 @@
-fakesnow/__init__.py,sha256=9tFJJKvowKNW3vfnlmza6hOLN1I52DwChgNc5Ew6CcA,3499
-fakesnow/__main__.py,sha256=GDrGyNTvBFuqn_UfDjKs7b3LPtU6gDv1KwosVDrukIM,76
-fakesnow/arrow.py,sha256=WLkr1nEiNxUcPdzadKSM33sRAiQJsN6LvuzTVIsi3D0,2766
-fakesnow/checks.py,sha256=-QMvdcrRbhN60rnzxLBJ0IkUBWyLR8gGGKKmCS0w9mA,2383
-fakesnow/cli.py,sha256=9qfI-Ssr6mo8UmIlXkUAOz2z2YPBgDsrEVaZv9FjGFs,2201
-fakesnow/conn.py,sha256=yR9SMGSKkLvdPfi5fDwj9PQggOrPcKkdBBnQy2y_Bak,6921
-fakesnow/cursor.py,sha256=lITuMMy_hA9_riC121Sv9bFDFbU9BP9NlOLdHlP0ahY,19371
-fakesnow/expr.py,sha256=CAxuYIUkwI339DQIBzvFF0F-m1tcVGKEPA5rDTzmH9A,892
-fakesnow/fakes.py,sha256=JQTiUkkwPeQrJ8FDWhPFPK6pGwd_aR2oiOrNzCWznlM,187
-fakesnow/fixtures.py,sha256=G-NkVeruSQAJ7fvSS2fR2oysUn0Yra1pohHlOvacKEk,455
-fakesnow/info_schema.py,sha256=DObVOrhzppAFHsdtj4YI9oRISn9SkJUG6ONjVleQQ_Y,6303
-fakesnow/instance.py,sha256=3cJvPRuFy19dMKXbtBLl6imzO48pEw8uTYhZyFDuwhk,3133
-fakesnow/macros.py,sha256=pX1YJDnQOkFJSHYUjQ6ErEkYIKvFI6Ncz_au0vv1csA,265
-fakesnow/pandas_tools.py,sha256=ecL0kxIVU5o2--P3bRLWWVhxXqq6Km4trFr36txukMg,1897
-fakesnow/py.typed,sha256=B-DLSjYBi7pkKjwxCSdpVj2J02wgfJr-E7B1wOUyxYU,80
-fakesnow/server.py,sha256=8dzaLUUXPzCMm6-ESn0CBws6XSwwOpnUuHQAZJ-4SwU,3011
-fakesnow/transforms.py,sha256=ellcY5OBc7mqgL9ChNolrqcCLWXF9RH21Jt88FcFl-I,54419
-fakesnow/types.py,sha256=9Tt83Z7ctc9_v6SYyayXYz4MEI4RZo4zq_uqdj4g3Dk,2681
-fakesnow/variables.py,sha256=WXyPnkeNwD08gy52yF66CVe2twiYC50tztNfgXV4q1k,3032
-fakesnow-0.9.23.dist-info/LICENSE,sha256=kW-7NWIyaRMQiDpryfSmF2DObDZHGR1cJZ39s6B1Svg,11344
-fakesnow-0.9.23.dist-info/METADATA,sha256=90vwAf40lg9ipw2urbu9nMq7cEARbtENE0aKd7qSpjo,18064
-fakesnow-0.9.23.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-fakesnow-0.9.23.dist-info/entry_points.txt,sha256=2riAUgu928ZIHawtO8EsfrMEJhi-EH-z_Vq7Q44xKPM,47
-fakesnow-0.9.23.dist-info/top_level.txt,sha256=500evXI1IFX9so82cizGIEMHAb_dJNPaZvd2H9dcKTA,24
-fakesnow-0.9.23.dist-info/RECORD,,