fakesnow 0.9.7__tar.gz → 0.9.8__tar.gz

{fakesnow-0.9.7 → fakesnow-0.9.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fakesnow
-Version: 0.9.7
+Version: 0.9.8
 Summary: Fake Snowflake Connector for Python. Run, mock and test Snowflake DB locally.
 License:                                  Apache License
                                    Version 2.0, January 2004

{fakesnow-0.9.7 → fakesnow-0.9.8}/fakesnow/fakes.py

@@ -11,6 +11,7 @@ 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
@@ -22,7 +23,7 @@ import sqlglot
 from duckdb import DuckDBPyConnection
 from snowflake.connector.cursor import DictCursor, ResultMetadata, SnowflakeCursor
 from snowflake.connector.result_batch import ResultBatch
-from sqlglot import exp, parse_one
+from sqlglot import parse_one
 from typing_extensions import Self
 
 import fakesnow.checks as checks
@@ -112,7 +113,9 @@ class FakeSnowflakeCursor:
     def description(self) -> list[ResultMetadata]:
         # use a separate cursor to avoid consuming the result set on this cursor
         with self._conn.cursor() as cur:
-            cur.execute(f"DESCRIBE {self._last_sql}", self._last_params)
+            # self._duck_conn.execute(sql, params)
+            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
@@ -126,43 +129,20 @@ class FakeSnowflakeCursor:
     ) -> FakeSnowflakeCursor:
         try:
             self._sqlstate = None
