feat: move full litestar plugin (#17)

Author: cofin

.gitignore

@@ -50,6 +50,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+test.sqlite
 
 # Translations
 *.mo

.pre-commit-config.yaml

@@ -17,12 +17,12 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/provinzkraut/unasyncd
-    rev: "v0.6.0"
+    rev: "v0.6.1"
     hooks:
       - id: unasyncd
         additional_dependencies: ["ruff"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.0.289"
+    rev: "v0.0.290"
     hooks:
       - id: ruff
         args: ["--fix"]

README.md

@@ -31,14 +31,26 @@ offering features such as:
 - Tested support for multiple database backends including:
 
   - SQLite via [aiosqlite](https://aiosqlite.omnilib.dev/en/stable/) or [sqlite](https://docs.python.org/3/library/sqlite3.html)
-  - Postgres via [asyncpg](https://magicstack.github.io/asyncpg/current/)\_ or [psycopg3 (async or sync)](https://www.psycopg.org/psycopg3/)
+  - Postgres via [asyncpg](https://magicstack.github.io/asyncpg/current/) or [psycopg3 (async or sync)](https://www.psycopg.org/psycopg3/)
   - MySQL via [asyncmy](https://github.com/long2ice/asyncmy)
   - Oracle via [oracledb](https://oracle.github.io/python-oracledb/)
   - Google Spanner via [spanner-sqlalchemy](https://github.com/googleapis/python-spanner-sqlalchemy/)
   - DuckDB via [duckdb_engine](https://github.com/Mause/duckdb_engine>)
 
 ## Usage
 
+### Litestar
+
+> [!NOTE]\
+> This section has not been completed (yet!)
+
+### Starlette/FastAPI
+
+> [!NOTE]\
+> This section has not been completed (yet!)
+
+### Sanic
+
 > [!NOTE]\
 > This section has not been completed (yet!)
 

advanced_alchemy/alembic/templates/asyncio/env.py

@@ -119,8 +119,7 @@ async def run_migrations_online() -> None:
     connectable = cast(
         "AsyncEngine",
         config.engine
-        if config.engine
-        else async_engine_from_config(
+        or async_engine_from_config(
             configuration,
             prefix="sqlalchemy.",
             poolclass=pool.NullPool,

advanced_alchemy/base.py

@@ -1,6 +1,8 @@
 """Application ORM configuration."""
+
 from __future__ import annotations
 
+import contextlib
 import re
 from datetime import date, datetime, timezone
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeVar, runtime_checkable
@@ -18,7 +20,7 @@
     registry,
 )
 
-from .types import GUID, BigIntIdentity, DateTimeUTC, JsonB
+from advanced_alchemy.types import GUID, BigIntIdentity, DateTimeUTC, JsonB
 
 if TYPE_CHECKING:
     from sqlalchemy.sql import FromClause
@@ -160,13 +162,10 @@ def create_registry() -> registry:
         date: Date,
         dict: JsonB,
     }
-    try:
+    with contextlib.suppress(ImportError):
         from pydantic import AnyHttpUrl, AnyUrl, EmailStr
 
         type_annotation_map.update({EmailStr: String, AnyUrl: String, AnyHttpUrl: String})
-    except ImportError:
-        pass
-
     return registry(metadata=meta, type_annotation_map=type_annotation_map)
 
 

advanced_alchemy/extensions/litestar/__init__.py

@@ -1,3 +0,0 @@
-from .plugin import AdvancedAlchemyPlugin
-
-__all__ = ["AdvancedAlchemyPlugin"]

advanced_alchemy/extensions/litestar/alembic.py

@@ -0,0 +1,29 @@
+from contextlib import suppress
+
+from litestar import Litestar
+
+from advanced_alchemy.alembic.commands import AlembicCommands as _AlembicCommands
+from advanced_alchemy.exceptions import ImproperConfigurationError
+from advanced_alchemy.extensions.litestar.plugins import SQLAlchemyInitPlugin
+
+
+def get_database_migration_plugin(app: Litestar) -> SQLAlchemyInitPlugin:
+    """Retrieve a database migration plugin from the Litestar application's plugins.
+
+    This function attempts to find and return either the SQLAlchemyPlugin or SQLAlchemyInitPlugin.
+    If neither plugin is found, it raises an ImproperlyConfiguredException.
+    """
+
+    with suppress(KeyError):
+        return app.plugins.get(SQLAlchemyInitPlugin)
+    msg = "Failed to initialize database migrations. The required plugin (SQLAlchemyPlugin or SQLAlchemyInitPlugin) is missing."
+    raise ImproperConfigurationError(
+        msg,
+    )
+
+
+class AlembicCommands(_AlembicCommands):
+    def __init__(self, app: Litestar) -> None:
+        self._app = app
+        self.plugin_config = get_database_migration_plugin(self._app)._config  # noqa: SLF001
+        self.config = self._get_alembic_command_config()

advanced_alchemy/extensions/litestar/cli.py

@@ -2,10 +2,11 @@
 
 from typing import TYPE_CHECKING, cast
 
+from click import argument, group, option
 from litestar.cli._utils import LitestarGroup, console
 from rich.prompt import Confirm, Prompt
 
-from advanced_alchemy.extensions.litestar.plugin import AlembicCommands, _get_advanced_alchemy_plugin
+from advanced_alchemy.extensions.litestar.alembic import AlembicCommands, get_database_migration_plugin
 
 if TYPE_CHECKING:
     from litestar import Litestar
@@ -14,9 +15,6 @@
     from alembic.operations.ops import MigrationScript, UpgradeOps
 
 
-from click import argument, group, option
-
-
 @group(cls=LitestarGroup, name="database")
 def database_group() -> None:
     """Manage SQLAlchemy database components."""
@@ -130,16 +128,16 @@ def upgrade_database(app: Litestar, revision: str, sql: bool, tag: str | None, n
 def init_alembic(app: Litestar, directory: str | None, multidb: bool, package: bool, no_prompt: bool) -> None:
     """Upgrade the database to the latest revision."""
     console.rule("[yellow]Initializing database migrations.", align="left")
-    plugin = _get_advanced_alchemy_plugin(app)
+    plugin = get_database_migration_plugin(app)
     if directory is None:
-        directory = plugin._config.script_location  # noqa: SLF001
+        directory = plugin._alembic_config.script_location  # noqa: SLF001
     input_confirmed = (
         True
         if no_prompt
         else Confirm.ask("[bold]Are you sure you you want initialize the project in `{directory}`?[/]")
     )
     if input_confirmed:
-        alembic_commands = AlembicCommands(app=app)
+        alembic_commands = AlembicCommands(app)
         alembic_commands.init(directory=directory, multidb=multidb, package=package)
 
 
@@ -207,6 +205,7 @@ def process_revision_directives(
     console.rule("[yellow]Starting database upgrade process[/]", align="left")
     if message is None:
         message = "autogenerated" if no_prompt else Prompt.ask("Please enter a message describing this revision")
+
     alembic_commands = AlembicCommands(app=app)
     alembic_commands.revision(
         message=message,