Vineeth/dev 516 hotfix

Dockerfile

@@ -34,6 +34,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV PATH="/app/.venv/bin:$PATH"
 
 COPY --chown=app:app src/ /app/src/
+COPY --chown=app:app migrations/ /app/migrations/
+COPY --chown=app:app alembic.ini /app/alembic.ini
 
 EXPOSE 8000
 

migrations/env.py

@@ -1,11 +1,17 @@
+import logging
 import os
 import sys
 from logging.config import fileConfig
 from pathlib import Path
 
 from alembic import context
 from dotenv import load_dotenv
-from sqlalchemy import engine_from_config, pool
+from sqlalchemy import engine_from_config, pool, text
+
+# Set up logging more verbosely
+logging.basicConfig()
+logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+logging.getLogger("alembic").setLevel(logging.DEBUG)
 
 # Import your models
 from src.db import Base
@@ -37,8 +43,18 @@
 # ... etc.
 
 
-def get_url():
-    return os.getenv("CONNECTION_URI")
+def include_name(name, type_, parent_names):
+    if type_ == "schema":
+        return name == target_metadata.schema
+    else:
+        return True
+
+
+def get_url() -> str:
+    url = os.getenv("CONNECTION_URI")
+    if url is None:
+        raise ValueError("CONNECTION_URI environment variable is not set")
+    return url
 
 
 def run_migrations_offline() -> None:
@@ -53,16 +69,23 @@ def run_migrations_offline() -> None:
     script output.
 
     """
-    # url = config.get_main_option("sqlalchemy.url")
     url = get_url()
+
+    print(target_metadata.schema)
+
     context.configure(
         url=url,
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        version_table_schema=target_metadata.schema,  # This sets schema for version table
+        include_schemas=True,
+        include_name=include_name,
     )
 
     with context.begin_transaction():
+        context.execute(f"create schema if not exists {target_metadata.schema};")
+        context.execute(f"SET search_path TO {target_metadata.schema}")
         context.run_migrations()
 
 
@@ -74,19 +97,42 @@ def run_migrations_online() -> None:
 
     """
     configuration = config.get_section(config.config_ini_section)
-    configuration["sqlalchemy.url"] = get_url()
+    if configuration is None:
+        configuration = {}
+
+    url = get_url()
+    configuration["sqlalchemy.url"] = url
+
+    print(f"Debug - Target metadata schema: {target_metadata.schema}")
 
     connectable = engine_from_config(
         configuration,
         prefix="sqlalchemy.",
+        echo=True,
         poolclass=pool.NullPool,
     )
 
     with connectable.connect() as connection:
-        context.configure(connection=connection, target_metadata=target_metadata)
+        # Set and verify search_path
+        connection.execute(text(f"SET search_path TO {target_metadata.schema}, public"))
+
+        # make use of non-supported SQLAlchemy attribute to ensure
+        # the dialect reflects tables in terms of the current tenant name
+        # connection.dialect.default_schema_name = target_metadata.schema
+
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table_schema=target_metadata.schema,
+            include_schemas=True,
+            include_name=include_name,
+            transaction_per_migration=True,
+            transactional_ddl=True,
+        )
 
         with context.begin_transaction():
             context.run_migrations()
+            connection.commit()
 
 
 if context.is_offline_mode():

migrations/versions/c3828084f472_add_indexes_for_messages_and_.py

