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

fakesnow/__init__.py

@@ -119,7 +119,7 @@ def server(port: int | None = None, session_parameters: dict[str, str | int | bo
 
     assert port
     server = uvicorn.Server(uvicorn.Config(fakesnow.server.app, port=port, log_level="info"))
-    thread = threading.Thread(target=server.run, name="Server", daemon=True)
+    thread = threading.Thread(target=server.run, name="fakesnow server", daemon=True)
     thread.start()
 
     while not server.started:

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
@@ -255,7 +256,6 @@ class FakeSnowflakeCursor:
             .transform(lambda e: transforms.show_keys(e, self._conn.database, kind="FOREIGN"))
             .transform(transforms.show_users)
             .transform(transforms.create_user)
-            .transform(transforms.copy_into)
             .transform(transforms.sha256)
             .transform(transforms.create_clone)
             .transform(transforms.alias_in_join)
@@ -286,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
@@ -307,10 +310,6 @@ class FakeSnowflakeCursor:
             raise snowflake.connector.errors.DatabaseError(msg=e.args[0], errno=250002, sqlstate="08003") from None
         except duckdb.ParserException as e:
             raise snowflake.connector.errors.ProgrammingError(msg=e.args[0], errno=1003, sqlstate="42000") from None
-        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
 
         affected_count = None
 
@@ -330,10 +329,6 @@ class FakeSnowflakeCursor:
             self._duck_conn.execute(info_schema.per_db_creation_sql(create_db_name))
             result_sql = SQL_CREATED_DATABASE.substitute(name=create_db_name)
 
-        elif copy_from := transformed.args.get("copy_from"):
-            (affected_count,) = self._duck_conn.fetchall()[0]
-            result_sql = SQL_COPY_ROWS.substitute(count=affected_count, file=copy_from)
-
         elif cmd == "INSERT":
             (affected_count,) = self._duck_conn.fetchall()[0]
             result_sql = SQL_INSERTED_ROWS.substitute(count=affected_count)
@@ -399,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()
@@ -409,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/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)};
     """