SNOW-619615: SQLAlchemy dialect compliance (#314)

Author: sfc-gh-aling

src/snowflake/sqlalchemy/base.py

@@ -359,13 +359,50 @@ def update_from_clause(
             for t in extra_froms
         )
 
+    def _get_regexp_args(self, binary, kw):
+        string = self.process(binary.left, **kw)
+        pattern = self.process(binary.right, **kw)
+        flags = binary.modifiers["flags"]
+        if flags is not None:
+            flags = self.process(flags, **kw)
+        return string, pattern, flags
+
+    def visit_regexp_match_op_binary(self, binary, operator, **kw):
+        string, pattern, flags = self._get_regexp_args(binary, kw)
+        if flags is None:
+            return f"REGEXP_LIKE({string}, {pattern})"
+        else:
+            return f"REGEXP_LIKE({string}, {pattern}, {flags})"
+
+    def visit_regexp_replace_op_binary(self, binary, operator, **kw):
+        string, pattern, flags = self._get_regexp_args(binary, kw)
+        replacement = self.process(binary.modifiers["replacement"], **kw)
+        if flags is None:
+            return "REGEXP_REPLACE({}, {}, {})".format(
+                string,
+                pattern,
+                replacement,
+            )
+        else:
+            return "REGEXP_REPLACE({}, {}, {}, {})".format(
+                string,
+                pattern,
+                replacement,
+                flags,
+            )
+
+    def visit_not_regexp_match_op_binary(self, binary, operator, **kw):
+        return f"NOT {self.visit_regexp_match_op_binary(binary, operator, **kw)}"
+
+    def render_literal_value(self, value, type_):
+        # escape backslash
+        return super().render_literal_value(value, type_).replace("\\", "\\\\")
+
 
 class SnowflakeExecutionContext(default.DefaultExecutionContext):
     def fire_sequence(self, seq, type_):
         return self._execute_scalar(
-            "SELECT "
-            + self.dialect.identifier_preparer.format_sequence(seq)
-            + ".nextval",
+            f"SELECT {self.identifier_preparer.format_sequence(seq)}.nextval",
             type_,
         )
 
@@ -387,6 +424,35 @@ def should_autocommit(self):
         else:
             return autocommit and not self.isddl
 
+    def pre_exec(self):
+        if self.compiled:
+            # for compiled statements, percent is doubled for escape, we turn on _interpolate_empty_sequences
+            if hasattr(self._dbapi_connection, "driver_connection"):
+                # _dbapi_connection is a _ConnectionFairy which proxies raw SnowflakeConnection
+                self._dbapi_connection.driver_connection._interpolate_empty_sequences = (
+                    True
+                )
+            else:
+                # _dbapi_connection is a raw SnowflakeConnection
+                self._dbapi_connection._interpolate_empty_sequences = True
+
+    def post_exec(self):
+        if self.compiled:
+            # for compiled statements, percent is doubled for escapeafter execution
+            # we reset _interpolate_empty_sequences to false which is turned on in pre_exec
+            if hasattr(self._dbapi_connection, "driver_connection"):
+                # _dbapi_connection is a _ConnectionFairy which proxies raw SnowflakeConnection
+                self._dbapi_connection.driver_connection._interpolate_empty_sequences = (
+                    False
+                )
+            else:
+                # _dbapi_connection is a raw SnowflakeConnection
+                self._dbapi_connection._interpolate_empty_sequences = False
+
+    @property
+    def rowcount(self):
+        return self.cursor.rowcount
+
 
 class SnowflakeDDLCompiler(compiler.DDLCompiler):
     def denormalize_column_name(self, name):
@@ -406,6 +472,10 @@ def get_column_specification(self, column, **kwargs):
             self.dialect.type_compiler.process(column.type, type_expression=column),
         ]
 
+        has_identity = (
+            column.identity is not None and self.dialect.supports_identity_columns
+        )
+
         if not column.nullable:
             colspec.append("NOT NULL")
 
@@ -422,10 +492,15 @@ def get_column_specification(self, column, **kwargs):
             and column.server_default is None
         ):
             if isinstance(column.default, Sequence):
-                colspec.append(f"DEFAULT {column.default.name}.nextval")
+                colspec.append(
+                    f"DEFAULT {self.dialect.identifier_preparer.format_sequence(column.default)}.nextval"
+                )
             else:
                 colspec.append("AUTOINCREMENT")
 
+        if has_identity:
+            colspec.append(self.process(column.identity))
+
         return " ".join(colspec)
 
     def post_create_table(self, table):
