Fixed SQLAlchemy DDL statements (#226)

* Fixed SQLAlchemy DDL statements

* Using black for code formatting

* Moved implementation to higher-level and created tests

Co-authored-by: Vadim Markovtsev <vadim@athenian.co>

Author: Yureien

databases/backends/aiopg.py

@@ -10,6 +10,7 @@
 from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
 from sqlalchemy.engine.result import ResultMetaData, RowProxy
 from sqlalchemy.sql import ClauseElement
+from sqlalchemy.sql.ddl import DDLElement
 from sqlalchemy.types import TypeEngine
 
 from databases.core import DatabaseURL
@@ -181,18 +182,23 @@ def _compile(
         self, query: ClauseElement
     ) -> typing.Tuple[str, dict, CompilationContext]:
         compiled = query.compile(dialect=self._dialect)
-        args = compiled.construct_params()
-        for key, val in args.items():
-            if key in compiled._bind_processors:
-                args[key] = compiled._bind_processors[key](val)
 
         execution_context = self._dialect.execution_ctx_cls()
         execution_context.dialect = self._dialect
-        execution_context.result_column_struct = (
-            compiled._result_columns,
-            compiled._ordered_columns,
-            compiled._textual_ordered_columns,
-        )
+
+        if not isinstance(query, DDLElement):
+            args = compiled.construct_params()
+            for key, val in args.items():
+                if key in compiled._bind_processors:
+                    args[key] = compiled._bind_processors[key](val)
+
+            execution_context.result_column_struct = (
+                compiled._result_columns,
+                compiled._ordered_columns,
+                compiled._textual_ordered_columns,
+            )
+        else:
+            args = {}
 
         logger.debug("Query: %s\nArgs: %s", compiled.string, args)
         return compiled.string, args, CompilationContext(execution_context)

databases/backends/mysql.py

@@ -8,6 +8,7 @@
 from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
 from sqlalchemy.engine.result import ResultMetaData, RowProxy
 from sqlalchemy.sql import ClauseElement
+from sqlalchemy.sql.ddl import DDLElement
 from sqlalchemy.types import TypeEngine
 
 from databases.core import LOG_EXTRA, DatabaseURL
@@ -171,18 +172,23 @@ def _compile(
         self, query: ClauseElement
     ) -> typing.Tuple[str, dict, CompilationContext]:
         compiled = query.compile(dialect=self._dialect)
-        args = compiled.construct_params()
-        for key, val in args.items():
-            if key in compiled._bind_processors:
-                args[key] = compiled._bind_processors[key](val)
 
         execution_context = self._dialect.execution_ctx_cls()
         execution_context.dialect = self._dialect
-        execution_context.result_column_struct = (
-            compiled._result_columns,
-            compiled._ordered_columns,
-            compiled._textual_ordered_columns,
-        )
+
+        if not isinstance(query, DDLElement):
+            args = compiled.construct_params()
+            for key, val in args.items():
+                if key in compiled._bind_processors:
+                    args[key] = compiled._bind_processors[key](val)
+
+            execution_context.result_column_struct = (
+                compiled._result_columns,
+                compiled._ordered_columns,
+                compiled._textual_ordered_columns,
+            )
+        else:
+            args = {}
 
         query_message = compiled.string.replace(" \n", " ").replace("\n", " ")
         logger.debug("Query: %s Args: %s", query_message, repr(args), extra=LOG_EXTRA)

databases/backends/postgres.py

@@ -6,6 +6,7 @@
 from sqlalchemy.dialects.postgresql import pypostgresql
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.sql import ClauseElement
+from sqlalchemy.sql.ddl import DDLElement
 from sqlalchemy.sql.schema import Column
 from sqlalchemy.types import TypeEngine
 
@@ -209,24 +210,32 @@ def transaction(self) -> TransactionBackend:
 
     def _compile(self, query: ClauseElement) -> typing.Tuple[str, list, tuple]:
         compiled = query.compile(dialect=self._dialect)
-        compiled_params = sorted(compiled.params.items())
 
-        mapping = {
-            key: "$" + str(i) for i, (key, _) in enumerate(compiled_params, start=1)
-        }
-        compiled_query = compiled.string % mapping
+        if not isinstance(query, DDLElement):
+            compiled_params = sorted(compiled.params.items())
 
-        processors = compiled._bind_processors
-        args = [
-            processors[key](val) if key in processors else val
-            for key, val in compiled_params
-        ]
+            mapping = {
+                key: "$" + str(i) for i, (key, _) in enumerate(compiled_params, start=1)
+            }
+            compiled_query = compiled.string % mapping
+
+            processors = compiled._bind_processors
+            args = [
+                processors[key](val) if key in processors else val
+                for key, val in compiled_params
+            ]
+
+            result_map = compiled._result_columns
+        else:
+            compiled_query = compiled.string
+            args = []
+            result_map = None
 
         query_message = compiled_query.replace(" \n", " ").replace("\n", " ")
         logger.debug(
             "Query: %s Args: %s", query_message, repr(tuple(args)), extra=LOG_EXTRA
         )
-        return compiled_query, args, compiled._result_columns
+        return compiled_query, args, result_map
 
     @staticmethod
     def _create_column_maps(

databases/backends/sqlite.py

@@ -7,6 +7,7 @@
 from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
 from sqlalchemy.engine.result import ResultMetaData, RowProxy
 from sqlalchemy.sql import ClauseElement
+from sqlalchemy.sql.ddl import DDLElement
 from sqlalchemy.types import TypeEngine
 
 from databases.core import LOG_EXTRA, DatabaseURL
@@ -139,21 +140,25 @@ def _compile(
         self, query: ClauseElement
     ) -> typing.Tuple[str, list, CompilationContext]:
         compiled = query.compile(dialect=self._dialect)
-        args = []
-        for key, raw_val in compiled.construct_params().items():
-            if key in compiled._bind_processors:
-                val = compiled._bind_processors[key](raw_val)
-            else:
-                val = raw_val
-            args.append(val)
 
         execution_context = self._dialect.execution_ctx_cls()
         execution_context.dialect = self._dialect
-        execution_context.result_column_struct = (
-            compiled._result_columns,
-            compiled._ordered_columns,
-            compiled._textual_ordered_columns,
-        )
+
+        args = []
+
+        if not isinstance(query, DDLElement):
+            for key, raw_val in compiled.construct_params().items():
+                if key in compiled._bind_processors:
+                    val = compiled._bind_processors[key](raw_val)
+                else:
+                    val = raw_val
+                args.append(val)
+
+            execution_context.result_column_struct = (
+                compiled._result_columns,
+                compiled._ordered_columns,
+                compiled._textual_ordered_columns,
+            )
 
         query_message = compiled.string.replace(" \n", " ").replace("\n", " ")
         logger.debug(

tests/test_databases.py

@@ -236,6 +236,24 @@ async def test_queries_raw(database_url):
             assert iterate_results[2]["completed"] == True
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_ddl_queries(database_url):
+    """
+    Test that the built-in DDL elements such as `DropTable()`,
+    `CreateTable()` are supported (using SQLAlchemy core).
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # DropTable()
+            query = sqlalchemy.schema.DropTable(notes)
+            await database.execute(query)
+
+            # CreateTable()
+            query = sqlalchemy.schema.CreateTable(notes)
+            await database.execute(query)
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_results_support_mapping_interface(database_url):