diff --git a/Dockerfile b/Dockerfile index 4a7e473..52ac0e0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -47,6 +47,7 @@ WORKDIR /app # Copy application files and migrations COPY --chown=agentkit:agentkit alembic.ini ./ COPY --chown=agentkit:agentkit migrations ./migrations +COPY --chown=agentkit:agentkit scripts ./scripts COPY --chown=agentkit:agentkit agentkit ./agentkit # Switch to non-root user diff --git a/docker-compose.yml b/docker-compose.yml index b153448..701ba51 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -43,7 +43,7 @@ services: dockerfile: Dockerfile container_name: agentkit-migrations restart: "no" - command: ["alembic", "upgrade", "head"] + command: ["sh", "-c", "alembic upgrade head && python scripts/seed.py"] depends_on: postgres: condition: service_healthy diff --git a/examples/demo.py b/examples/demo.py index d550da6..ec57942 100644 --- a/examples/demo.py +++ b/examples/demo.py @@ -27,10 +27,8 @@ import asyncio import logging from dataclasses import dataclass -from pathlib import Path from redis.asyncio import Redis -from sqlalchemy import text from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, @@ -71,16 +69,9 @@ class DemoResources: async def seed_demo_database(engine: AsyncEngine) -> None: """Apply scripts/seed_demo_db.sql, creating products and orders tables.""" - seed_path = Path(__file__).parents[1] / "scripts" / "seed_demo_db.sql" - if not seed_path.exists(): - raise FileNotFoundError(f"Seed script not found at {seed_path}") + from scripts.seed import seed_database - statements = [s.strip() for s in seed_path.read_text(encoding="utf-8").split(";") if s.strip()] - async with engine.begin() as conn: - for stmt in statements: - await conn.execute(text(stmt)) - - logger.info("Seeded products and orders into the demo database.") + await seed_database(engine) async def build_persistence(db_url: str, offline: bool) -> DemoResources: diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..de1f3f4 --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""AgentKit utility and database seed scripts.""" diff --git a/scripts/seed.py b/scripts/seed.py new file mode 100644 index 0000000..bf697b8 --- /dev/null +++ b/scripts/seed.py @@ -0,0 +1,63 @@ +"""Database seeding script for AgentKit demo tables.""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +from pathlib import Path + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine + +from agentkit.config import get_settings + +logger = logging.getLogger("agentkit.seed") + + +async def seed_database(engine: AsyncEngine, script_path: Path | None = None) -> None: + """Execute demo SQL seed script against the provided database engine. + + Args: + engine: Connected AsyncEngine instance. + script_path: Optional custom path to SQL seed file (defaults to scripts/seed_demo_db.sql). + """ + if script_path is None: + script_path = Path(__file__).parent / "seed_demo_db.sql" + + if not script_path.exists(): + raise FileNotFoundError(f"Seed script not found at {script_path}") + + sql_content = script_path.read_text(encoding="utf-8") + statements = [stmt.strip() for stmt in sql_content.split(";") if stmt.strip()] + + logger.info("Applying %d SQL statements from %s", len(statements), script_path.name) + async with engine.begin() as conn: + for stmt in statements: + await conn.execute(text(stmt)) + + logger.info("Successfully seeded demo database tables (products, orders).") + + +async def main() -> None: + """CLI entrypoint for standalone database seeding.""" + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + parser = argparse.ArgumentParser(description="Seed demo products and orders database tables.") + parser.add_argument( + "--db-url", + type=str, + default=None, + help="Target database URL (defaults to DATABASE_URL in environment/settings).", + ) + args = parser.parse_args() + + db_url = args.db_url or get_settings().DATABASE_URL + engine = create_async_engine(db_url, pool_pre_ping=True) + try: + await seed_database(engine) + finally: + await engine.dispose() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/scripts/seed_demo_db.sql b/scripts/seed_demo_db.sql index 0d07e88..2388e33 100644 --- a/scripts/seed_demo_db.sql +++ b/scripts/seed_demo_db.sql @@ -27,7 +27,8 @@ INSERT INTO products (id, name, category, price, stock) VALUES (4, 'Ergonomic Desk Chair', 'Furniture', 350.00, 15), (5, 'Standing Desk Converter', 'Furniture', 280.00, 20), (6, 'Designing Data-Intensive Applications', 'Books', 45.00, 100), - (7, 'Site Reliability Engineering', 'Books', 40.00, 120); + (7, 'Site Reliability Engineering', 'Books', 40.00, 120) +ON CONFLICT (id) DO NOTHING; -- Seed Orders (Electronics total in 2024 = $2,400 + $500 + $300 = $3,200.00) INSERT INTO orders (id, product_id, quantity, total_price, order_date, customer_name) VALUES @@ -35,4 +36,5 @@ INSERT INTO orders (id, product_id, quantity, total_price, order_date, customer_ (2, 2, 2, 500.00, '2024-04-10', 'Bob Smith'), (3, 3, 2, 300.00, '2024-05-20', 'Charlie Brown'), (4, 4, 1, 350.00, '2024-02-11', 'Diana Prince'), - (5, 6, 2, 90.00, '2024-01-25', 'Evan Wright'); + (5, 6, 2, 90.00, '2024-01-25', 'Evan Wright') +ON CONFLICT (id) DO NOTHING; diff --git a/tests/integration/test_seed.py b/tests/integration/test_seed.py new file mode 100644 index 0000000..4a958e0 --- /dev/null +++ b/tests/integration/test_seed.py @@ -0,0 +1,37 @@ +"""Integration tests for demo database seeding script and idempotency.""" + +import pytest +from sqlalchemy import text +from sqlalchemy.ext.asyncio import create_async_engine + +from scripts.seed import seed_database + + +@pytest.mark.asyncio +async def test_seed_database_idempotent() -> None: + """Verify seed_database creates tables and rows, and second run does not fail.""" + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + + # First run + await seed_database(engine) + + async with engine.connect() as conn: + res_prod = await conn.execute(text("SELECT COUNT(*) FROM products")) + count_prod = res_prod.scalar() + assert count_prod == 7 + + res_orders = await conn.execute(text("SELECT COUNT(*) FROM orders")) + count_orders = res_orders.scalar() + assert count_orders == 5 + + # Second run should succeed cleanly with ON CONFLICT DO NOTHING + await seed_database(engine) + + async with engine.connect() as conn: + res_prod2 = await conn.execute(text("SELECT COUNT(*) FROM products")) + assert res_prod2.scalar() == 7 + + res_orders2 = await conn.execute(text("SELECT COUNT(*) FROM orders")) + assert res_orders2.scalar() == 5 + + await engine.dispose()