Enhancement: Way to hook into the database CLI

[DataDaoDe]

Summary

I would like to extend the database command to be able to do something like the following:

class DatabaseCLIPlugin(CLIPluginProtocol):
    def on_cli_init(self, cli: Group) -> None:
        db_group = get_alchemy_group()

        @db_group.command(name="create-database", help="Uses the config connection_string to create a new database if none exists")
        def create_database():
            ....

I can't seem to figure out a way to hook into the cli and extend the advanced alchemy command though. I can roll my db cli command and integrate that with litestar, but I'd rather not have two commands i.e. one called database and then my custom one called db.

Does anyone have any ideas?

Alternatively or additionally, is there any interest in taking up two new commands: create-db and drop-db which create and drop the database respectively. I tend to have these in a lot of my projects.

Thanks!

Basic Example

class CLIPlugin(CLIPluginProtocol):
    def on_cli_init(self, cli: Group) -> None:
        console = get_console()

        # for create/drop commands we cannot be connected to the db we want to create/drop
        app_db_url = urlparse(settings.DATABASE_URL)
        app_db_name = app_db_url.path.replace("/", "")
        pg_db_url_result = app_db_url._replace(path="/postgres")
        pg_db_url = urlunparse(pg_db_url_result)

        async_engine = create_async_engine(
            url=pg_db_url,
            execution_options={"isolation_level": "AUTOCOMMIT"},
        )

        @cli.command(name="db-create")
        def create_database(app: Litestar) -> None:
            async def _create_db():
                async with async_engine.begin() as conn:
                    try:
                        sql = "CREATE DATABASE {} ENCODING '{}'".format(
                            app_db_name, "utf-8"
                        )
                        await conn.execute(text(sql))
                    except DatabaseError as err:
                        console.print(err)

            asyncio.run(_create_db())

        @cli.command(name="db-drop")
        def drop_database(app: Litestar) -> None:
            async def _drop_db():
                async with async_engine.connect() as conn:
                    try:
                        conn = await conn.execution_options(
                            isolation_level="AUTOCOMMIT"
                        )

                        # We must first disconnect all users from the database we are dropping.
                        version = conn.dialect.server_version_info
                        pid_column = "pid" if (version >= (9, 2)) else "procpid"
                        drop_user_conn_sql = """SELECT pg_terminate_backend(pg_stat_activity.{pid_column})
                          FROM pg_stat_activity
                          WHERE pg_stat_activity.datname = '{database}'
                          AND {pid_column} <> pg_backend_pid();""".format(
                            pid_column=pid_column, database=app_db_name
                        )
                        await conn.execute(text(drop_user_conn_sql))

                        drop_sql = "DROP DATABASE {}".format(app_db_name)
                        await conn.execute(text(drop_sql))

                        console.print(f"dropped database '{app_db_name}'")
                    except DatabaseError as err:
                        console.print(f"Error dropping database: {err}")
                    except Exception as ex:
                        console.print(f"Unexpected error: {ex}")
                await async_engine.dispose()

            asyncio.run(_drop_db())

Drawbacks and Impact

No response

Unresolved questions

No response

[cofin]

Hey, @DataDaoDe, these are useful commands. Since they are Postgres specific though, I think we should leave them as user added. However, there's a very easy way to register the commands to the existing group. It works exactly like any other click group. Here's how I add a load-fixtures command to the database group in my own code:

class CommandLinePlugin(InitPluginProtocol, CLIPluginProtocol):
    """Adds custom commands to the CLI."""

    def on_cli_init(self, cli: click.Group) -> None:  # noqa: PLR0915
        from advanced_alchemy.extensions.litestar.cli import database_group
 
        @database_group.command("load-fixtures")
        def load_database_fixtures() -> None:
            import anyio
            from rich import get_console

            console = get_console()

            async def _load_database_fixtures() -> None:
                """Import/Synchronize Database Fixtures."""
                from pathlib import Path

                from advanced_alchemy.utils.fixtures import open_fixture_async
                from rich import get_console

                from app import config, services
                from config import get_settings

                console, settings = get_console(), get_settings()
                fixture_path = Path(settings.db.FIXTURE_PATH)
                async with services.UserService.new(config=config.alchemy) as service:
                    for fixture in await open_fixture_async(fixture_path, "users"):
                        _obj, _created = await service.get_or_upsert(match_fields=["email"], upsert=False, **fixture)
                    await service.repository.session.commit()
                    console.print("User fixtures loaded.")

            console.rule("Loading database fixtures.")
            anyio.run(_load_database_fixtures)

[cofin]

I should also add that if you want override an existing command, just extend the group with a new command that's named the same as the old one.

Lastly, if there's interest in making the command group name configurable, let me know. I find myself wanting to use app db instead of app database sometimes, and I've thought about just changing the name to that.

[DataDaoDe]

@cofin awesome! That worked like a charm. A few thoughts I had while adding:

From the time of my commenting and before you commented I also added essentially your load-fixtures command, also leveraging the advanced-alchemy functions. I just called it seed. In general, these are three "problems" that basically always need to be solved: "creating a db", "dropping a db" and "seeding a db". My typical approach in projects is to have utility functions to do all of these and then to give them a CLI wrapper. (Nicely all the other major db use cases are already solved with the other cli commands in advanced alchemy).

For what its worth I prefer db over database, less typing for no semantic loss :)

What are your thoughts about addressing these three problem areas with advanced-alchemy? I'd be glad to help out / or take on the implementation if its something people would be interested in. As far as the supported database, I'd be glad to try to write a more generic solution which I'm thinking might not be too difficult for most of the other supported databases. Of course, this might be better addressed by simply building another library...

[Kumzy]

Hello @DataDaoDe,

I'll close this issue as it has been documented in the documentation Please do not hesitate to reopen if you find the documentation unclear or insufficient for your needs.

Thank you.