fakesnow 0.9.35__py3-none-any.whl → 0.9.37__py3-none-any.whl

fakesnow/__init__.py

@@ -91,3 +91,54 @@ def patch(
     finally:
         stack.close()
         fs.duck_conn.close()
+
+
+@contextmanager
+def server(port: int | None = None, session_parameters: dict[str, str | int | bool] | None = None) -> Iterator[dict]:
+    """Start a fake snowflake server in a separate thread and yield connection kwargs.
+
+    Args:
+        port (int | None, optional): Port to run the server on. If None, an available port is chosen. Defaults to None.
+
+    Yields:
+        Iterator[dict]: Connection parameters for the fake snowflake server.
+    """
+    import socket
+    import threading
+    from time import sleep
+
+    import uvicorn
+
+    import fakesnow.server
+
+    # find an unused TCP port between 1024-65535
+    if not port:
+        with contextlib.closing(socket.socket(type=socket.SOCK_STREAM)) as sock:
+            sock.bind(("127.0.0.1", 0))
+            port = sock.getsockname()[1]
+
+    assert port
+    server = uvicorn.Server(uvicorn.Config(fakesnow.server.app, port=port, log_level="info"))
+    thread = threading.Thread(target=server.run, name="fakesnow server", daemon=True)
+    thread.start()
+
+    while not server.started:
+        sleep(0.1)
+
+    try:
+        yield dict(
+            user="fake",
+            password="snow",
+            account="fakesnow",
+            host="localhost",
+            port=port,
+            protocol="http",
+            # disable telemetry
+            session_parameters={"CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED": False} | (session_parameters or {}),
+            # disable retries on error
+            network_timeout=1,
+        )
+    finally:
+        server.should_exit = True
+        # wait for server thread to end
+        thread.join()

fakesnow/conn.py

@@ -6,8 +6,7 @@ from pathlib import Path
 from types import TracebackType
 from typing import Any
 
-import snowflake.connector.converter
-import snowflake.connector.errors
+import snowflake.connector
 import sqlglot
 from duckdb import DuckDBPyConnection
 from snowflake.connector.cursor import DictCursor, SnowflakeCursor

fakesnow/converter.py

@@ -8,20 +8,20 @@ from datetime import date, time, timezone
 
 
 def from_binding(binding: dict[str, str]) -> int | bytes | bool | date | time | datetime.datetime | str:
-    typ = binding["type"]
+    type_ = binding["type"]
     value = binding["value"]
-    if typ == "FIXED":
+    if type_ == "FIXED":
         return int(value)
-    elif typ == "BINARY":
+    elif type_ == "BINARY":
         return from_binary(value)
     # TODO: not strictly needed
-    elif typ == "BOOLEAN":
+    elif type_ == "BOOLEAN":
         return value.lower() == "true"
-    elif typ == "DATE":
+    elif type_ == "DATE":
         return from_date(value)
-    elif typ == "TIME":
+    elif type_ == "TIME":
         return from_time(value)
-    elif typ == "TIMESTAMP_NTZ":
+    elif type_ == "TIMESTAMP_NTZ":
         return from_datetime(value)
     else:
         # For other types, return str

fakesnow/copy_into.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any, Protocol, cast
+from urllib.parse import urlparse, urlunparse
+
+import duckdb
+import snowflake.connector.errors
+from duckdb import DuckDBPyConnection
+from sqlglot import exp
+
+from fakesnow import logger
+
+
+def copy_into(
+    duck_conn: DuckDBPyConnection, expr: exp.Copy, params: Sequence[Any] | dict[Any, Any] | None = None
+) -> str:
+    cparams = _params(expr)
+    urls = _source_urls(expr, cparams.files)
+    inserts = _inserts(expr, cparams, urls)
+
+    results = []
+    try:
+        # TODO: fetch files last modified dates and check if file exists in load_history already
+        for i, url in zip(inserts, urls):
+            sql = i.sql(dialect="duckdb")
+            logger.log_sql(sql, params)
+            duck_conn.execute(sql, params)
+            (affected_count,) = duck_conn.fetchall()[0]
+            results.append(f"('{url}', 'LOADED', {affected_count}, {affected_count}, 1, 0, NULL, NULL, NULL, NULL)")
+
+            # TODO: update load_history with the results if loaded
+
+        columns = "file, status, rows_parsed, rows_loaded, error_limit, errors_seen, first_error, first_error_line, first_error_character, first_error_column_name"  # noqa: E501
+        values = "\n, ".join(results)
+        sql = f"SELECT * FROM (VALUES\n  {values}\n) AS t({columns})"
+        duck_conn.execute(sql)
+        return sql
+    except duckdb.HTTPException as e:
+        raise snowflake.connector.errors.ProgrammingError(msg=e.args[0], errno=91016, sqlstate="22000") from None
+    except duckdb.ConversionException as e:
+        raise snowflake.connector.errors.ProgrammingError(msg=e.args[0], errno=100038, sqlstate="22018") from None
+
+
+def _params(expr: exp.Copy) -> Params:
+    kwargs = {}
+    force = False
+
+    params = cast(list[exp.CopyParameter], expr.args.get("params", []))
+    cparams = Params()
+    for param in params:
+        assert isinstance(param.this, exp.Var), f"{param.this.__class__} is not a Var"
+        var = param.this.name.upper()
+        if var == "FILE_FORMAT":
+            if kwargs.get("file_format"):
+                raise ValueError(cparams)
+
+            var_type = next((e.args["value"].this for e in param.expressions if e.this.this == "TYPE"), None)
+            if not var_type:
+                raise NotImplementedError("FILE_FORMAT without TYPE is not currently implemented")
+
+            if var_type == "CSV":
+                kwargs["file_format"] = handle_csv(param.expressions)
+            else:
+                raise NotImplementedError(f"{var_type} FILE_FORMAT is not currently implemented")
+        elif var == "FORCE":
+            force = True
+        elif var == "FILES":
+            kwargs["files"] = [lit.name for lit in param.find_all(exp.Literal)]
+        else:
+            raise ValueError(f"Unknown copy parameter: {param.this}")
+
+    if not force:
+        raise NotImplementedError("COPY INTO with FORCE=false (default) is not currently implemented")
+
+    return Params(**kwargs)
+
+
+def _source_urls(expr: exp.Copy, files: list[str]) -> list[str]:
+    """
+    Given a COPY statement and a list of files, return a list of URLs with each file appended as a fragment.
+    Checks that the source is a valid URL.
+    """
+    source = expr.args["files"][0].this
+    assert isinstance(source, exp.Literal), f"{source} is not a exp.Literal"
+
+    scheme, netloc, path, params, query, fragment = urlparse(source.name)
+    if not scheme:
+        raise snowflake.connector.errors.ProgrammingError(
+            msg=f"SQL compilation error:\ninvalid URL prefix found in: '{source.name}'", errno=1011, sqlstate="42601"
+        )
+
+    # rebuild url from components to ensure correct handling of host slash
+    return [_urlunparse(scheme, netloc, path, params, query, fragment, file) for file in files] or [source.name]
+
+
+def _urlunparse(scheme: str, netloc: str, path: str, params: str, query: str, fragment: str, suffix: str) -> str:
+    """Construct a URL from its components appending suffix to the last used component."""
+    if fragment:
+        fragment += suffix
+    elif query:
+        query += suffix
+    elif params:
+        params += suffix
+    else:
+        path += suffix
+    return urlunparse((scheme, netloc, path, params, query, fragment))
+
+
+def _inserts(expr: exp.Copy, params: Params, urls: list[str]) -> list[exp.Expression]:
+    # INTO expression
+    target = expr.this
+    columns = [exp.Column(this=exp.Identifier(this=f"column{i}")) for i in range(len(target.expressions))] or [
+        exp.Column(this=exp.Star())
+    ]
+
+    return [
+        exp.Insert(
+            this=target,
+            expression=exp.Select(expressions=columns).from_(exp.Table(this=params.file_format.read_expression(url))),
+        )
+        for url in urls
+    ]
+
+
+def handle_csv(expressions: list[exp.Property]) -> ReadCSV:
+    skip_header = ReadCSV.skip_header
+    quote = ReadCSV.quote
+    delimiter = ReadCSV.delimiter
+
+    for expression in expressions:
+        exp_type = expression.name
+        if exp_type in {"TYPE"}:
+            continue
+
+        elif exp_type == "SKIP_HEADER":
+            skip_header = True
+        elif exp_type == "FIELD_OPTIONALLY_ENCLOSED_BY":
+            quote = expression.args["value"].this
+        elif exp_type == "FIELD_DELIMITER":
+            delimiter = expression.args["value"].this
+        else:
+            raise NotImplementedError(f"{exp_type} is not currently implemented")
+
+    return ReadCSV(
+        skip_header=skip_header,
+        quote=quote,
+        delimiter=delimiter,
+    )
+
+
+@dataclass
+class FileTypeHandler(Protocol):
+    def read_expression(self, url: str) -> exp.Expression: ...
+
+    @staticmethod
+    def make_eq(name: str, value: list | str | int | bool) -> exp.EQ:
+        if isinstance(value, list):
+            expression = exp.array(*[exp.Literal(this=str(v), is_string=isinstance(v, str)) for v in value])
+        elif isinstance(value, bool):
+            expression = exp.Boolean(this=value)
+        else:
+            expression = exp.Literal(this=str(value), is_string=isinstance(value, str))
+
+        return exp.EQ(this=exp.Literal(this=name, is_string=False), expression=expression)
+
+
+@dataclass
+class ReadCSV(FileTypeHandler):
+    skip_header: bool = False
+    quote: str | None = None
+    delimiter: str = ","
+
+    def read_expression(self, url: str) -> exp.Expression:
+        args = []
+
+        # don't parse header and use as column names, keep them as column0, column1, etc
+        args.append(self.make_eq("header", False))
+
+        if self.skip_header:
+            args.append(self.make_eq("skip", 1))
+
+        if self.quote:
+            quote = self.quote.replace("'", "''")
+            args.append(self.make_eq("quote", quote))
+
+        if self.delimiter and self.delimiter != ",":
+            delimiter = self.delimiter.replace("'", "''")
+            args.append(self.make_eq("sep", delimiter))
+
+        return exp.func("read_csv", exp.Literal(this=url, is_string=True), *args)
+
+
+@dataclass
+class Params:
+    files: list[str] = field(default_factory=list)
+    # Snowflake defaults to CSV when no file format is specified
+    file_format: FileTypeHandler = field(default_factory=ReadCSV)

fakesnow/cursor.py

@@ -25,12 +25,13 @@ import fakesnow.checks as checks
 import fakesnow.expr as expr
 import fakesnow.info_schema as info_schema
 import fakesnow.transforms as transforms
+from fakesnow import logger
+from fakesnow.copy_into import copy_into
 from fakesnow.rowtype 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
@@ -46,6 +47,12 @@ 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'")
+SQL_COPY_ROWS = Template(
+    "SELECT '${file}' as file, 'LOADED' as status, ${count} as rows_parsed, "
+    "${count} as rows_loaded, 1 as error_limit, 0 as errors_seen, "
+    "NULL as first_error, NULL as first_error_line, NULL as first_error_character, "
+    "NULL as first_error_column_name"
+)
 
 
 class FakeSnowflakeCursor:
@@ -279,8 +286,11 @@ class FakeSnowflakeCursor:
         result_sql = None
 
         try:
-            self._log_sql(sql, params)
-            self._duck_conn.execute(sql, params)
+            if isinstance(transformed, exp.Copy):
+                sql = copy_into(self._duck_conn, transformed, params)
+            else:
+                logger.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
@@ -384,7 +394,7 @@ class FakeSnowflakeCursor:
             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)
