Merge pull request #12 from saritasa-nest/feature/alembic

Add alembic feature

alembic.ini

@@ -0,0 +1,54 @@
+# A generic, single database configuration.
+# https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+
+[alembic]
+# path to migration scripts
+script_location = tests/alembic
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# version path separator; This is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console, rich
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = rich
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[handler_rich]
+class = rich.logging.RichHandler
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

pyproject.toml

@@ -40,10 +40,15 @@ factory-boy  = {version= "<4", optional = true}
 # Data validation using Python type hints
 # https://docs.pydantic.dev/latest/
 pydantic = {version= "<3", optional = true}
+# Alembic is a lightweight database migration tool for usage with the
+# SQLAlchemy Database Toolkit for Python
+# https://alembic.sqlalchemy.org/
+alembic = {version= "<2", optional = true}
 
 [tool.poetry.extras]
 factories = ["factory-boy"]
 auto_schema = ["pydantic"]
+migrations = ["alembic"]
 
 [tool.poetry.group.dev.dependencies]
 # Improved REPL
@@ -251,6 +256,7 @@ asyncio_mode = "auto"
 [tool.coverage.run]
 omit = [
   "saritasa_sqlalchemy_tools/session.py",
+  "saritasa_sqlalchemy_tools/alembic.py",
 ]
 
 # https://docformatter.readthedocs.io/en/latest/configuration.html#

saritasa_sqlalchemy_tools/__init__.py

@@ -1,8 +1,5 @@
-from .auto_schema import (
-    ModelAutoSchema,
-    ModelAutoSchemaError,
-    ModelAutoSchemaT,
-)
+import contextlib
+
 from .models import (
     BaseIDModel,
     BaseModel,
@@ -50,9 +47,22 @@
     get_async_engine,
     get_async_session_factory,
 )
