diff --git a/opengauss_sqlalchemy/psycopg2.py b/opengauss_sqlalchemy/psycopg2.py index 007565a9eee54c9d72d02830a00b9e9945f58874..f4a7be97f6ec74aa0ebe8242ff69fa6691fe70c4 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 0000000000000000000000000000000000000000..c91ae76cf390f6195a3d5aa0316ea96ba6d2b567 --- /dev/null +++ b/test/test_alembic_dialect.py @@ -0,0 +1,36 @@ +# -*- coding: utf-8 -*- +# Copyright (C) 2005-2022 the SQLAlchemy authors and contributors +# +# +# 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 AlembicDialectTest(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 220480a373308a017ad787b55d7591e87b5dd37f..d3f310c19d8fb8e9c1bfa780e914a88a1cfb112e 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