Database Tables

Agno templates come pre-configured with SqlAlchemy and Alembic to manage databases. You can use these tables to store data for your Agents, Teams and Workflows. The general way to add a table is:

1. Add table definition to the db/tables directory.
2. Import the table class in the db/tables/__init__.py file.
3. Create a database migration.
4. Run database migration.

Table Definition

Let’s create a UsersTable, copy the following code to db/tables/user.py

db/tables/user.py

from datetime import datetime
from typing import Optional

from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.sql.expression import text
from sqlalchemy.types import BigInteger, DateTime, String

from db.tables.base import Base


class UsersTable(Base):
    """Table for storing user data."""

    __tablename__ = "dim_users"

    id_user: Mapped[int] = mapped_column(
        BigInteger, primary_key=True, autoincrement=True, nullable=False, index=True
    )
    email: Mapped[str] = mapped_column(String)
    is_active: Mapped[bool] = mapped_column(default=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=text("now()")
    )
    updated_at: Mapped[Optional[datetime]] = mapped_column(
        DateTime(timezone=True), onupdate=text("now()")
    )

Update the db/tables/__init__.py file:

db/tables/__init__.py

from db.tables.base import Base
from db.tables.user import UsersTable

Create a database revision

Run the alembic command to create a database migration in the dev container:

docker exec -it {infra_name}-api alembic -c db/alembic.ini revision --autogenerate -m "Initialize DB"

Migrate dev database

Run the alembic command to migrate the dev database:

docker exec -it {infra_name}-api alembic -c db/alembic.ini upgrade head

Optional: Add test user

Now lets’s add a test user. Copy the following code to db/tables/test_add_user.py

db/tables/test_add_user.py

from typing import Optional
from sqlalchemy.orm import Session

from db.session import SessionLocal
from db.tables.user import UsersTable
from utils.log import logger


def create_user(db_session: Session, email: str) -> UsersTable:
    """Create a new user."""
    new_user = UsersTable(email=email)
    db_session.add(new_user)
    return new_user


def get_user(db_session: Session, email: str) -> Optional[UsersTable]:
    """Get a user by email."""
    return db_session.query(UsersTable).filter(UsersTable.email == email).first()


if __name__ == "__main__":
    test_user_email = "[e-mail]"
    with SessionLocal() as sess, sess.begin():
        logger.info(f"Creating user: {test_user_email}")
        create_user(db_session=sess, email=test_user_email)
        logger.info(f"Getting user: {test_user_email}")
        user = get_user(db_session=sess, email=test_user_email)
        if user:
            logger.info(f"User created: {user.id_user}")
        else:
            logger.info(f"User not found: {test_user_email}")

Run the script to add a test adding a user:

docker exec -it {infra_name}-api python db/tables/test_add_user.py

Migrate production database

We recommended migrating the production database by setting the environment variable MIGRATE_DB = True and restarting the production service. This runs alembic -c db/alembic.ini upgrade head from the entrypoint script at container startup.

Update infra/prd_resources.py

infra/prd_resources.py

container_env = {
    ...
    # Migrate database on startup using alembic
    "MIGRATE_DB": True,
}

Update the ECS Task Definition

Because we updated the Environment Variables, we need to update the Task Definition:

Full Options

ag infra patch --env prd --infra aws --name td

ag infra patch prd:aws::td

Update the ECS Service

After updating the task definition, redeploy the production application:

Full Options

ag infra patch --env prd --infra aws --name service

ag infra patch prd:aws::service

Manually migrate prodution database

Another approach is to SSH into the production container to run the migration manually. Your ECS tasks are already enabled with SSH access. Run the alembic command to migrate the production database:

ECS_CLUSTER={infra_name}-prd
TASK_ARN=$(aws ecs list-tasks --cluster $ECS_CLUSTER --query "taskArns[0]" --output text)
CONTAINER_NAME={infra_name}-prd-api

aws ecs execute-command --cluster $ECS_CLUSTER \
    --task $TASK_ARN \
    --container $CONTAINER_NAME \
    --interactive \
    --command "alembic -c db/alembic.ini upgrade head"

---

How the migrations directory was created

These commands have been run and are described for completeness

The migrations directory was created using:

docker exec -it {infra_name}-api sh -c "cd db && alembic init migrations"

• After running the above command, the db/migrations directory should be created.
• Update alembic.ini
  • set script_location = db/migrations
  • uncomment black hook in [post_write_hooks]
• Update db/migrations/env.py file following this link
• Add the following function to configure to only include tables in the target_metadata

db/migrations/env.py

# -*- Only include tables that are in the target_metadata
def include_name(name, type_, parent_names):
    if type_ == "table":
        return name in target_metadata.tables
    else:
        return True
...