-from .testing import AsyncSQLAlchemyModelFactory, AsyncSQLAlchemyOptions
+
+with contextlib.suppress(ImportError):
+    from .testing import AsyncSQLAlchemyModelFactory, AsyncSQLAlchemyOptions
+
+with contextlib.suppress(ImportError):
+    from .alembic import AlembicMigrations
+
+with contextlib.suppress(ImportError):
+    from .auto_schema import (
+        ModelAutoSchema,
+        ModelAutoSchemaError,
+        ModelAutoSchemaT,
+    )
 
 __all__ = (
+    "AlembicMigrations",
     "AsyncSQLAlchemyModelFactory",
     "AsyncSQLAlchemyOptions",
     "Session",

saritasa_sqlalchemy_tools/alembic.py

@@ -0,0 +1,163 @@
+import asyncio
+import dataclasses
+import importlib
+import logging.config
+import typing
+
+import alembic
+import alembic.config
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+import sqlalchemy.sql.schema
+
+
+@dataclasses.dataclass
+class AlembicMigrations:
+    """Class for managing alembic migrations."""
+
+    target_metadata: sqlalchemy.sql.schema.MetaData
+    db_url: sqlalchemy.engine.URL | None = None
+    db_driver: str = ""
+    db_user: str = ""
+    db_name: str = ""
+    db_password: str = ""
+    db_host: str = ""
+    db_port: int = 0
+    db_schema: str = ""
+    query: dict[str, typing.Any] | None = None
+    plugins: tuple[str, ...] = ()
+
+    @property
+    def alembic_config(self) -> alembic.config.Config:
+        """Get alembic config."""
+        return alembic.context.config
+
+    @property
+    def database_url(self) -> str | sqlalchemy.engine.URL:
+        """Get database url."""
+        if self.db_url:
+            return self.db_url
+        if database_url := self.alembic_config.get_main_option(
+            "sqlalchemy.url",
+        ):
+            return database_url
+        return sqlalchemy.engine.URL(
+            drivername=self.db_driver,
+            username=self.db_user,
+            password=self.db_password,
+            host=self.db_host,
+            port=self.db_port,
+            database=self.db_name,
+            query=self.query or {},  # type: ignore
+        )
+
+    @property
+    def logger(self) -> logging.Logger:
+        """Get logger."""
+        return logging.getLogger("alembic")
+
+    def setup_config(self) -> None:
+        """Set up config."""
+        if self.alembic_config.config_file_name is not None:
+            logging.config.fileConfig(self.alembic_config.config_file_name)
+
+    def import_plugins(self) -> None:
+        """Import plugins."""
+        self.logger.info("Set up plugins.")
+        for plugin in self.plugins:
+            self.logger.info(f"Importing plugin '{plugin}'.")
+            importlib.import_module(plugin)
+
+    def run_migrations_offline(self) -> None:
+        """Run migrations in 'offline' mode.
+
+        This configures the context with just a URL
+        and not an Engine, though an Engine is acceptable
+        here as well.  By skipping the Engine creation
+        we don't even need a DBAPI to be available.
+
+        Calls to context.execute() here emit the given string to the
+        script output.
+
+        """
+        alembic.context.configure(
+            url=self.database_url,
+            target_metadata=self.target_metadata,
+            literal_binds=True,
+            dialect_opts={
+                "paramstyle": "named",
+            },
+        )
+
+        with alembic.context.begin_transaction():
+            alembic.context.run_migrations()
+
+    def set_up_schema(
+        self,
+        connection: sqlalchemy.engine.Connection,
+    ) -> None:
+        """Set up schema for migrations."""
+        self.logger.info(f"Setting up schema '{self.db_schema}'.")
+        connection.execute(
+            sqlalchemy.text(f"create schema if not exists {self.db_schema}"),
+        )
+        # set search path on the connection, which ensures that
+        # PostgreSQL will emit all CREATE / ALTER / DROP statements
+        # in terms of this schema by default
+        connection.execute(
+            sqlalchemy.text(f"set search_path to {self.db_schema}"),
+        )
+        connection.commit()
+
+        # make use of non-supported SQLAlchemy attribute to ensure
+        # the dialect reflects tables in terms of the current schema name
+        connection.dialect.default_schema_name = self.db_schema
+
+    def do_run_migrations(
+        self,
+        connection: sqlalchemy.engine.Connection,
+    ) -> None:
+        """Apply migrations."""
+        if self.db_schema:
+            self.set_up_schema(connection=connection)
+        alembic.context.configure(
+            connection=connection,
+            target_metadata=self.target_metadata,
+        )
+        with alembic.context.begin_transaction():
+            alembic.context.run_migrations()
+
+    async def run_async_migrations(self) -> None:
+        """Run async migrations.
+
+        In this scenario we need to create an Engine and associate a connection
+        with the context.
+
+        """
+        connectable = sqlalchemy.ext.asyncio.async_engine_from_config(
+            self.alembic_config.get_section(
+                self.alembic_config.config_ini_section,
+                {},
+            ),
+            prefix="sqlalchemy.",
+            poolclass=sqlalchemy.pool.NullPool,
+            url=self.database_url,
+        )
+
+        async with connectable.connect() as connection:
+            await connection.run_sync(self.do_run_migrations)
+
+        await connectable.dispose()
+
+    def run_migrations_online(self) -> None:
+        """Run migrations in 'online' mode."""
+        asyncio.run(self.run_async_migrations())
+
+    def run(self) -> None:
+        """Run migrations."""
+        self.setup_config()
+        self.import_plugins()
+        if alembic.context.is_offline_mode():
+            self.run_migrations_offline()
+        else:
+            self.run_migrations_online()

tasks.py

@@ -14,6 +14,7 @@
     saritasa_invocations.poetry,
     saritasa_invocations.mypy,
     saritasa_invocations.pytest,
+    saritasa_invocations.alembic,
 )
 
 # Configurations for run command
@@ -28,6 +29,9 @@
             docker=saritasa_invocations.DockerSettings(
                 main_containers=("postgres",),
             ),
+            alembic=saritasa_invocations.AlembicSettings(
+                migrations_folder="tests/alembic/versions",
+            ),
         ),
     },
 )

tests/alembic/README

@@ -0,0 +1 @@
+Generic single-database configuration with an async dbapi.

tests/alembic/env.py

@@ -0,0 +1,13 @@
+import saritasa_sqlalchemy_tools
+import tests.models  # noqa
+
+saritasa_sqlalchemy_tools.AlembicMigrations(
+    target_metadata=saritasa_sqlalchemy_tools.models.BaseModel.metadata,
+    db_driver="postgresql+asyncpg",
+    db_user="saritasa-sqlalchemy-tools-user",
+    db_password="manager",
+    db_host="postgres",
+    db_port=5432,
+    db_name="saritasa-sqlalchemy-tools-dev",
+    query={},  # type: ignore
+).run()

tests/alembic/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: str | None = ${repr(down_revision)}
+branch_labels: str | None = ${repr(branch_labels)}
+depends_on: str | None = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    """Apply migrations to database."""
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    """Roll back migrations from database."""
+    ${downgrades if downgrades else "pass"}