From d68343103306dbf69a8f8434d854ece7d7441883 Mon Sep 17 00:00:00 2001
From: anyeshenxing <zhangxiaolei21@huawei.com>
Date: Fri, 18 Nov 2022 14:45:29 +0800
Subject: [PATCH 1/2] adapter for alembic and sqlalchemy-migrate

---
 opengauss_sqlalchemy/psycopg2.py | 64 +++++++++++++++++++++++++++++++-
 test/test_alembic_dialect.py     | 36 ++++++++++++++++++
 tox.ini                          |  2 +
 3 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 test/test_alembic_dialect.py

diff --git a/opengauss_sqlalchemy/psycopg2.py b/opengauss_sqlalchemy/psycopg2.py
index 007565a..f4a7be9 100644
--- a/opengauss_sqlalchemy/psycopg2.py
+++ b/opengauss_sqlalchemy/psycopg2.py
@@ -6,13 +6,73 @@
 #
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
-
-from sqlalchemy.dialects.postgresql.psycopg2 import PGCompiler_psycopg2, PGDialect_psycopg2
 from sqlalchemy import schema
 from sqlalchemy import util
+from sqlalchemy.dialects.postgresql.psycopg2 import PGCompiler_psycopg2, PGDialect_psycopg2
+from sqlalchemy.ext.compiler import compiles
 
 from opengauss_sqlalchemy.base import OpenGaussDDLCompiler, OpenGaussIdentifierPreparer
 
+# If alembic is installed, register an alias in its dialect mapping.
+try:
+    import alembic
+except ImportError:
+    pass
+else:
+    from alembic.ddl import postgresql
+    from alembic.ddl.base import RenameTable
+
+    compiles(RenameTable, 'opengauss')(postgresql.visit_rename_table)
+    compiles(postgresql.PostgresqlColumnType, "opengauss")(postgresql.visit_column_type)
+
+
+    class OpenGaussImpl(postgresql.PostgresqlImpl):
+        __dialect__ = 'opengauss'
+
+# If sqlalchemy-migrate is installed, register there too.
+try:
+    from migrate.changeset.databases.visitor import DIALECTS as migrate_dialects
+except ImportError:
+    pass
+else:
+    from migrate.changeset import ansisql
+
+
+    class OGColumnGenerator(OpenGaussDDLCompiler, ansisql.ANSIColumnGenerator):
+        """OpenGauss column generator implementation."""
+        pass
+
+
+    class OGColumnDropper(ansisql.ANSIColumnDropper):
+        """OpenGauss column dropper implementation."""
+        pass
+
+
+    class OGSchemaChanger(ansisql.ANSISchemaChanger):
+        """OpenGauss schema changer implementation."""
+        pass
+
+
+    class OGConstraintGenerator(ansisql.ANSIConstraintGenerator):
+        """OpenGauss constraint generator implementation."""
+        pass
+
+
+    class OGConstraintDropper(ansisql.ANSIConstraintDropper):
+        """OpenGauss constaint dropper implementation."""
+        pass
+
+
+    class OGDialect(ansisql.ANSIDialect):
+        columngenerator = OGColumnGenerator
+        columndropper = OGColumnDropper
+        schemachanger = OGSchemaChanger
+        constraintgenerator = OGConstraintGenerator
+        constraintdropper = OGConstraintDropper
+
+
+    migrate_dialects["opengauss"] = OGDialect
+
 
 class OpenGaussCompiler_psycopg2(PGCompiler_psycopg2):
     def get_cte_preamble(self, recursive):
diff --git a/test/test_alembic_dialect.py b/test/test_alembic_dialect.py
new file mode 100644
index 0000000..1f18e05
--- /dev/null
+++ b/test/test_alembic_dialect.py
@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2005-2022 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# Copyright (C) 2021-2022 Huawei Technologies Co.,Ltd.
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+from alembic import migration
+from alembic.ddl.base import RenameTable
+from alembic.ddl.postgresql import PostgresqlColumnType
+from sqlalchemy import types
+from sqlalchemy.testing import fixtures
+
+from opengauss_sqlalchemy import psycopg2
+
+
+class AlebicDialectTest(fixtures.TestBase):
+
+    def test_configure_migration_context(self):
+        context = migration.MigrationContext.configure(
+            url='opengauss+psycopg2://mydb'
+        )
+        assert isinstance(context.impl, psycopg2.OpenGaussImpl)
+
+    def test_rename_table(self):
+        compiler = psycopg2.OpenGaussDDLCompiler(psycopg2.dialect(), None)
+        sql = compiler.process(RenameTable("old", "new", "schema"))
+        assert sql == 'ALTER TABLE schema."old" RENAME TO "new"'
+
+    def test_alter_column_type(self):
+        compiler = psycopg2.OpenGaussDDLCompiler(psycopg2.dialect(), None)
+        sql = compiler.process(
+            PostgresqlColumnType("table_name", "column_name", types.INTEGER)
+        )
+        assert sql == "ALTER TABLE table_name ALTER COLUMN column_name TYPE INTEGER "
\ No newline at end of file
diff --git a/tox.ini b/tox.ini
index 220480a..d3f310c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -8,6 +8,8 @@ deps =
     pytest>=6.2,<8; python_version >= '3'
     mock; python_version < '3.3'
     psycopg2>=2.8.6,<2.9; python_version < '3'
+    alembic>=0.9.7
+    sqlalchemy-migrate==0.13.0
 
 base_command=python -m pytest --rootdir {toxinidir} --maxfail 1
 
-- 
Gitee


From 2f79a83bbbd841c0b8405a413e4e05a7c6542a03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=B4=BE=E5=B3=BB=E8=8B=8F?= <jiajunsu@huawei.com>
Date: Fri, 18 Nov 2022 10:51:16 +0000
Subject: [PATCH 2/2] Update test/test_alembic_dialect.py

---
 test/test_alembic_dialect.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_alembic_dialect.py b/test/test_alembic_dialect.py
index 1f18e05..c91ae76 100644
--- a/test/test_alembic_dialect.py
+++ b/test/test_alembic_dialect.py
@@ -15,7 +15,7 @@ from sqlalchemy.testing import fixtures
 from opengauss_sqlalchemy import psycopg2
 
 
-class AlebicDialectTest(fixtures.TestBase):
+class AlembicDialectTest(fixtures.TestBase):
 
     def test_configure_migration_context(self):
         context = migration.MigrationContext.configure(
-- 
Gitee