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

fakesnow/arrow.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import pyarrow as pa
 
 
@@ -8,10 +10,21 @@ def with_sf_metadata(schema: pa.Schema) -> pa.Schema:
     for i, t in enumerate(schema.types):
         f = schema.field(i)
 
-        if isinstance(t, pa.Decimal128Type):
+        # 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():
-            fm = f.with_metadata({"logicalType": "TEXT"})
+            # 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)
@@ -26,7 +39,29 @@ def to_ipc(table: pa.Table) -> pa.Buffer:
 
     sink = pa.BufferOutputStream()
 
-    with pa.ipc.new_stream(sink, with_sf_metadata(batch.schema)) as writer:
+    with pa.ipc.new_stream(sink, with_sf_metadata(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
+    ]

fakesnow/conn.py

@@ -0,0 +1,175 @@
+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
+from duckdb import DuckDBPyConnection
+from snowflake.connector.cursor import DictCursor, SnowflakeCursor
+from sqlglot import exp
+from typing_extensions import Self
+
+import fakesnow.info_schema as info_schema
+import fakesnow.macros as macros
+from fakesnow.cursor import FakeSnowflakeCursor
+from fakesnow.variables import Variables
+
+
+class FakeSnowflakeConnection:
+    def __init__(
+        self,
+        duck_conn: DuckDBPyConnection,
+        database: str | None = None,
+        schema: str | None = None,
+        create_database: bool = True,
+        create_schema: bool = True,
+        db_path: str | os.PathLike | None = None,
+        nop_regexes: list[str] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        self._duck_conn = duck_conn
+        self._is_closed = False
+        # upper case database and schema like snowflake unquoted identifiers
+        # so they appear as upper-cased in information_schema
+        # catalog and schema names are not actually case-sensitive in duckdb even though
+        # they are as cased in information_schema.schemata, so when selecting from
+        # information_schema.schemata below we use upper-case to match any existing duckdb
+        # catalog or schemas like "information_schema"
+        self.database = database and database.upper()
+        self.schema = schema and schema.upper()
+
+        self.database_set = False
+        self.schema_set = False
+        self.db_path = Path(db_path) if db_path else None
+        self.nop_regexes = nop_regexes
+        self._paramstyle = snowflake.connector.paramstyle
+        self.variables = Variables()
+
+        # create database if needed
+        if (
+            create_database
+            and self.database
+            and not duck_conn.execute(
+                f"""select * from information_schema.schemata
+                where upper(catalog_name) = '{self.database}'"""
+            ).fetchone()
+        ):
+            db_file = f"{self.db_path/self.database}.db" if self.db_path else ":memory:"
+            duck_conn.execute(f"ATTACH DATABASE '{db_file}' AS {self.database}")
+            duck_conn.execute(info_schema.creation_sql(self.database))
+            duck_conn.execute(macros.creation_sql(self.database))
+
+        # create schema if needed
+        if (
+            create_schema
+            and self.database
+            and self.schema
+            and not duck_conn.execute(
+                f"""select * from information_schema.schemata
+                where upper(catalog_name) = '{self.database}' and upper(schema_name) = '{self.schema}'"""
+            ).fetchone()
+        ):
+            duck_conn.execute(f"CREATE SCHEMA {self.database}.{self.schema}")
+
+        # set database and schema if both exist
+        if (
+            self.database
+            and self.schema
+            and duck_conn.execute(
+                f"""select * from information_schema.schemata
+                where upper(catalog_name) = '{self.database}' and upper(schema_name) = '{self.schema}'"""
+            ).fetchone()
+        ):
+            duck_conn.execute(f"SET schema='{self.database}.{self.schema}'")
+            self.database_set = True
+            self.schema_set = True
+        # set database if only that exists
+        elif (
+            self.database
+            and duck_conn.execute(
+                f"""select * from information_schema.schemata
+                where upper(catalog_name) = '{self.database}'"""
+            ).fetchone()
+        ):
+            duck_conn.execute(f"SET schema='{self.database}.main'")
+            self.database_set = True
+
+        # use UTC instead of local time zone for consistent testing
+        duck_conn.execute("SET GLOBAL TimeZone = 'UTC'")
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        pass
+
+    def close(self, retry: bool = True) -> None:
+        self._duck_conn.close()
+        self._is_closed = True
+
+    def commit(self) -> None:
+        self.cursor().execute("COMMIT")
+
+    def cursor(self, cursor_class: type[SnowflakeCursor] = SnowflakeCursor) -> FakeSnowflakeCursor:
+        # TODO: use duck_conn cursor for thread-safety
+        return FakeSnowflakeCursor(conn=self, duck_conn=self._duck_conn, use_dict_result=cursor_class == DictCursor)
+
+    def execute_string(
+        self,
+        sql_text: str,
+        remove_comments: bool = False,
+        return_cursors: bool = True,
+        cursor_class: type[SnowflakeCursor] = SnowflakeCursor,
+        **kwargs: dict[str, Any],
+    ) -> Iterable[FakeSnowflakeCursor]:
+        cursors = [
+            self.cursor(cursor_class).execute(e.sql(dialect="snowflake"))
+            for e in sqlglot.parse(sql_text, read="snowflake")
+            if e and not isinstance(e, exp.Semicolon)  # ignore comments
+        ]
+        return cursors if return_cursors else []
+
+    def is_closed(self) -> bool:
+        return self._is_closed
+
+    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]