fakesnow 0.9.18__py3-none-any.whl → 0.9.20__py3-none-any.whl

fakesnow/__init__.py

@@ -13,6 +13,7 @@ import snowflake.connector
 import snowflake.connector.pandas_tools
 
 import fakesnow.fakes as fakes
+from fakesnow.global_database import create_global_database
 
 
 @contextmanager
@@ -52,6 +53,7 @@ def patch(
     assert not isinstance(snowflake.connector.connect, mock.MagicMock), "Snowflake connector is already patched"
 
     duck_conn = duckdb.connect(database=":memory:")
+    create_global_database(duck_conn)
 
     fake_fns = {
         # every time we connect, create a new cursor (ie: connection) so we can isolate each connection's

fakesnow/fakes.py

@@ -32,7 +32,7 @@ import fakesnow.expr as expr
 import fakesnow.info_schema as info_schema
 import fakesnow.macros as macros
 import fakesnow.transforms as transforms
-from fakesnow.global_database import create_global_database
+from fakesnow.variables import Variables
 
 SCHEMA_UNSET = "schema_unset"
 SQL_SUCCESS = "SELECT 'Statement executed successfully.' as 'status'"
@@ -134,6 +134,7 @@ class FakeSnowflakeCursor:
             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
@@ -148,6 +149,7 @@ class FakeSnowflakeCursor:
     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)
@@ -501,6 +503,9 @@ class FakeSnowflakeCursor:
 
         return command, params
 
+    def _inline_variables(self, sql: str) -> str:
+        return self._conn.variables.inline_variables(sql)
+
 
 class FakeSnowflakeConnection:
     def __init__(
@@ -525,8 +530,7 @@ class FakeSnowflakeConnection:
         self.db_path = Path(db_path) if db_path else None
         self.nop_regexes = nop_regexes
         self._paramstyle = snowflake.connector.paramstyle
-
-        create_global_database(duck_conn)
+        self.variables = Variables()
 
         # create database if needed
         if (

fakesnow/transforms.py

@@ -8,6 +8,7 @@ import sqlglot
 from sqlglot import exp
 
 from fakesnow.global_database import USERS_TABLE_FQ_NAME
+from fakesnow.variables import Variables
 
 MISSING_DATABASE = "missing_database"
 SUCCESS_NOP = sqlglot.parse_one("SELECT 'Statement executed successfully.'")
@@ -1401,12 +1402,27 @@ def show_keys(
 
                 if schema:
                     statement += f"AND schema_name = '{schema}' "
+            elif scope_kind == "TABLE":
+                if not table:
+                    raise ValueError(f"SHOW PRIMARY KEYS with {scope_kind} scope requires a table")
+
+                statement += f"AND table_name = '{table.name}' "
             else:
                 raise NotImplementedError(f"SHOW PRIMARY KEYS with {scope_kind} not yet supported")
         return sqlglot.parse_one(statement)
     return expression
 
 
+def update_variables(
+    expression: exp.Expression,
+    variables: Variables,
+) -> exp.Expression:
+    if Variables.is_variable_modifier(expression):
+        variables.update_variables(expression)
+        return SUCCESS_NOP  # Nothing further to do if its a SET/UNSET operation.
+    return expression
+
+
 class SHA256(exp.Func):
     _sql_names: ClassVar = ["SHA256"]
     arg_types: ClassVar = {"this": True}

fakesnow/variables.py

@@ -0,0 +1,56 @@
+import re
+
+import snowflake.connector.errors
+from sqlglot import exp
+
+
+# Implements snowflake variables: https://docs.snowflake.com/en/sql-reference/session-variables#using-variables-in-sql
+class Variables:
+    @classmethod
+    def is_variable_modifier(cls, expr: exp.Expression) -> bool:
+        return isinstance(expr, exp.Set) or cls._is_unset_expression(expr)
+
+    @classmethod
+    def _is_unset_expression(cls, expr: exp.Expression) -> bool:
+        if isinstance(expr, exp.Alias):
+            this_expr = expr.this.args.get("this")
+            return isinstance(this_expr, exp.Expression) and this_expr.this == "UNSET"
+        return False
+
+    def __init__(self) -> None:
+        self._variables = {}
+
+    def update_variables(self, expr: exp.Expression) -> None:
+        if isinstance(expr, exp.Set):
+            unset = expr.args.get("unset")
+            if not unset:  # SET varname = value;
+                unset_expressions = expr.args.get("expressions")
+                assert unset_expressions, "SET without values in expression(s) is unexpected."
+                eq = unset_expressions[0].this
+                name = eq.this.sql()
+                value = eq.args.get("expression").sql()
+                self._set(name, value)
+            else:
+                # Haven't been able to produce this in tests yet due to UNSET being parsed as an Alias expression.
+                raise NotImplementedError("UNSET not supported yet")
+        elif self._is_unset_expression(expr):  # Unfortunately UNSET varname; is parsed as an Alias expression :(
+            alias = expr.args.get("alias")
+            assert alias, "UNSET without value in alias attribute is unexpected."
+            name = alias.this
+            self._unset(name)
+
+    def _set(self, name: str, value: str) -> None:
+        self._variables[name] = value
+
+    def _unset(self, name: str) -> None:
+        self._variables.pop(name)
+
+    def inline_variables(self, sql: str) -> str:
+        for name, value in self._variables.items():
+            sql = re.sub(rf"\${name}", value, sql, flags=re.IGNORECASE)
+
+        if remaining_variables := re.search(r"(?<!\$)\$\w+", sql):
+            raise snowflake.connector.errors.ProgrammingError(
+                msg=f"Session variable '{remaining_variables.group().upper()}' does not exist"
+            )
+        return sql

{fakesnow-0.9.18.dist-info → fakesnow-0.9.20.dist-info}/METADATA

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

{fakesnow-0.9.18.dist-info → fakesnow-0.9.20.dist-info}/RECORD

@@ -1,18 +1,19 @@
-fakesnow/__init__.py,sha256=7040jPjD-bNJDU4Lel42CD_DU47oeu1IsY9Fx_DlslQ,3799
+fakesnow/__init__.py,sha256=gEo6Jg6f8tJUwRTsrfb9q-FpDOufelf8Y3XeWtEJ9wg,3897
 fakesnow/__main__.py,sha256=GDrGyNTvBFuqn_UfDjKs7b3LPtU6gDv1KwosVDrukIM,76
 fakesnow/checks.py,sha256=-QMvdcrRbhN60rnzxLBJ0IkUBWyLR8gGGKKmCS0w9mA,2383
 fakesnow/cli.py,sha256=9qfI-Ssr6mo8UmIlXkUAOz2z2YPBgDsrEVaZv9FjGFs,2201
 fakesnow/expr.py,sha256=CAxuYIUkwI339DQIBzvFF0F-m1tcVGKEPA5rDTzmH9A,892
-fakesnow/fakes.py,sha256=u2EDQ_fXQFFx5bq5pgLrUXV33Flgg5U82L7LIvFiqGw,30556
+fakesnow/fakes.py,sha256=MYzj3XqVokPbOdOcU4WFep1CRsuXlBIwaz9EWaNCGo4,30778
 fakesnow/fixtures.py,sha256=G-NkVeruSQAJ7fvSS2fR2oysUn0Yra1pohHlOvacKEk,455
 fakesnow/global_database.py,sha256=WTVIP1VhNvdCeX7TQncX1TRpGQU5rBf5Pbxim40zeSU,1399
 fakesnow/info_schema.py,sha256=CdIcGXHEQ_kmEAzdQKvA-PX41LA6wlK-4p1J45qgKYA,6266
 fakesnow/macros.py,sha256=pX1YJDnQOkFJSHYUjQ6ErEkYIKvFI6Ncz_au0vv1csA,265
 fakesnow/py.typed,sha256=B-DLSjYBi7pkKjwxCSdpVj2J02wgfJr-E7B1wOUyxYU,80
-fakesnow/transforms.py,sha256=1xok38tft_VowwSm8kKA3P3OEZEeGF09hqfpawxXNaY,52648
-fakesnow-0.9.18.dist-info/LICENSE,sha256=kW-7NWIyaRMQiDpryfSmF2DObDZHGR1cJZ39s6B1Svg,11344
-fakesnow-0.9.18.dist-info/METADATA,sha256=cmS6VqDOHlkSZgfhT7biB03DuSktq4DNtv3EYbFDrrg,17839
-fakesnow-0.9.18.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-fakesnow-0.9.18.dist-info/entry_points.txt,sha256=2riAUgu928ZIHawtO8EsfrMEJhi-EH-z_Vq7Q44xKPM,47
-fakesnow-0.9.18.dist-info/top_level.txt,sha256=500evXI1IFX9so82cizGIEMHAb_dJNPaZvd2H9dcKTA,24
-fakesnow-0.9.18.dist-info/RECORD,,
+fakesnow/transforms.py,sha256=7XS42Ehv3SRciQLpAL4swCFCkNpOTwarXEIX6jUzO4Y,53230
+fakesnow/variables.py,sha256=iE8fnyaMnIoWh_ftT2qUq39fCNglwtQoGbC7MkTqHGU,2341
+fakesnow-0.9.20.dist-info/LICENSE,sha256=kW-7NWIyaRMQiDpryfSmF2DObDZHGR1cJZ39s6B1Svg,11344
+fakesnow-0.9.20.dist-info/METADATA,sha256=RSk43Zbbb8LQsG3G03DwUO8nVbaxs5qNd5VKTHqg4Xw,17839
+fakesnow-0.9.20.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+fakesnow-0.9.20.dist-info/entry_points.txt,sha256=2riAUgu928ZIHawtO8EsfrMEJhi-EH-z_Vq7Q44xKPM,47
+fakesnow-0.9.20.dist-info/top_level.txt,sha256=500evXI1IFX9so82cizGIEMHAb_dJNPaZvd2H9dcKTA,24
+fakesnow-0.9.20.dist-info/RECORD,,