@@ -512,6 +587,14 @@ def visit_drop_column_comment(self, drop, **kw):
             self.preparer.format_column(drop.element),
         )
 
+    def visit_identity_column(self, identity, **kw):
+        text = " IDENTITY"
+        if identity.start is not None or identity.increment is not None:
+            start = 1 if identity.start is None else identity.start
+            increment = 1 if identity.increment is None else identity.increment
+            text += f"({start},{increment})"
+        return text
+
 
 class SnowflakeTypeCompiler(compiler.GenericTypeCompiler):
     def visit_BYTEINT(self, type_, **kw):

src/snowflake/sqlalchemy/custom_types.py

@@ -1,6 +1,9 @@
 #
 # Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
 #
+import datetime
+import decimal
+import re
 
 import sqlalchemy.types as sqltypes
 import sqlalchemy.util as util
@@ -17,6 +20,16 @@
 VARBINARY = sqltypes.BINARY
 
 
+def _process_float(value):
+    if value == float("inf"):
+        return "inf"
+    elif value == float("-inf"):
+        return "-inf"
+    elif value is not None:
+        return float(value)
+    return value
+
+
 class SnowflakeType(sqltypes.TypeEngine):
     def _default_dialect(self):
         # Get around circular import
@@ -51,7 +64,98 @@ class GEOGRAPHY(SnowflakeType):
     __visit_name__ = "GEOGRAPHY"
 
 
+class _CUSTOM_Date(SnowflakeType, sqltypes.Date):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                return f"'{value.isoformat()}'"
+
+        return process
+
+    _reg = re.compile(r"(\d+)-(\d+)-(\d+)")
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if isinstance(value, str):
+                m = self._reg.match(value)
+                if not m:
+                    raise ValueError(f"could not parse {value!r} as a date value")
+                return datetime.date(*[int(x or 0) for x in m.groups()])
+            else:
+                return value
+
+        return process
+
+
+class _CUSTOM_DateTime(SnowflakeType, sqltypes.DateTime):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                datetime_str = value.isoformat(" ", timespec="microseconds")
+                return f"'{datetime_str}'"
+
+        return process
+
+    _reg = re.compile(r"(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)(?:\.(\d{0,6}))?")
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if isinstance(value, str):
+                m = self._reg.match(value)
+                if not m:
+                    raise ValueError(f"could not parse {value!r} as a datetime value")
+                return datetime.datetime(*[int(x or 0) for x in m.groups()])
+            else:
+                return value
+
+        return process
+
+
+class _CUSTOM_Time(SnowflakeType, sqltypes.Time):
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is not None:
+                time_str = value.isoformat(timespec="microseconds")
+                return f"'{time_str}'"
+
+        return process
+
+    _reg = re.compile(r"(\d+):(\d+):(\d+)(?:\.(\d{0,6}))?")
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if isinstance(value, str):
+                m = self._reg.match(value)
+                if not m:
+                    raise ValueError(f"could not parse {value!r} as a time value")
+                return datetime.time(*[int(x or 0) for x in m.groups()])
+            else:
+                return value
+
+        return process
+
+
+class _CUSTOM_Float(SnowflakeType, sqltypes.Float):
+    def bind_processor(self, dialect):
+        return _process_float
+
+
 class _CUSTOM_DECIMAL(SnowflakeType, sqltypes.DECIMAL):
     @util.memoized_property
     def _type_affinity(self):
         return sqltypes.INTEGER if self.scale == 0 else sqltypes.DECIMAL
+
+
+class _CUSTOM_Numeric(SnowflakeType, sqltypes.Numeric):
+    def result_processor(self, dialect, coltype):
+        if self.asdecimal:
+
+            def process(value):
+                if value:
+                    return decimal.Decimal(value)
+                else:
+                    return None
+
+            return process
+        else:
+            return _process_float

src/snowflake/sqlalchemy/provision.py

@@ -0,0 +1,12 @@
+#
+# Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
+#
+from sqlalchemy.testing.provision import set_default_schema_on_connection
+
+
+# This is only for test purpose required by Requirement "default_schema_name_switch"
+@set_default_schema_on_connection.for_db("snowflake")
+def _snowflake_set_default_schema_on_connection(cfg, dbapi_connection, schema_name):
+    cursor = dbapi_connection.cursor()
+    cursor.execute(f"USE SCHEMA {dbapi_connection.database}.{schema_name};")
+    cursor.close()