+            logger.log_sql(result_sql)
             self._duck_conn.execute(result_sql)
 
         self._arrow_table = self._duck_conn.fetch_arrow_table()
@@ -394,10 +404,6 @@ class FakeSnowflakeCursor:
         self._last_sql = result_sql or sql
         self._last_params = None if result_sql else params
 
-    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,

fakesnow/fixtures.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterator
+from typing import Any
 
 import pytest
 
@@ -11,6 +12,12 @@ def _fakesnow() -> Iterator[None]:
         yield
 
 
+@pytest.fixture(scope="session")
+def fakesnow_server() -> Iterator[dict[str, Any]]:
+    with fakesnow.server() as conn_kwargs:
+        yield conn_kwargs
+
+
 @pytest.fixture
 def _fakesnow_no_auto_create() -> Iterator[None]:
     with fakesnow.patch(create_database_on_connect=False, create_schema_on_connect=False):

fakesnow/logger.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import Sequence
+from typing import Any
+
+
+def log_sql(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/macros.py

@@ -6,8 +6,32 @@ CREATE MACRO IF NOT EXISTS ${catalog}.equal_null(a, b) AS a IS NOT DISTINCT FROM
 """
 )
 
+# emulate the Snowflake FLATTEN function for ARRAYs
+# see https://docs.snowflake.com/en/sql-reference/functions/flatten.html
+FS_FLATTEN = Template(
+    """
+CREATE OR REPLACE MACRO ${catalog}._fs_flatten(input) AS TABLE
+    SELECT
+        NULL AS SEQ, -- TODO use a sequence and nextval
+        CAST(NULL AS VARCHAR) AS KEY,
+        '[' || GENERATE_SUBSCRIPTS(
+            CAST(TO_JSON(input) AS JSON []),
+            1
+        ) - 1 || ']' AS PATH,
+        GENERATE_SUBSCRIPTS(
+            CAST(TO_JSON(input) AS JSON []),
+            1
+        ) - 1 AS INDEX,
+        UNNEST(
+            CAST(TO_JSON(input) AS JSON [])
+        ) AS VALUE,
+        TO_JSON(input) AS THIS;
+    """
+)
+
 
 def creation_sql(catalog: str) -> str:
     return f"""
         {EQUAL_NULL.substitute(catalog=catalog)};
+        {FS_FLATTEN.substitute(catalog=catalog)};
     """

fakesnow/server.py

@@ -59,7 +59,10 @@ async def login_request(request: Request) -> JSONResponse:
         {
             "data": {
                 "token": token,
-                "parameters": [{"name": "AUTOCOMMIT", "value": True}],
+                "parameters": [
+                    {"name": "AUTOCOMMIT", "value": True},
+                    {"name": "CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY", "value": 3600},
+                ],
             },
             "success": True,
         }