Add database support with migrations

Author: arthurhenrique

cookiecutter.json

@@ -5,6 +5,7 @@
     "machine_learn_model_path": "./ml/model/",
     "machine_learn_model_name": "model.pkl",
     "input_example_path": "./ml/model/examples/example.json",
+    "database_url": "sqlite:///./app.db",
     "full_name": "Your name",
     "email": "[email protected]",
     "release_date": "{% now 'local' %}",

sample/pregnancy-model/.env.example

@@ -1,4 +1,6 @@
 SECRET_KEY=secret
 DEBUG=True
 MODEL_PATH=/Users/arthur.dasilva/repos/arthurhenrique/n
-MODEL_NAME=pregnancy_model_local.joblib
+MODEL_NAME=pregnancy_model_local.joblib
+MEMOIZATION_FLAG=False
+DATABASE_URL=sqlite:///./app.db

sample/pregnancy-model/README.md

@@ -14,6 +14,28 @@ MODEL_PATH=./ml/model/
 MODEL_NAME=model.pkl
 ```
 
+### Database Configuration
+
+Set your database url in `.env` using `DATABASE_URL`. The default uses SQLite:
+
+```sh
+DATABASE_URL=sqlite:///./app.db
+```
+
+### Migrations
+
+Create a new migration with:
+
+```sh
+alembic revision --autogenerate -m "message"
+```
+
+Apply migrations with:
+
+```sh
+alembic upgrade head
+```
+
 ### Update `/predict`
 
 To update your machine learning model, add your `load` and `method` [change here](app/api/routes/predictor.py#L19) at `predictor.py`

sample/pregnancy-model/alembic.ini

@@ -0,0 +1,36 @@
+[alembic]
+script_location = migrations
+sqlalchemy.url = sqlite:///./app.db
+
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname = 
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s

sample/pregnancy-model/app/core/config.py

@@ -5,7 +5,7 @@
 from starlette.config import Config
 from starlette.datastructures import Secret
 
-from core.logging import InterceptHandler
+from .logging import InterceptHandler
 
 config = Config(".env")
 
@@ -29,3 +29,4 @@
 MODEL_PATH = config("MODEL_PATH", default="/Users/arthur.dasilva/repos/arthurhenrique/cookiecutter-fastapi/sample/pregnancy-model")
 MODEL_NAME = config("MODEL_NAME", default="pregnancy_model_local.joblib")
 INPUT_EXAMPLE = config("INPUT_EXAMPLE", default="./ml/model/examples/example.json")
+DATABASE_URL = config("DATABASE_URL", default="sqlite:///./app.db")

sample/pregnancy-model/app/core/database.py

@@ -0,0 +1,20 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+from .config import DATABASE_URL
+
+connect_args = {}
+if DATABASE_URL.startswith("sqlite"):
+    connect_args["check_same_thread"] = False
+
+engine = create_engine(DATABASE_URL, connect_args=connect_args)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+Base = declarative_base()
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()

sample/pregnancy-model/app/models/item.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Column, Integer, String
+
+from app.core.database import Base
+
+
+class Item(Base):
+    __tablename__ = "items"
+
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String, index=True)

sample/pregnancy-model/migrations/env.py

@@ -0,0 +1,51 @@
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+
+from alembic import context
+
+from core.database import Base
+from models import item
+from core.config import DATABASE_URL
+
+config = context.config
+
+if config.get_main_option("sqlalchemy.url") is None:
+    config.set_main_option("sqlalchemy.url", DATABASE_URL)
+
+fileConfig(config.config_file_name)
+
+target_metadata = Base.metadata
+
+
+def run_migrations_offline() -> None:
+    context.configure(
+        url=config.get_main_option("sqlalchemy.url"),
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(connection=connection, target_metadata=target_metadata)
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

sample/pregnancy-model/migrations/versions/0001_create_items_table.py

@@ -0,0 +1,25 @@
+"""create items table
+
+Revision ID: 0001
+Revises: 
+Create Date: 2020-01-01 00:00:00.000000
+"""
+from alembic import op
+import sqlalchemy as sa
+
+revision = "0001"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "items",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("name", sa.String(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("items")

sample/pregnancy-model/pyproject.toml

@@ -14,6 +14,8 @@ dependencies = [
     "loguru>=0.7.0",
     "joblib>=1.2.0",
     "scikit-learn>=1.1.3",
+    "SQLAlchemy>=2.0.0",
+    "alembic>=1.12.0",
     "pandas>=2.2.3",
 ]