Fix not fetching cursor (#59)

* Fix not fetching cursor on insert/update

* Fetch the cursor on insert/update/delete/copy into

* Enable CTE tests

* Enable further tests

* Support for table and column comments

* Include CTE test now bug is fixed

* Run against nightly

* Update pipenv

* Work in SQLAlchemy 1.4

Author: rad-pat

.github/workflows/test.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     services:
       databend:
-        image: datafuselabs/databend
+        image: datafuselabs/databend:nightly
         env:
           QUERY_DEFAULT_USER: databend
           QUERY_DEFAULT_PASSWORD: databend

Pipfile.lock

@@ -49,8 +49,13 @@ def escape_item(self, item):
             return self.escape_number(item)
         elif isinstance(item, timedelta):
             return self.escape_string(f"{item.total_seconds()} seconds") + "::interval"
-        elif isinstance(item, (datetime, date, time, timedelta)):
-            return self.escape_string(item.strftime("%Y-%m-%d %H:%M:%S"))
+        elif isinstance(item, time):
+            # N.B. Date here must match date in DatabendTime.literal_processor - 1970-01-01
+            return self.escape_string(item.strftime("1970-01-01 %H:%M:%S.%f")) + "::timestamp"
+        elif isinstance(item, datetime):
+            return self.escape_string(item.strftime("%Y-%m-%d %H:%M:%S.%f")) + "::timestamp"
+        elif isinstance(item, date):
+            return self.escape_string(item.strftime("%Y-%m-%d")) + "::date"
         else:
             return self.escape_string(item)
 

databend_sqlalchemy/connector.py

@@ -30,6 +30,7 @@
 import datetime
 from types import NoneType
 
+import sqlalchemy.engine.reflection
 import sqlalchemy.types as sqltypes
 from typing import Any, Dict, Optional, Union
 from sqlalchemy import util as sa_util
@@ -44,6 +45,7 @@
     Subquery,
 )
 from sqlalchemy.dialects.postgresql.base import PGCompiler, PGIdentifierPreparer
+from sqlalchemy import Table, MetaData, Column
 from sqlalchemy.types import (
     BIGINT,
     INTEGER,
@@ -670,7 +672,7 @@ def process(value):
 class DatabendDateTime(sqltypes.DATETIME):
     __visit_name__ = "DATETIME"
 
-    _reg = re.compile(r"(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)")
+    _reg = re.compile(r"(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)\.(\d+)")
 
     def result_processor(self, dialect, coltype):
         def process(value):
@@ -698,7 +700,7 @@ def process(value):
 class DatabendTime(sqltypes.TIME):
     __visit_name__ = "TIME"
 
-    _reg = re.compile(r"(?:\d+)-(?:\d+)-(?:\d+) (\d+):(\d+):(\d+)")
+    _reg = re.compile(r"(?:\d+)-(?:\d+)-(?:\d+) (\d+):(\d+):(\d+)\.(\d+)")
 
     def result_processor(self, dialect, coltype):
         def process(value):
@@ -720,7 +722,7 @@ def literal_processor(self, dialect):
         def process(value):
             if value is not None:
                 from_min_value = datetime.datetime.combine(
-                    datetime.date(1000, 1, 1), value
+                    datetime.date(1970, 1, 1), value
                 )
                 time_str = from_min_value.isoformat(timespec="microseconds")
                 return f"'{time_str}'"
@@ -800,6 +802,9 @@ class DatabendIdentifierPreparer(PGIdentifierPreparer):
 
 
 class DatabendCompiler(PGCompiler):
+    iscopyintotable: bool = False
+    iscopyintolocation: bool = False
+
     def get_select_precolumns(self, select, **kw):
         # call the base implementation because Databend doesn't support DISTINCT ON
         return super(PGCompiler, self).get_select_precolumns(select, **kw)
@@ -971,6 +976,11 @@ def visit_when_merge_unmatched(self, merge_unmatched, **kw):
         )
 
     def visit_copy_into(self, copy_into, **kw):
+        if isinstance(copy_into.target, (TableClause,)):
+            self.iscopyintotable = True
+        else:
+            self.iscopyintolocation = True
+
         target = (
             self.preparer.format_table(copy_into.target)
             if isinstance(copy_into.target, (TableClause,))
@@ -1090,8 +1100,21 @@ def visit_google_cloud_storage(self, gcs: GoogleCloudStorage, **kw):
             f")"
         )
 
+    def visit_stage(self, stage, **kw):
+        if stage.path:
+            return f"@{stage.name}/{stage.path}"
+        return f"@{stage.name}"
+
 
 class DatabendExecutionContext(default.DefaultExecutionContext):
