Moved alembic to inside mcpgateway and fixed wheel (#310)

* Moved alembic to inside mcpgateway and fixed wheel
Signed-off-by: Madhav Kandukuri <[email protected]>

* Ruff isort fixes
Signed-off-by: Madhav Kandukuri <[email protected]>

* Flake8 and black fixes
Signed-off-by: Madhav Kandukuri <[email protected]>

* Moved _pubsub to __init__
Signed-off-by: Madhav Kandukuri <[email protected]>

* Update readme for alembic current
Signed-off-by: Madhav Kandukuri <[email protected]>

Author: madhav165

alembic.ini renamed to mcpgateway/alembic.ini

@@ -70,7 +70,7 @@ alembic/
 ### Other Common Commands
 
 ```bash
-alembic current             # Show current DB revision
+alembic -c mcpgateway/alembic.ini current             # Show current DB revision
 alembic history --verbose   # Show all migrations and their order
 alembic downgrade -1        # Roll back one revision
 alembic downgrade <rev>     # Roll back to a specific revision hash
@@ -158,7 +158,7 @@ make db-history
 | **`Can't locate revision ...`**    | You deleted or renamed a revision file that the DB is pointing to. Either restore it or run `alembic stamp base` and recreate the revision.           |
 | **`script.py.mako` missing**       | This file is required. Run `alembic init alembic` in a temp folder and copy the missing template into your project.                                   |
 | **SQLite foreign key limitations** | SQLite doesn't allow dropping constraints. Use `create table → copy → drop` flow manually, or plan around it.                                         |
-| **DB not updating**                | Did you forget to run `alembic upgrade head`? Check with `alembic current`.                                                                           |
+| **DB not updating**                | Did you forget to run `alembic upgrade head`? Check with `alembic -c mcpgateway/alembic.ini current`.                                                                           |
 | **Wrong DB URL or config errors**  | Confirm `settings.database_url` is valid. Check `env.py` and your `.env`/config settings. Alembic ignores `alembic.ini` for URLs in your setup.       |
 | **Model changes not detected**     | Alembic only picks up declarative models in `Base.metadata`. Ensure all models are imported and not behind `if TYPE_CHECKING:` or other lazy imports. |
 

alembic/README.md renamed to mcpgateway/alembic/README.md

@@ -3,11 +3,12 @@
 from logging.config import fileConfig
 
 # Third-Party
+from alembic import context
 from sqlalchemy import engine_from_config, pool
 
 # First-Party
-from alembic import context
 from mcpgateway.config import settings
+from mcpgateway.db import Base
 
 # from mcpgateway.db import get_metadata
 # target_metadata = get_metadata()
@@ -28,7 +29,6 @@
 # add your model's MetaData object here
 # for 'autogenerate' support
 # from myapp import mymodel
-from mcpgateway.db import Base
 
 target_metadata = Base.metadata
 
@@ -81,9 +81,7 @@ def run_migrations_online() -> None:
     )
 
     with connectable.connect() as connection:
-        context.configure(
-            connection=connection, target_metadata=target_metadata
-        )
+        context.configure(connection=connection, target_metadata=target_metadata)
 
         with context.begin_transaction():
             context.run_migrations()

alembic/env.py renamed to mcpgateway/alembic/env.py

@@ -11,16 +11,16 @@
 import uuid
 
 # Third-Party
+from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.orm import Session
 
 # First-Party
-from alembic import op
 from mcpgateway.config import settings
 from mcpgateway.utils.create_slug import slugify
 
 # revision identifiers, used by Alembic.
-revision: str = 'b77ca9d2de7e'
+revision: str = "b77ca9d2de7e"
 down_revision: Union[str, Sequence[str], None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -70,7 +70,6 @@ def upgrade() -> None:
     op.add_column("server_resource_association", sa.Column("server_id_new", sa.String(36), nullable=True))
     op.add_column("server_prompt_association", sa.Column("server_id_new", sa.String(36), nullable=True))
 
-
     # ── STAGE 2: POPULATE THE NEW COLUMNS (DATA MIGRATION) ───────────────
     gateways = sess.execute(sa.select(sa.text("id, name")).select_from(sa.text("gateways"))).all()
     for gid, gname in gateways:
@@ -80,9 +79,7 @@ def upgrade() -> None:
             {"u": g_uuid, "s": slugify(gname), "i": gid},
         )
 
