Merge pull request #74 from Tanzania-AI-Community/feature/db-seed-procedure

New Database Seed Script

Author: jurmy24

Makefile

@@ -14,6 +14,6 @@ restart: stop run
 
 generate-local-data:
 	@echo 'Generating local data ...'
-	@docker-compose -f docker/dev/docker-compose.yml --env-file .env run --rm app bash -c "PYTHONPATH=/app uv run python scripts/database/seed.py --sample-data --vector-data chunks_BAAI.json"
+	@docker-compose -f docker/dev/docker-compose.yml --env-file .env run --rm app bash -c "PYTHONPATH=/app uv run python scripts/database/seed.py --create --sample-data --vector-data chunks_BAAI.json"
 
 setup-env: build generate-local-data run

app/config.py

@@ -41,6 +41,11 @@ class Settings(BaseSettings):
     # Database settings
     database_url: SecretStr
 
+    @property
+    def sync_database_url(self) -> str:
+        """Get synchronous database URL for migrations"""
+        return self.database_url.get_secret_value().replace("+asyncpg", "")
+
     # Business environment
     business_env: bool = False  # Default if not found in .env
 

docs/en/GETTING_STARTED.md

@@ -137,10 +137,6 @@ Next up, let's build all Docker images and local data, needed for further steps
 make setup-env
 ```
 
-> [!Note]
->
-> You can try out [TablePlus](https://tableplus.com/) to visualize your databases.
-
 ## 🖥️ Set up the FastAPI application
 
 Run the following command to run the project.

docs/en/MIGRATIONS.md

@@ -0,0 +1,13 @@
+# Migrations
+
+You might run into some database troubles that require you to do database migrations. In the folder `migrations/versions/` you find the list of past database migrations. We're using [Alembic](https://alembic.sqlalchemy.org/en/latest/). They're docs aren't great so here's a beginner [article](https://medium.com/@kasperjuunge/how-to-get-started-with-alembic-and-sqlmodel-288700002543) on it.
+
+By default, our Docker images use the alembic versioning system to initialize the database. If you wan't to rebuild the database to your needs, you can run new migrations and rebuild the Docker containers.
+
+If you're not using Docker to run Twiga, then you can initialize the database and inject seed data with the command:
+
+```
+uv run python -m scripts.database.seed --create --sample-data --vector-data chunks_BAAI.json
+```
+
+This will remove all tables in the database if they exist, create new ones, install pgvector and inject sample data and vector data so that the database is ready to accept new users.

migrations/env.py

@@ -15,7 +15,7 @@
 config = context.config
 
 # Load PostgreSQL migrations URL from environment variables
-db_url = settings.database_url.get_secret_value()
+db_url = settings.sync_database_url
 print(f"Connecting to: {db_url}")
 config.set_main_option("sqlalchemy.url", db_url)
 

scripts/database/seed.py

@@ -1,45 +1,70 @@
 from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
-import asyncio
 import json
 from pathlib import Path
-from sqlalchemy import text
+from sqlmodel import text
 from sqlmodel import SQLModel, select
 import logging
 from typing import List, Dict, Any
-import argparse
-
 import yaml
+import argparse
+import asyncio
+import sys
 
 # Import all your models
 import app.database.models as models
 from app.database.enums import ChunkType
 from app.database.utils import get_database_url
 
+
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-async def init_db(drop_all: bool = False):
+async def reset_db():
+    """Reset database using async operations"""
     try:
-        """Initialize the database with tables."""
         engine = create_async_engine(get_database_url())
-
         async with engine.begin() as conn:
-            if drop_all:
-                logger.info("Dropping all existing tables...")
-                await conn.run_sync(SQLModel.metadata.drop_all)
-
-            logger.info("Creating tables and pgvector extension...")
+            logger.info("Dropping all existing tables...")
+            await conn.run_sync(SQLModel.metadata.drop_all)
+
+            # Drop alembic_version table
+            logger.info("Dropping alembic_version table...")
+            await conn.execute(text("DROP TABLE IF EXISTS alembic_version CASCADE"))
+
+            logger.info("Dropping existing enum types...")
+            enum_query = """
+                SELECT t.typname FROM pg_type t
+                JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+                WHERE t.typtype = 'e' AND n.nspname = 'public'
+            """
+            result = await conn.execute(text(enum_query))
+            enum_types = result.scalars().all()
+
+            for enum_name in enum_types:
+                await conn.execute(text(f"DROP TYPE IF EXISTS {enum_name} CASCADE"))
+
+            logger.info("Creating pgvector extension...")
             await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
-            await conn.run_sync(SQLModel.metadata.create_all)
     except Exception as e:
         logger.error(f"Error initializing database: {str(e)}")
         raise
     finally:
         await engine.dispose()
 
 
+def run_migrations():
+    """Run alembic migrations"""
+    from alembic.config import Config
+    from alembic import command
+
+    logger.info("Running migrations...")
+    alembic_cfg = Config("alembic.ini")
+    command.upgrade(alembic_cfg, "head")
+    logger.info("Migrations complete")
+
+
 async def inject_sample_data():
     try:
         """Load sample data into the database."""
@@ -194,39 +219,45 @@ async def inject_vector_data(file: str):
         await engine.dispose()
 
 
-if __name__ == "__main__":
-    try:
-        # Parse command line arguments
-        parser = argparse.ArgumentParser(
-            description="Initialize the database for Twiga development."
-        )
-        parser.add_argument(
-            "--drop",
-            action="store_true",
-            help="Drop the existing tables before creation",
-        )
-        parser.add_argument(
-            "--sample-data", action="store_true", help="Add sample data to the database"
-        )
-        parser.add_argument(
-            "--vector-data",
-            type=str,
-            metavar="FILENAME",
-            help="Populate the vector database using the specified chunks file (chunks_OPENAI.json or chunks_BAAI.json)",
-        )
+async def main():
+    """Initialize and setup the Twiga development database."""
+    # Set up argument parser
+    parser = argparse.ArgumentParser(
+        description="Initialize and setup the Twiga development database."
+    )
+    parser.add_argument(
+        "--create", action="store_true", help="Reset and run alembic migrations"
+    )
+    parser.add_argument("--sample-data", action="store_true", help="Add sample data")
+    parser.add_argument(
+        "--vector-data",
+        type=str,
+        help="Vector database chunks file (chunks_OPENAI.json or chunks_BAAI.json)",
+    )
+
+    # Parse arguments
+    args = parser.parse_args()
 
-        args = parser.parse_args()
+    try:
 
-        asyncio.run(init_db(drop_all=args.drop))
+        if args.create:
+            logger.info("Starting database setup...")
+            await reset_db()
+            run_migrations()
 
         if args.sample_data:
-            logger.info("Starting data injection...")
-            asyncio.run(inject_sample_data())
+            logger.info("Starting sample data injection...")
+            await inject_sample_data()
 
         if args.vector_data:
             logger.info("Starting vector data injection...")
-            asyncio.run(inject_vector_data(args.vector_data))
+            await inject_vector_data(args.vector_data)
 
-        logger.info("Complete.")
+        logger.info("Database setup complete")
     except Exception as e:
-        logger.error(f"Setup failed: {e}")
+        print(f"Setup failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())