fakesnow 0.9.22__py3-none-any.whl → 0.9.24__py3-none-any.whl

fakesnow/fakes.py

@@ -1,750 +1,3 @@
-from __future__ import annotations
-
-import json
-import os
-import re
-import sys
-from collections.abc import Iterable, Iterator, Sequence
-from pathlib import Path
-from string import Template
-from types import TracebackType
-from typing import TYPE_CHECKING, Any, Literal, Optional, cast
-
-import duckdb
-from sqlglot import exp
-
-if TYPE_CHECKING:
-    import pandas as pd
-    import pyarrow.lib
-import numpy as np
-import pyarrow
-import snowflake.connector.converter
-import snowflake.connector.errors
-import sqlglot
-from duckdb import DuckDBPyConnection
-from snowflake.connector.cursor import DictCursor, ResultMetadata, SnowflakeCursor
-from snowflake.connector.result_batch import ResultBatch
-from sqlglot import parse_one
-from typing_extensions import Self
-
-import fakesnow.checks as checks
-import fakesnow.expr as expr
-import fakesnow.info_schema as info_schema
-import fakesnow.macros as macros
-import fakesnow.transforms as transforms
-from fakesnow.variables import Variables
-
-SCHEMA_UNSET = "schema_unset"
-SQL_SUCCESS = "SELECT 'Statement executed successfully.' as 'status'"
-SQL_CREATED_DATABASE = Template("SELECT 'Database ${name} successfully created.' as 'status'")
-SQL_CREATED_SCHEMA = Template("SELECT 'Schema ${name} successfully created.' as 'status'")
-SQL_CREATED_TABLE = Template("SELECT 'Table ${name} successfully created.' as 'status'")
-SQL_CREATED_VIEW = Template("SELECT 'View ${name} successfully created.' as 'status'")
-SQL_DROPPED = Template("SELECT '${name} successfully dropped.' as 'status'")
-SQL_INSERTED_ROWS = Template("SELECT ${count} as 'number of rows inserted'")
-SQL_UPDATED_ROWS = Template("SELECT ${count} as 'number of rows updated', 0 as 'number of multi-joined rows updated'")
-SQL_DELETED_ROWS = Template("SELECT ${count} as 'number of rows deleted'")
-
-
-class FakeSnowflakeCursor:
-    def __init__(
-        self,
-        conn: FakeSnowflakeConnection,
-        duck_conn: DuckDBPyConnection,
-        use_dict_result: bool = False,
-    ) -> None:
-        """Create a fake snowflake cursor backed by DuckDB.
-
-        Args:
-            conn (FakeSnowflakeConnection): Used to maintain current database and schema.
-            duck_conn (DuckDBPyConnection): DuckDB connection.
-            use_dict_result (bool, optional): If true rows are returned as dicts otherwise they
-                are returned as tuples. Defaults to False.
-        """
-        self._conn = conn
-        self._duck_conn = duck_conn
-        self._use_dict_result = use_dict_result
-        self._last_sql = None
-        self._last_params = None
-        self._sqlstate = None
-        self._arraysize = 1
-        self._arrow_table = None
-        self._arrow_table_fetch_index = None
-        self._rowcount = None
-        self._converter = snowflake.connector.converter.SnowflakeConverter()
-
-    def __enter__(self) -> Self:
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> None:
-        pass
-
-    @property
-    def arraysize(self) -> int:
-        return self._arraysize
-
-    @arraysize.setter
-    def arraysize(self, value: int) -> None:
-        self._arraysize = value
-
-    def close(self) -> bool:
-        self._last_sql = None
-        self._last_params = None
-        return True
-
-    def describe(self, command: str, *args: Any, **kwargs: Any) -> list[ResultMetadata]:
-        """Return the schema of the result without executing the query.
-
-        Takes the same arguments as execute
-
-        Returns:
-            list[ResultMetadata]: _description_
-        """
-
-        describe = f"DESCRIBE {command}"
-        self.execute(describe, *args, **kwargs)
-        return FakeSnowflakeCursor._describe_as_result_metadata(self.fetchall())
-
-    @property
-    def description(self) -> list[ResultMetadata]:
-        # use a separate cursor to avoid consuming the result set on this cursor
-        with self._conn.cursor() as cur:
-            expression = sqlglot.parse_one(f"DESCRIBE {self._last_sql}", read="duckdb")
-            cur._execute(expression, self._last_params)  # noqa: SLF001
-            meta = FakeSnowflakeCursor._describe_as_result_metadata(cur.fetchall())
-
-        return meta
-
-    def execute(
-        self,
-        command: str,
-        params: Sequence[Any] | dict[Any, Any] | None = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> FakeSnowflakeCursor:
-        try:
-            self._sqlstate = None
-
-            if os.environ.get("FAKESNOW_DEBUG") == "snowflake":
-                print(f"{command};{params=}" if params else f"{command};", file=sys.stderr)
-
-            command = self._inline_variables(command)
-            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)
-        except snowflake.connector.errors.ProgrammingError as e:
-            self._sqlstate = e.sqlstate
-            raise e
-
-    def _transform(self, expression: exp.Expression) -> exp.Expression:
-        return (
-            expression.transform(transforms.upper_case_unquoted_identifiers)
-            .transform(transforms.update_variables, variables=self._conn.variables)
-            .transform(transforms.set_schema, current_database=self._conn.database)
-            .transform(transforms.create_database, db_path=self._conn.db_path)
-            .transform(transforms.extract_comment_on_table)
-            .transform(transforms.extract_comment_on_columns)
-            .transform(transforms.information_schema_fs_columns_snowflake)
-            .transform(transforms.information_schema_fs_tables_ext)
-            .transform(transforms.drop_schema_cascade)
-            .transform(transforms.tag)
-            .transform(transforms.semi_structured_types)
-            .transform(transforms.try_parse_json)
-            .transform(transforms.split)
-            # NOTE: trim_cast_varchar must be before json_extract_cast_as_varchar
-            .transform(transforms.trim_cast_varchar)
-            # indices_to_json_extract must be before regex_substr
-            .transform(transforms.indices_to_json_extract)
-            .transform(transforms.json_extract_cast_as_varchar)
-            .transform(transforms.json_extract_cased_as_varchar)
-            .transform(transforms.json_extract_precedence)
-            .transform(transforms.flatten_value_cast_as_varchar)
-            .transform(transforms.flatten)
-            .transform(transforms.regex_replace)
-            .transform(transforms.regex_substr)
-            .transform(transforms.values_columns)
-            .transform(transforms.to_date)
-            .transform(transforms.to_decimal)
-            .transform(transforms.try_to_decimal)
-            .transform(transforms.to_timestamp_ntz)
-            .transform(transforms.to_timestamp)
-            .transform(transforms.object_construct)
-            .transform(transforms.timestamp_ntz)
-            .transform(transforms.float_to_double)
-            .transform(transforms.integer_precision)
-            .transform(transforms.extract_text_length)
-            .transform(transforms.sample)
-            .transform(transforms.array_size)
-            .transform(transforms.random)
-            .transform(transforms.identifier)
-            .transform(transforms.array_agg_within_group)
-            .transform(transforms.array_agg)
-            .transform(transforms.dateadd_date_cast)
-            .transform(transforms.dateadd_string_literal_timestamp_cast)
-            .transform(transforms.datediff_string_literal_timestamp_cast)
-            .transform(lambda e: transforms.show_schemas(e, self._conn.database))
-            .transform(lambda e: transforms.show_objects_tables(e, self._conn.database))
-            # TODO collapse into a single show_keys function
-            .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="PRIMARY"))
-            .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="UNIQUE"))
-            .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="FOREIGN"))
-            .transform(transforms.show_users)
-            .transform(transforms.create_user)
-            .transform(transforms.sha256)
-            .transform(transforms.create_clone)
-            .transform(transforms.alias_in_join)
-            .transform(transforms.alter_table_strip_cluster_by)
-        )
-
-    def _execute(
-        self, transformed: exp.Expression, params: Sequence[Any] | dict[Any, Any] | None = None
-    ) -> FakeSnowflakeCursor:
-        self._arrow_table = None
-        self._arrow_table_fetch_index = None
-        self._rowcount = None
-
-        cmd = expr.key_command(transformed)
-
-        no_database, no_schema = checks.is_unqualified_table_expression(transformed)
-
-        if no_database and not self._conn.database_set:
-            raise snowflake.connector.errors.ProgrammingError(
-                msg=f"Cannot perform {cmd}. This session does not have a current database. Call 'USE DATABASE', or use a qualified name.",  # noqa: E501
-                errno=90105,
-                sqlstate="22000",
-            )
-        elif no_schema and not self._conn.schema_set:
-            raise snowflake.connector.errors.ProgrammingError(
-                msg=f"Cannot perform {cmd}. This session does not have a current schema. Call 'USE SCHEMA', or use a qualified name.",  # noqa: E501
-                errno=90106,
-                sqlstate="22000",
-            )
-
-        sql = transformed.sql(dialect="duckdb")
-
-        if transformed.find(exp.Select) and (seed := transformed.args.get("seed")):
-            sql = f"SELECT setseed({seed}); {sql}"
-
-        result_sql = None
-
-        try:
-            self._log_sql(sql, params)
-            self._duck_conn.execute(sql, params)
-        except duckdb.BinderException as e:
-            msg = e.args[0]
-            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2043, sqlstate="02000") from None
-        except duckdb.CatalogException as e:
-            # minimal processing to make it look like a snowflake exception, message content may differ
-            msg = cast(str, e.args[0]).split("\n")[0]
-            raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=2003, sqlstate="42S02") from None
-        except duckdb.TransactionException as e:
-            if "cannot rollback - no transaction is active" in str(
-                e
-            ) or "cannot commit - no transaction is active" in str(e):
-                # snowflake doesn't error on rollback or commit outside a tx
-                result_sql = SQL_SUCCESS
-            else:
-                raise e
-        except duckdb.ConnectionException as e:
-            raise snowflake.connector.errors.DatabaseError(msg=e.args[0], errno=250002, sqlstate="08003") from None
-
-        affected_count = None
-
-        if set_database := transformed.args.get("set_database"):
-            self._conn.database = set_database
-            self._conn.database_set = True
-
-        elif set_schema := transformed.args.get("set_schema"):
-            self._conn.schema = set_schema
-            self._conn.schema_set = True
-
-        elif create_db_name := transformed.args.get("create_db_name"):
-            # we created a new database, so create the info schema extensions
-            self._duck_conn.execute(info_schema.creation_sql(create_db_name))
-            result_sql = SQL_CREATED_DATABASE.substitute(name=create_db_name)
-
-        elif cmd == "INSERT":
-            (affected_count,) = self._duck_conn.fetchall()[0]
-            result_sql = SQL_INSERTED_ROWS.substitute(count=affected_count)
-
-        elif cmd == "UPDATE":
-            (affected_count,) = self._duck_conn.fetchall()[0]
-            result_sql = SQL_UPDATED_ROWS.substitute(count=affected_count)
-
-        elif cmd == "DELETE":
-            (affected_count,) = self._duck_conn.fetchall()[0]
-            result_sql = SQL_DELETED_ROWS.substitute(count=affected_count)
-
-        elif cmd in ("DESCRIBE TABLE", "DESCRIBE VIEW"):
-            # DESCRIBE TABLE/VIEW has already been run above to detect and error if the table exists
-            # We now rerun DESCRIBE TABLE/VIEW but transformed with columns to match Snowflake
-            result_sql = transformed.transform(
-                lambda e: transforms.describe_table(e, self._conn.database, self._conn.schema)
-            ).sql(dialect="duckdb")
-
-        elif (eid := transformed.find(exp.Identifier, bfs=False)) and isinstance(eid.this, str):
-            ident = eid.this if eid.quoted else eid.this.upper()
-            if cmd == "CREATE SCHEMA" and ident:
-                result_sql = SQL_CREATED_SCHEMA.substitute(name=ident)
-
-            elif cmd == "CREATE TABLE" and ident:
-                result_sql = SQL_CREATED_TABLE.substitute(name=ident)
-
-            elif cmd.startswith("ALTER") and ident:
-                result_sql = SQL_SUCCESS
-
-            elif cmd == "CREATE VIEW" and ident:
-                result_sql = SQL_CREATED_VIEW.substitute(name=ident)
-
-            elif cmd.startswith("DROP") and ident:
-                result_sql = SQL_DROPPED.substitute(name=ident)
-
-                # if dropping the current database/schema then reset conn metadata
-                if cmd == "DROP DATABASE" and ident == self._conn.database:
-                    self._conn.database = None
-                    self._conn.schema = None
-
-                elif cmd == "DROP SCHEMA" and ident == self._conn.schema:
-                    self._conn.schema = None
-
-        if table_comment := cast(tuple[exp.Table, str], transformed.args.get("table_comment")):
-            # record table comment
-            table, comment = table_comment
-            catalog = table.catalog or self._conn.database
-            schema = table.db or self._conn.schema
-            assert catalog and schema
-            self._duck_conn.execute(info_schema.insert_table_comment_sql(catalog, schema, table.name, comment))
-
-        if (text_lengths := cast(list[tuple[str, int]], transformed.args.get("text_lengths"))) and (
-            table := transformed.find(exp.Table)
-        ):
-            # record text lengths
-            catalog = table.catalog or self._conn.database
-            schema = table.db or self._conn.schema
-            assert catalog and schema
-            self._duck_conn.execute(info_schema.insert_text_lengths_sql(catalog, schema, table.name, text_lengths))
-
-        if result_sql:
-            self._log_sql(result_sql, params)
-            self._duck_conn.execute(result_sql)
-
-        self._arrow_table = self._duck_conn.fetch_arrow_table()
-        self._rowcount = affected_count or self._arrow_table.num_rows
-
-        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)
-
-    def executemany(
-        self,
-        command: str,
-        seqparams: Sequence[Any] | dict[str, Any],
-        **kwargs: Any,
-    ) -> FakeSnowflakeCursor:
-        if isinstance(seqparams, dict):
-            # see https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api
-            raise NotImplementedError("dict params not supported yet")
-
-        # TODO: support insert optimisations
-        # the snowflake connector will optimise inserts into a single query
-        # unless num_statements != 1 .. but for simplicity we execute each
-        # query one by one, which means the response differs
-        for p in seqparams:
-            self.execute(command, p)
-
-        return self
-
-    def fetchall(self) -> list[tuple] | list[dict]:
-        if self._arrow_table is None:
-            # mimic snowflake python connector error type
-            raise TypeError("No open result set")
-        return self.fetchmany(self._arrow_table.num_rows)
-
-    def fetch_pandas_all(self, **kwargs: dict[str, Any]) -> pd.DataFrame:
-        if self._arrow_table is None:
-            # mimic snowflake python connector error type
-            raise snowflake.connector.NotSupportedError("No open result set")
-        return self._arrow_table.to_pandas()
-
-    def fetchone(self) -> dict | tuple | None:
-        result = self.fetchmany(1)
-        return result[0] if result else None
-
-    def fetchmany(self, size: int | None = None) -> list[tuple] | list[dict]:
-        # https://peps.python.org/pep-0249/#fetchmany
-        size = size or self._arraysize
-
-        if self._arrow_table is None:
-            # mimic snowflake python connector error type
-            raise TypeError("No open result set")
-        tslice = self._arrow_table.slice(offset=self._arrow_table_fetch_index or 0, length=size).to_pylist()
-
-        if self._arrow_table_fetch_index is None:
-            self._arrow_table_fetch_index = size
-        else:
-            self._arrow_table_fetch_index += size
-
-        return tslice if self._use_dict_result else [tuple(d.values()) for d in tslice]
-
-    def get_result_batches(self) -> list[ResultBatch] | None:
-        if self._arrow_table is None:
-            return None
-        return [FakeResultBatch(self._use_dict_result, b) for b in self._arrow_table.to_batches(max_chunksize=1000)]
-
-    @property
-    def rowcount(self) -> int | None:
-        return self._rowcount
-
-    @property
-    def sfqid(self) -> str | None:
-        return "fakesnow"
-
-    @property
-    def sqlstate(self) -> str | None:
-        return self._sqlstate
-
-    @staticmethod
-    def _describe_as_result_metadata(describe_results: list) -> list[ResultMetadata]:
-        # fmt: off
-        def as_result_metadata(column_name: str, column_type: str, _: str) -> ResultMetadata:
-            # see https://docs.snowflake.com/en/user-guide/python-connector-api.html#type-codes
-            # and https://arrow.apache.org/docs/python/api/datatypes.html#type-checking
-            if column_type in {"BIGINT", "INTEGER"}:
-                return ResultMetadata(
-                    name=column_name, type_code=0, display_size=None, internal_size=None, precision=38, scale=0, is_nullable=True               # noqa: E501
-                )
-            elif column_type.startswith("DECIMAL"):
-                match = re.search(r'\((\d+),(\d+)\)', column_type)
-                if match:
-                    precision = int(match[1])
-                    scale = int(match[2])
-                else:
-                    precision = scale = None
-                return ResultMetadata(
-                    name=column_name, type_code=0, display_size=None, internal_size=None, precision=precision, scale=scale, is_nullable=True    # noqa: E501
-                )
-            elif column_type == "VARCHAR":
-                # TODO: fetch internal_size from varchar size
-                return ResultMetadata(
-                    name=column_name, type_code=2, display_size=None, internal_size=16777216, precision=None, scale=None, is_nullable=True      # noqa: E501
-                )
-            elif column_type == "DOUBLE":
-                return ResultMetadata(
-                    name=column_name, type_code=1, display_size=None, internal_size=None, precision=None, scale=None, is_nullable=True          # noqa: E501
-                )
-            elif column_type == "BOOLEAN":
-                return ResultMetadata(
-                    name=column_name, type_code=13, display_size=None, internal_size=None, precision=None, scale=None, is_nullable=True         # noqa: E501
-                )
-            elif column_type == "DATE":
-                return ResultMetadata(
-                    name=column_name, type_code=3, display_size=None, internal_size=None, precision=None, scale=None, is_nullable=True          # noqa: E501
-                )
-            elif column_type in {"TIMESTAMP", "TIMESTAMP_NS"}:
-                return ResultMetadata(
-                    name=column_name, type_code=8, display_size=None, internal_size=None, precision=0, scale=9, is_nullable=True                # noqa: E501
-                )
-            elif column_type == "TIMESTAMP WITH TIME ZONE":
-                return ResultMetadata(
-                    name=column_name, type_code=7, display_size=None, internal_size=None, precision=0, scale=9, is_nullable=True                # noqa: E501
-                )
-            elif column_type == "BLOB":
-                return ResultMetadata(
-                    name=column_name, type_code=11, display_size=None, internal_size=8388608, precision=None, scale=None, is_nullable=True      # noqa: E501
-                )
-            elif column_type == "TIME":
-                return ResultMetadata(
-                    name=column_name, type_code=12, display_size=None, internal_size=None, precision=0, scale=9, is_nullable=True               # noqa: E501
-                )
-            elif column_type == "JSON":
-                # TODO: correctly map OBJECT and ARRAY see https://github.com/tekumara/fakesnow/issues/26
-                return ResultMetadata(
-                    name=column_name, type_code=5, display_size=None, internal_size=None, precision=None, scale=None, is_nullable=True               # noqa: E501
-                )
-            else:
-                # TODO handle more types
-                raise NotImplementedError(f"for column type {column_type}")
-
-        # fmt: on
-
-        meta = [
-            as_result_metadata(column_name, column_type, null)
-            for (column_name, column_type, null, _, _, _) in describe_results
-        ]
-        return meta
-
-    def _rewrite_with_params(
-        self,
-        command: str,
-        params: Sequence[Any] | dict[Any, Any] | None = None,
-    ) -> tuple[str, Sequence[Any] | dict[Any, Any] | None]:
-        if params and self._conn._paramstyle in ("pyformat", "format"):  # noqa: SLF001
-            # handle client-side in the same manner as the snowflake python connector
-
-            def convert(param: Any) -> Any:  # noqa: ANN401
-                return self._converter.quote(self._converter.escape(self._converter.to_snowflake(param)))
-
-            if isinstance(params, dict):
-                params = {k: convert(v) for k, v in params.items()}
-            else:
-                params = tuple(convert(v) for v in params)
-
-            return command % params, None
-
-        return command, params
-
-    def _inline_variables(self, sql: str) -> str:
-        return self._conn.variables.inline_variables(sql)
-
-
-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
-        # 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()
-
-    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 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]
-
-
-class FakeResultBatch(ResultBatch):
-    def __init__(self, use_dict_result: bool, batch: pyarrow.RecordBatch):
-        self._use_dict_result = use_dict_result
-        self._batch = batch
-
-    def create_iter(
-        self, **kwargs: dict[str, Any]
-    ) -> Iterator[dict | Exception] | Iterator[tuple | Exception] | Iterator[pyarrow.Table] | Iterator[pd.DataFrame]:
-        if self._use_dict_result:
-            return iter(self._batch.to_pylist())
-
-        return iter(tuple(d.values()) for d in self._batch.to_pylist())
-
-    @property
-    def rowcount(self) -> int:
-        return self._batch.num_rows
-
-    def to_pandas(self) -> pd.DataFrame:
-        return self._batch.to_pandas()
-
-    def to_arrow(self) -> pyarrow.Table:
-        raise NotImplementedError()
-
-
-CopyResult = tuple[
-    str,
-    str,
-    int,
-    int,
-    int,
-    int,
-    Optional[str],
-    Optional[int],
-    Optional[int],
-    Optional[str],
-]
-
-WritePandasResult = tuple[
-    bool,
-    int,
-    int,
-    Sequence[CopyResult],
-]
-
-
-def sql_type(dtype: np.dtype) -> str:
-    if str(dtype) == "int64":
-        return "NUMBER"
-    elif str(dtype) == "object":
-        return "VARCHAR"
-    else:
-        raise NotImplementedError(f"sql_type {dtype=}")
-
-
-def write_pandas(
-    conn: FakeSnowflakeConnection,
-    df: pd.DataFrame,
-    table_name: str,
-    database: str | None = None,
-    schema: str | None = None,
-    chunk_size: int | None = None,
-    compression: str = "gzip",
-    on_error: str = "abort_statement",
-    parallel: int = 4,
-    quote_identifiers: bool = True,
-    auto_create_table: bool = False,
-    create_temp_table: bool = False,
-    overwrite: bool = False,
-    table_type: Literal["", "temp", "temporary", "transient"] = "",
-    **kwargs: Any,
-) -> WritePandasResult:
-    name = table_name
-    if schema:
-        name = f"{schema}.{name}"
-    if database:
-        name = f"{database}.{name}"
-
-    if auto_create_table:
-        cols = [f"{c} {sql_type(t)}" for c, t in df.dtypes.to_dict().items()]
-
-        conn.cursor().execute(f"CREATE TABLE IF NOT EXISTS {name} ({','.join(cols)})")
-
-    count = conn._insert_df(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)
+from .conn import FakeSnowflakeConnection as FakeSnowflakeConnection
+from .cursor import FakeSnowflakeCursor as FakeSnowflakeCursor
+from .pandas_tools import write_pandas as write_pandas