-    tools = sess.execute(
-        sa.select(sa.text("id, name, gateway_id")).select_from(sa.text("tools"))
-    ).all()
+    tools = sess.execute(sa.select(sa.text("id, name, gateway_id")).select_from(sa.text("tools"))).all()
     for tid, tname, g_old in tools:
         t_uuid = uuid.uuid4().hex
         tool_slug = slugify(tname)
@@ -102,8 +99,12 @@ def upgrade() -> None:
                 """
             ),
             {
-                "u": t_uuid, "on": tname, "ons": tool_slug,
-                "sep": settings.gateway_tool_name_separator, "g": g_old, "i": tid,
+                "u": t_uuid,
+                "on": tname,
+                "ons": tool_slug,
+                "sep": settings.gateway_tool_name_separator,
+                "g": g_old,
+                "i": tid,
             },
         )
 
@@ -123,7 +124,10 @@ def upgrade() -> None:
         sess.execute(sa.text("UPDATE prompts SET gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) WHERE id=:i"), {"g": g_old, "i": pid})
     sta = sess.execute(sa.select(sa.text("server_id, tool_id")).select_from(sa.text("server_tool_association"))).all()
     for s_old, t_old in sta:
-        sess.execute(sa.text("UPDATE server_tool_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s), tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE server_id=:s AND tool_id=:t"), {"s": s_old, "t": t_old})
+        sess.execute(
+            sa.text("UPDATE server_tool_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s), tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE server_id=:s AND tool_id=:t"),
+            {"s": s_old, "t": t_old},
+        )
     tool_metrics = sess.execute(sa.select(sa.text("id, tool_id")).select_from(sa.text("tool_metrics"))).all()
     for tmid, t_old in tool_metrics:
         sess.execute(sa.text("UPDATE tool_metrics SET tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE id=:i"), {"t": t_old, "i": tmid})
@@ -220,6 +224,7 @@ def upgrade() -> None:
     with op.batch_alter_table("server_prompt_association") as batch_op:
         batch_op.create_foreign_key("fk_server_prompt_association_server_id", "servers", ["server_id"], ["id"])
 
+
 # def upgrade() -> None:
 #     bind = op.get_bind()
 #     sess = Session(bind=bind)

alembic/script.py.mako renamed to mcpgateway/alembic/script.py.mako

@@ -10,15 +10,12 @@
 from typing import Sequence, Union
 
 # Third-Party
-import sqlalchemy as sa
-from sqlalchemy.orm import Session
-
-# First-Party
 from alembic import op
+import sqlalchemy as sa
 
 # revision identifiers, used by Alembic.
-revision: str = 'e4fc04d1a442'
-down_revision: Union[str, Sequence[str], None] = 'b77ca9d2de7e'
+revision: str = "e4fc04d1a442"
+down_revision: Union[str, Sequence[str], None] = "b77ca9d2de7e"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
@@ -32,14 +29,13 @@ def upgrade() -> None:
     that existing rows get a non-null default value.
     """
     bind = op.get_bind()
-    sess = Session(bind=bind)
     inspector = sa.inspect(bind)
 
     if not inspector.has_table("gateways"):
         print("Fresh database detected. Skipping migration.")
         return
 
-    op.add_column('tools', sa.Column('annotations', sa.JSON(), server_default=sa.text("'{}'"), nullable=False))
+    op.add_column("tools", sa.Column("annotations", sa.JSON(), server_default=sa.text("'{}'"), nullable=False))
 
 
 def downgrade() -> None:
@@ -50,11 +46,10 @@ def downgrade() -> None:
     'annotations' column from the 'tool' table.
     """
     bind = op.get_bind()
-    sess = Session(bind=bind)
     inspector = sa.inspect(bind)
 
     if not inspector.has_table("gateways"):
         print("Fresh database detected. Skipping migration.")
         return
 
-    op.drop_column('tools', 'annotations')
+    op.drop_column("tools", "annotations")

alembic/versions/b77ca9d2de7e_uuid_pk_and_slug_refactor.py renamed to mcpgateway/alembic/versions/b77ca9d2de7e_uuid_pk_and_slug_refactor.py

@@ -10,11 +10,9 @@
 from typing import Sequence, Union
 
 # Third-Party
-import sqlalchemy as sa
-
-# First-Party
 # Alembic / SQLAlchemy
 from alembic import op
+import sqlalchemy as sa
 
 # Revision identifiers.
 revision: str = "e75490e949b1"
@@ -28,20 +26,20 @@ def upgrade():
     Renames 'is_active' to 'enabled' and adds a new 'reachable' column (default True)
     in both 'tools' and 'gateways' tables.
     """