@@ -6,12 +6,12 @@
 
 """
 
-from typing import Sequence, Union
+from collections.abc import Sequence
+from os import getenv
+from typing import Union
 
 import sqlalchemy as sa
-from sqlalchemy import text
 from alembic import op
-from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
 revision: str = "c3828084f472"
@@ -21,9 +21,14 @@
 
 
 def upgrade() -> None:
+    schema = getenv("DATABASE_SCHEMA", "public")
     # Add new indexes
-    op.create_index("idx_users_app_lookup", "users", ["app_id", "public_id"])
-    op.create_index("idx_sessions_user_lookup", "sessions", ["user_id", "public_id"])
+    op.create_index(
+        "idx_users_app_lookup", "users", ["app_id", "public_id"], schema=schema
+    )
+    op.create_index(
+        "idx_sessions_user_lookup", "sessions", ["user_id", "public_id"], schema=schema
+    )
 
     op.create_index(
         "idx_messages_session_lookup",
@@ -32,10 +37,9 @@ def upgrade() -> None:
         postgresql_include=[
             "public_id",
             "is_user",
-            "content",
-            "metadata",
             "created_at",
         ],
+        schema=schema,
     )
 
     op.create_index(
@@ -44,17 +48,18 @@ def upgrade() -> None:
         ["metamessage_type", sa.text("id DESC")],
         postgresql_include=[
             "public_id",
-            "content",
             "message_id",
             "created_at",
-            "metadata",
         ],
+        schema=schema,
     )
 
 
 def downgrade() -> None:
+    schema = getenv("DATABASE_SCHEMA", "public")
+
     # Remove new indexes
-    op.drop_index("idx_users_app_lookup", table_name="users")
-    op.drop_index("idx_sessions_user_lookup", table_name="sessions")
-    op.drop_index("idx_messages_session_lookup", table_name="messages")
-    op.drop_index("idx_metamessages_lookup", table_name="metamessages")
+    op.drop_index("idx_users_app_lookup", table_name="users", schema=schema)
+    op.drop_index("idx_sessions_user_lookup", table_name="sessions", schema=schema)
+    op.drop_index("idx_messages_session_lookup", table_name="messages", schema=schema)
+    op.drop_index("idx_metamessages_lookup", table_name="metamessages", schema=schema)

src/db.py

@@ -3,7 +3,7 @@
 from alembic import command
 from alembic.config import Config
 from dotenv import load_dotenv
-from sqlalchemy import MetaData
+from sqlalchemy import MetaData, create_engine, inspect
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.orm import declarative_base
 
@@ -45,5 +45,27 @@ def scaffold_db():
     """use a sync engine for scaffolding the database. ddl operations are unavailable
     with async engines
     """
+    # Create engine
+    engine = create_engine(
+        os.environ["CONNECTION_URI"],
+        pool_pre_ping=True,
+        echo=True,
+    )
+
+    # Create inspector to check if database exists
+    inspector = inspect(engine)
+
+    print(inspector.get_table_names(Base.metadata.schema))
+
+    # If no tables exist, create them with SQLAlchemy
+    if not inspector.get_table_names(Base.metadata.schema):
+        print("No tables found. Creating database schema...")
+        Base.metadata.create_all(bind=engine)
+
+    # Clean up
+    engine.dispose()
+
+    # Run Alembic migrations regardless
+    print("Running database migrations...")
     alembic_cfg = Config("alembic.ini")
     command.upgrade(alembic_cfg, "head")

src/security.py

@@ -1,12 +1,11 @@
 import os
 from typing import Annotated
 
-import httpx
-from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
+from fastapi import Depends, HTTPException
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 USE_AUTH_SERVICE = os.getenv("USE_AUTH_SERVICE", "False").lower() == "true"
-AUTH_SERVICE_URL = os.getenv("AUTH_SERVICE_URL", "http://localhost:8001")
+SECRET_KEY = os.getenv("SECRET_KEY", "test")
 
 security = HTTPBearer(
     auto_error=False,
@@ -18,15 +17,6 @@ async def auth(
 ):
     if not USE_AUTH_SERVICE:
         return True
-    print(credentials)
-    if not credentials or credentials.credentials != "test":
+    if not credentials or credentials.credentials != SECRET_KEY:
         raise HTTPException(status_code=401, detail="Invalid access token")
-    # payload = {"token": token}
-    # res = httpx.get(
-    #     f"{AUTH_SERVICE_URL}/validate",
-    #     params=payload,
-    # )
-    # data = res.json()
     return {"message": "OK"}
-
-    # return {"scheme": credentials.scheme, "token": credentials.credentials}