+    iscopyintotable = False
+    iscopyintolocation = False
+
+    _copy_input_bytes: Optional[int] = None
+    _copy_output_bytes: Optional[int] = None
+    _copy_into_table_results: Optional[list[dict]] = None
+    _copy_into_location_results: dict = None
+
     @sa_util.memoized_property
     def should_autocommit(self):
         return False  # No DML supported, never autocommit
@@ -1102,6 +1125,38 @@ def create_server_side_cursor(self):
     def create_default_cursor(self):
         return self._dbapi_connection.cursor()
 
+    def post_exec(self):
+        self.iscopyintotable = getattr(self.compiled, 'iscopyintotable', False)
+        self.iscopyintolocation = getattr(self.compiled, 'iscopyintolocation', False)
+        if (self.isinsert or self.isupdate or self.isdelete or
+            self.iscopyintolocation or self.iscopyintotable):
+            result = self.cursor.fetchall()
+            if self.iscopyintotable:
+                self._copy_into_table_results = [
+                    {
+                        'file': row[0],
+                        'rows_loaded': row[1],
+                        'errors_seen': row[2],
+                        'first_error': row[3],
+                        'first_error_line': row[4],
+                    } for row in result
+                ]
+                self._rowcount = sum(c['rows_loaded'] for c in self._copy_into_table_results)
+            else:
+                self._rowcount = result[0][0]
+                if self.iscopyintolocation:
+                    self._copy_into_location_results = {
+                        'rows_unloaded': result[0][0],
+                        'input_bytes': result[0][1],
+                        'output_bytes': result[0][2],
+                    }
+
+    def copy_into_table_results(self) -> list[dict]:
+        return self._copy_into_table_results
+
+    def copy_into_location_results(self) -> dict:
+        return self._copy_into_location_results
+
 
 class DatabendTypeCompiler(compiler.GenericTypeCompiler):
     def visit_ARRAY(self, type_, **kw):
@@ -1171,6 +1226,12 @@ def post_create_table(self, table):
         if engine is not None:
             table_opts.append(f" ENGINE={engine}")
 
+        if table.comment is not None:
+            comment = self.sql_compiler.render_literal_value(
+                table.comment, sqltypes.String()
+            )
+            table_opts.append(f" COMMENT={comment}")
+
         cluster_keys = db_opts.get("cluster_by")
         if cluster_keys is not None:
             if isinstance(cluster_keys, str):
@@ -1192,6 +1253,37 @@ def post_create_table(self, table):
 
         return " ".join(table_opts)
 