-    op.alter_column('tools', 'is_active', new_column_name='enabled')
-    op.add_column('tools', sa.Column('reachable', sa.Boolean(), nullable=False, server_default=sa.true()))
+    op.alter_column("tools", "is_active", new_column_name="enabled")
+    op.add_column("tools", sa.Column("reachable", sa.Boolean(), nullable=False, server_default=sa.true()))
 
-    op.alter_column('gateways', 'is_active', new_column_name='enabled')
-    op.add_column('gateways', sa.Column('reachable', sa.Boolean(), nullable=False, server_default=sa.true()))
+    op.alter_column("gateways", "is_active", new_column_name="enabled")
+    op.add_column("gateways", sa.Column("reachable", sa.Boolean(), nullable=False, server_default=sa.true()))
 
 
 def downgrade():
     """
     Reverts the changes by renaming 'enabled' back to 'is_active'
     and dropping the 'reachable' column in both 'tools' and 'gateways' tables.
     """
-    op.alter_column('tools', 'enabled', new_column_name='is_active')
-    op.drop_column('tools', 'reachable')
+    op.alter_column("tools", "enabled", new_column_name="is_active")
+    op.drop_column("tools", "reachable")
 
-    op.alter_column('gateways', 'enabled', new_column_name='is_active')
-    op.drop_column('gateways', 'reachable')
+    op.alter_column("gateways", "enabled", new_column_name="is_active")
+    op.drop_column("gateways", "reachable")

alembic/versions/e4fc04d1a442_add_annotations_to_tables.py renamed to mcpgateway/alembic/versions/e4fc04d1a442_add_annotations_to_tables.py

@@ -21,15 +21,15 @@
 
 # Standard
 import asyncio
+from importlib.resources import files
 import logging
-from pathlib import Path
 
 # Third-Party
+from alembic import command
+from alembic.config import Config
 from sqlalchemy import create_engine, inspect
 
 # First-Party
-from alembic import command
-from alembic.config import Config
 from mcpgateway.config import settings
 from mcpgateway.db import Base
 
@@ -47,9 +47,8 @@ async def main() -> None:
         None
     """
     engine = create_engine(settings.database_url)
-    project_root = Path(__file__).resolve().parents[1]
-    ini_path = project_root / "alembic.ini"
-    cfg = Config(ini_path)  # path in container
+    ini_path = files("mcpgateway").joinpath("alembic.ini")
+    cfg = Config(str(ini_path))  # path in container
     cfg.attributes["configure_logger"] = False
 
     command.ensure_version(cfg)

alembic/versions/e75490e949b1_add_improved_status_to_tables.py renamed to mcpgateway/alembic/versions/e75490e949b1_add_improved_status_to_tables.py

@@ -94,6 +94,7 @@ def __init__(
                 raise ValueError("Redis backend requires redis_url")
 
             self._redis = Redis.from_url(redis_url)
+            self._pubsub = self._redis.pubsub()
 
         elif self._backend == "database":
             if not SQLALCHEMY_AVAILABLE:
@@ -151,7 +152,6 @@ async def initialize(self) -> None:
             logger.info("Database cleanup task started")
 
         elif self._backend == "redis":
-            self._pubsub = self._redis.pubsub()
             await self._pubsub.subscribe("mcp_session_events")
 
         elif self._backend == "none":

mcpgateway/bootstrap_db.py

@@ -145,7 +145,13 @@
 wait_for_db_ready(max_tries=int(settings.db_max_retries), interval=int(settings.db_retry_interval_ms) / 1000, sync=True)  # Converting ms to s
 
 # Create database tables
-asyncio.run(bootstrap_db())
+try:
+    loop = asyncio.get_running_loop()
+except RuntimeError:
+    asyncio.run(bootstrap_db())
+else:
+    loop.create_task(bootstrap_db())
+
 
 # Initialize services
 tool_service = ToolService()