-            return self._execute(command, params, *args, **kwargs)
+
+            if os.environ.get("FAKESNOW_DEBUG") == "snowflake":
+                print(f"{command};{params=}" if params else f"{command};", file=sys.stderr)
+
+            command, params = self._rewrite_with_params(command, params)
+            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 _execute(
-        self,
-        command: str,
-        params: Sequence[Any] | dict[Any, Any] | None = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> FakeSnowflakeCursor:
-        self._arrow_table = None
-        self._arrow_table_fetch_index = None
-        self._rowcount = None
-
-        command, params = self._rewrite_with_params(command, params)
-        expression = parse_one(command, read="snowflake")
-
-        cmd = expr.key_command(expression)
-
-        no_database, no_schema = checks.is_unqualified_table_expression(expression)
-
-        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",
-            )
-
-        transformed = (
+    def _transform(self, expression: exp.Expression) -> exp.Expression:
+        return (
             expression.transform(transforms.upper_case_unquoted_identifiers)
             .transform(transforms.set_schema, current_database=self._conn.database)
             .transform(transforms.create_database, db_path=self._conn.db_path)
@@ -174,6 +154,8 @@ class FakeSnowflakeCursor:
             .transform(transforms.tag)
             .transform(transforms.semi_structured_types)
             .transform(transforms.try_parse_json)
+            # 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)
@@ -212,15 +194,40 @@ class FakeSnowflakeCursor:
             .transform(transforms.create_user)
             .transform(transforms.sha256)
         )
+
+    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")
-        result_sql = None
 
         if transformed.find(exp.Select) and (seed := transformed.args.get("seed")):
             sql = f"SELECT setseed({seed}); {sql}"
 
-        if fs_debug := os.environ.get("FAKESNOW_DEBUG"):
-            debug = command if fs_debug == "snowflake" else sql
-            print(f"{debug};{params=}" if params else f"{debug};", file=sys.stderr)
+        if (fs_debug := os.environ.get("FAKESNOW_DEBUG")) and fs_debug != "snowflake":
+            print(f"{sql};{params=}" if params else f"{sql};", file=sys.stderr)
+
+        result_sql = None
 
         try:
             self._duck_conn.execute(sql, params)
@@ -244,17 +251,12 @@ class FakeSnowflakeCursor:
 
         affected_count = None
 
-        if (maybe_ident := expression.find(exp.Identifier, bfs=False)) and isinstance(maybe_ident.this, str):
-            ident = maybe_ident.this if maybe_ident.quoted else maybe_ident.this.upper()
-        else:
-            ident = None
-
-        if cmd == "USE DATABASE" and ident:
-            self._conn.database = ident
+        if set_database := transformed.args.get("set_database"):
+            self._conn.database = set_database
             self._conn.database_set = True
 
-        elif cmd == "USE SCHEMA" and ident:
-            self._conn.schema = ident
+        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"):
@@ -262,26 +264,6 @@ class FakeSnowflakeCursor:
             self._duck_conn.execute(info_schema.creation_sql(create_db_name))
             result_sql = SQL_CREATED_DATABASE.substitute(name=create_db_name)
 
-        elif 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 == "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
-
         elif cmd == "INSERT":
             (affected_count,) = self._duck_conn.fetchall()[0]
             result_sql = SQL_INSERTED_ROWS.substitute(count=affected_count)
@@ -301,6 +283,28 @@ class FakeSnowflakeCursor:
                 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 == "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

{fakesnow-0.9.7 → fakesnow-0.9.8}/fakesnow/transforms.py

@@ -566,6 +566,9 @@ def json_extract_cast_as_varchar(expression: exp.Expression) -> exp.Expression:
     """
     if (
         isinstance(expression, exp.Cast)
+        and (to := expression.to)
+        and isinstance(to, exp.DataType)
+        and to.this in {exp.DataType.Type.VARCHAR, exp.DataType.Type.TEXT}
         and (je := expression.this)
         and isinstance(je, exp.JSONExtract)
         and (path := je.expression)
@@ -580,7 +583,7 @@ def json_extract_precedence(expression: exp.Expression) -> exp.Expression:
 
     See https://github.com/tekumara/fakesnow/issues/53
     """
-    if isinstance(expression, exp.JSONExtract):
+    if isinstance(expression, (exp.JSONExtract, exp.JSONExtractScalar)):
         return exp.Paren(this=expression)
     return expression
 
@@ -779,7 +782,10 @@ def set_schema(expression: exp.Expression, current_database: str | None) -> exp.
 
         if kind.name.upper() == "DATABASE":
             # duckdb's default schema is main
-            name = f"{expression.this.name}.main"
+            database = expression.this.name
+            return exp.Command(
+                this="SET", expression=exp.Literal.string(f"schema = '{database}.main'"), set_database=database
+            )
         else:
             # SCHEMA
             if db := expression.this.args.get("db"):  # noqa: SIM108
@@ -788,9 +794,10 @@ def set_schema(expression: exp.Expression, current_database: str | None) -> exp.
                 # isn't qualified with a database
                 db_name = current_database or MISSING_DATABASE
 
-            name = f"{db_name}.{expression.this.name}"
-
-        return exp.Command(this="SET", expression=exp.Literal.string(f"schema = '{name}'"))
+            schema = expression.this.name
+            return exp.Command(
+                this="SET", expression=exp.Literal.string(f"schema = '{db_name}.{schema}'"), set_schema=schema
+            )
 
     return expression
 
@@ -1095,6 +1102,21 @@ def timestamp_ntz_ns(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def trim_cast_varchar(expression: exp.Expression) -> exp.Expression:
+    """Snowflake's TRIM casts input to VARCHAR implicitly."""
+
+    if not (isinstance(expression, exp.Trim)):
+        return expression
+
+    operand = expression.this
+    if isinstance(operand, exp.Cast) and operand.to.this in [exp.DataType.Type.VARCHAR, exp.DataType.Type.TEXT]:
+        return expression
+
+    return exp.Trim(
+        this=exp.Cast(this=operand, to=exp.DataType(this=exp.DataType.Type.VARCHAR, nested=False, prefix=False))
+    )
+
+
 def try_parse_json(expression: exp.Expression) -> exp.Expression:
     """Convert TRY_PARSE_JSON() to TRY_CAST(... as JSON).
 

{fakesnow-0.9.7 → fakesnow-0.9.8}/fakesnow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fakesnow
-Version: 0.9.7
+Version: 0.9.8
 Summary: Fake Snowflake Connector for Python. Run, mock and test Snowflake DB locally.
 License:                                  Apache License
                                    Version 2.0, January 2004

{fakesnow-0.9.7 → fakesnow-0.9.8}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "fakesnow"
 description = "Fake Snowflake Connector for Python. Run, mock and test Snowflake DB locally."
-version = "0.9.7"
+version = "0.9.8"
 readme = "README.md"
 license = { file = "LICENSE" }
 classifiers = ["License :: OSI Approved :: MIT License"]

{fakesnow-0.9.7 → fakesnow-0.9.8}/tests/test_fakes.py

@@ -612,6 +612,17 @@ def test_description_delete(dcur: snowflake.connector.cursor.DictCursor):
     # fmt: on
 
 
+def test_description_select(dcur: snowflake.connector.cursor.DictCursor):
+    dcur.execute("SELECT DATEDIFF( DAY, '2023-04-02'::DATE, '2023-04-05'::DATE) as days")
+    assert dcur.fetchall() == [{"DAYS": 3}]
+    # TODO: Snowflake is actually precision=9, is_nullable=False
+    # fmt: off
+    assert dcur.description == [
+        ResultMetadata(name='DAYS', type_code=0, display_size=None, internal_size=None, precision=38, scale=0, is_nullable=True),
+    ]
+    # fmt: on
+
+
 def test_equal_null(cur: snowflake.connector.cursor.SnowflakeCursor):
     cur.execute("select equal_null(NULL, NULL), equal_null(1, 1), equal_null(1, 2), equal_null(1, NULL)")
     assert cur.fetchall() == [(True, True, False, False)]
@@ -792,6 +803,11 @@ def test_get_path_precedence(cur: snowflake.connector.cursor.SnowflakeCursor):
     cur.execute("select {'K1': {'K2': 1}} as col where col:K1:K2 > 0")
     assert indent(cur.fetchall()) == [('{\n  "K1": {\n    "K2": 1\n  }\n}',)]
 
+    cur.execute(
+        """select parse_json('{"K1": "a", "K2": "b"}') as col, case when col:K1::VARCHAR = 'a' and col:K2::VARCHAR = 'b' then 'yes' end"""
+    )
+    assert indent(cur.fetchall()) == [('{\n  "K1": "a",\n  "K2": "b"\n}', "yes")]
+
 
 def test_get_result_batches(cur: snowflake.connector.cursor.SnowflakeCursor):
     # no result set
@@ -1393,6 +1409,14 @@ def test_transactions(conn: snowflake.connector.SnowflakeConnection):
         assert cur.fetchall() == [("Statement executed successfully.",)]
 
 
+def test_trim_cast_varchar(cur: snowflake.connector.cursor.SnowflakeCursor):
+    cur.execute("select trim(1), trim('  name 1  ')")
+    assert cur.fetchall() == [("1", "name 1")]
+
+    cur.execute("""select trim(parse_json('{"k1": "   v11  "}'):k1), trim(parse_json('{"k1": 21}'):k1)""")
+    assert cur.fetchall() == [("v11", "21")]
+
+
 def test_unquoted_identifiers_are_upper_cased(dcur: snowflake.connector.cursor.SnowflakeCursor):
     dcur.execute("create table customers (id int, first_name varchar, last_name varchar)")
     dcur.execute("insert into customers values (1, 'Jenny', 'P')")
@@ -1450,6 +1474,17 @@ def test_values(conn: snowflake.connector.SnowflakeConnection):
         ]
 
 
+def test_json_extract_cast_as_varchar(dcur: snowflake.connector.cursor.DictCursor):
+    dcur.execute("CREATE TABLE example (j VARIANT)")
+    dcur.execute("""INSERT INTO example SELECT PARSE_JSON('{"str": "100", "number" : 100}')""")
+
+    dcur.execute("SELECT j:str::varchar as c_str_varchar, j:number::varchar as c_num_varchar  FROM example")
+    assert dcur.fetchall() == [{"C_STR_VARCHAR": "100", "C_NUM_VARCHAR": "100"}]
+
+    dcur.execute("SELECT j:str::number as c_str_number, j:number::number as c_num_number  FROM example")
+    assert dcur.fetchall() == [{"C_STR_NUMBER": 100, "C_NUM_NUMBER": 100}]
+
+
 def test_write_pandas_quoted_column_names(conn: snowflake.connector.SnowflakeConnection):
     with conn.cursor(snowflake.connector.cursor.DictCursor) as dcur:
         # colunmn names with spaces

{fakesnow-0.9.7 → fakesnow-0.9.8}/tests/test_transforms.py

@@ -44,6 +44,7 @@ from fakesnow.transforms import (
     to_decimal,
     to_timestamp,
     to_timestamp_ntz,
+    trim_cast_varchar,
     try_parse_json,
     try_to_decimal,
     upper_case_unquoted_identifiers,
@@ -446,6 +447,16 @@ def test_json_extract_cast_as_varchar() -> None:
         == """SELECT JSON('{"fruit":"banana"}') ->> '$.fruit'"""
     )
 
+    assert (
+        sqlglot.parse_one(
+            """select parse_json('{"fruit":"9000"}'):fruit::number""",
+            read="snowflake",
+        )
+        .transform(json_extract_cast_as_varchar)
+        .sql(dialect="duckdb")
+        == """SELECT CAST(JSON('{"fruit":"9000"}') -> '$.fruit' AS DECIMAL)"""
+    )
+
 
 def test_json_extract_precedence() -> None:
     assert (
@@ -770,6 +781,23 @@ def test_try_parse_json() -> None:
     )
 
 
+def test_trim_cast_varchar() -> None:
+    assert (
+        sqlglot.parse_one("SELECT TRIM(col) FROM table1").transform(trim_cast_varchar).sql(dialect="duckdb")
+        == "SELECT TRIM(CAST(col AS TEXT)) FROM table1"
+    )
+
+    assert (
+        sqlglot.parse_one("SELECT TRIM(col::varchar) FROM table1").transform(trim_cast_varchar).sql(dialect="duckdb")
+        == "SELECT TRIM(CAST(col AS TEXT)) FROM table1"
+    )
+
+    assert (
+        sqlglot.parse_one("SELECT TRIM(col::TEXT) FROM table1").transform(trim_cast_varchar).sql(dialect="duckdb")
+        == "SELECT TRIM(CAST(col AS TEXT)) FROM table1"
+    )
+
+
 def test_upper_case_unquoted_identifiers() -> None:
     assert (
         sqlglot.parse_one("select name, name as fname from table1").transform(upper_case_unquoted_identifiers).sql()