+    def get_column_specification(self, column, **kwargs):
+        colspec = super().get_column_specification(column, **kwargs)
+        comment = column.comment
+        if comment is not None:
+            literal = self.sql_compiler.render_literal_value(
+                comment, sqltypes.String()
+            )
+            colspec += " COMMENT " + literal
+
+        return colspec
+
+    def visit_set_table_comment(self, create, **kw):
+        return "ALTER TABLE %s COMMENT = %s" % (
+            self.preparer.format_table(create.element),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_table_comment(self, create, **kw):
+        return "ALTER TABLE %s COMMENT = ''" % (
+            self.preparer.format_table(create.element)
+        )
+
+    def visit_set_column_comment(self, create, **kw):
+        return "ALTER TABLE %s MODIFY %s %s" % (
+            self.preparer.format_table(create.element.table),
+            self.preparer.format_column(create.element),
+            self.get_column_specification(create.element),
+        )
+
 
 class DatabendDialect(default.DefaultDialect):
     name = "databend"
@@ -1204,7 +1296,7 @@ class DatabendDialect(default.DefaultDialect):
     supports_alter = True
     supports_comments = False
     supports_empty_insert = False
-    supports_is_distinct_from = False
+    supports_is_distinct_from = True
     supports_multivalues_insert = True
 
     supports_statement_cache = False
@@ -1316,7 +1408,7 @@ def has_table(self, connection, table_name, schema=None, **kw):
     def get_columns(self, connection, table_name, schema=None, **kw):
         query = text(
             """
-            select column_name, column_type, is_nullable
+            select column_name, column_type, is_nullable, nullif(column_comment, '')
             from information_schema.columns
             where table_name = :table_name
             and table_schema = :schema_name
@@ -1337,6 +1429,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
                 "type": self._get_column_type(row[1]),
                 "nullable": get_is_nullable(row[2]),
                 "default": None,
+                "comment": row[3],
             }
             for row in result
         ]
@@ -1416,6 +1509,23 @@ def get_table_names(self, connection, schema=None, **kw):
         result = connection.execute(query, dict(schema_name=schema))
         return [row[0] for row in result]
 
+    @reflection.cache
+    def get_temp_table_names(self, connection, schema=None, **kw):
+        table_name_query = """
+            select name
+            from system.temporary_tables
+            where database = :schema_name
+            """
+        query = text(table_name_query).bindparams(
+            bindparam("schema_name", type_=sqltypes.Unicode)
+        )
+        if schema is None:
+            schema = self.default_schema_name
+
+        result = connection.execute(query, dict(schema_name=schema))
+        return [row[0] for row in result]
+
+
     @reflection.cache
     def get_view_names(self, connection, schema=None, **kw):
         view_name_query = """
@@ -1510,6 +1620,82 @@ def get_table_options(self, connection, table_name, schema=None, **kw):
 
         return options
 
+    @reflection.cache
+    def get_table_comment(self, connection, table_name, schema, **kw):
+        query_text = """
+            SELECT comment
+            FROM system.tables
+            WHERE database = :schema_name
+            and name = :table_name
+            """
+        query = text(query_text).bindparams(
+            bindparam("table_name", type_=sqltypes.Unicode),
+            bindparam("schema_name", type_=sqltypes.Unicode),
+        )
+        if schema is None:
+            schema = self.default_schema_name
+
+        result = connection.execute(
+            query, dict(table_name=table_name, schema_name=schema)
+        ).one_or_none()
+        if not result:
+            raise NoSuchTableError(
+                f"{self.identifier_preparer.quote_identifier(schema)}."
+                f"{self.identifier_preparer.quote_identifier(table_name)}"
+            )
+        return {'text': result[0]} if result[0] else reflection.ReflectionDefaults.table_comment() if hasattr(reflection, 'ReflectionDefault') else {'text': None}
+
+    def _prepare_filter_names(self, filter_names):
+        if filter_names:
+            fn = [name for name in filter_names]
+            return True, {"filter_names": fn}
+        else:
+            return False, {}
+
+    def get_multi_table_comment(
+        self, connection, schema, filter_names, scope, kind, **kw
+    ):
+        meta = MetaData()
+        all_tab_comments=Table(
+            "tables",
+            meta,
+            Column("database", VARCHAR, nullable=False),
+            Column("name", VARCHAR, nullable=False),
+            Column("comment", VARCHAR),
+            Column("table_type", VARCHAR),
+            schema='system',
+        ).alias("a_tab_comments")
+
+
+        has_filter_names, params = self._prepare_filter_names(filter_names)
+        owner = schema or self.default_schema_name
+
+        table_types = set()
+        if reflection.ObjectKind.TABLE in kind:
+            table_types.add('BASE TABLE')
+        if reflection.ObjectKind.VIEW in kind:
+            table_types.add('VIEW')
+
+        query = select(
+            all_tab_comments.c.name, all_tab_comments.c.comment
+        ).where(
+            all_tab_comments.c.database == owner,
+            all_tab_comments.c.table_type.in_(table_types),
+            sqlalchemy.true() if reflection.ObjectScope.DEFAULT in scope else sqlalchemy.false(),
+        )
+        if has_filter_names:
+            query = query.where(all_tab_comments.c.name.in_(bindparam("filter_names")))
+
+        result = connection.execute(query, params)
+        default_comment = reflection.ReflectionDefaults.table_comment
+        return (
+            (
+                (schema, table),
+                {"text": comment} if comment else default_comment(),
+            )
+            for table, comment in result
+        )
+
     def do_rollback(self, dbapi_connection):
         # No transactions
         pass

databend_sqlalchemy/databend_dialect.py

@@ -1,7 +1,7 @@
 
 from sqlalchemy.testing.provision import create_db
 from sqlalchemy.testing.provision import drop_db
-from sqlalchemy.testing.provision import configure_follower, update_db_opts
+from sqlalchemy.testing.provision import configure_follower, update_db_opts, temp_table_keyword_args
 
 
 @create_db.for_db("databend")
@@ -31,6 +31,10 @@ def _databend_drop_db(cfg, eng, ident):
         conn.exec_driver_sql("DROP DATABASE IF EXISTS %s_test_schema_2" % ident)
         conn.exec_driver_sql("DROP DATABASE IF EXISTS %s" % ident)
 
+@temp_table_keyword_args.for_db("databend")
+def _databend_temp_table_keyword_args(cfg, eng):
+    return {"prefixes": ["TEMPORARY"]}
+
 
 @configure_follower.for_db("databend")
 def _databend_configure_follower(config, ident):
@@ -39,5 +43,5 @@ def _databend_configure_follower(config, ident):
 
 # Uncomment to debug SQL Statements in tests
 # @update_db_opts.for_db("databend")
-# def _mssql_update_db_opts(db_url, db_opts):
+# def _databend_update_db_opts(db_url, db_opts):
 #     db_opts["echo"] = True