Backup local DB to S3

Author: nolan1999

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:system"
+}

api/core/auth.py

@@ -0,0 +1,33 @@
+from typing import Any
+
+from fastapi import Header, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials
+from fastapi_cloudauth.cognito import CognitoClaims, CognitoCurrentUser  # type: ignore
+from pydantic import Field
+
+
+# https://github.com/iwpnd/fastapi-key-auth/blob/main/fastapi_key_auth/dependency/authorizer.py
+class APIKeyDependency:
+    def __init__(self, key: str | None):
+        self.key = key
+
+    def __call__(self, x_api_key: str | None = Header(...)) -> str | None:
+        if x_api_key != self.key:
+            raise HTTPException(status_code=401, detail="unauthorized")
+        return x_api_key
+
+
+class GroupClaims(CognitoClaims):  # type: ignore
+    cognito_groups: list[str] | None = Field(alias="cognito:groups")
+
+
+class UserGroupCognitoCurrentUser(CognitoCurrentUser):  # type: ignore
+    user_info = GroupClaims
+
+    async def call(self, http_auth: HTTPAuthorizationCredentials) -> Any:
+        user_info = await super().call(http_auth)
+        if "users" not in (getattr(user_info, "cognito_groups") or []):
+            raise HTTPException(
+                status_code=403, detail="Not a member of the 'users' group"
+            )
+        return user_info

api/core/database.py

@@ -0,0 +1,121 @@
+import gzip
+import os
+import sqlite3
+import tempfile
+import time
+from typing import Any, cast
+
+from botocore.exceptions import ClientError
+from mypy_boto3_s3 import S3Client
+from sqlalchemy import Engine, create_engine
+
+from api.models import Base
+
+
+class Database:
+    """Database wrapper."""
+
+    def __init__(
+        self,
+        db_url: str,
+        connect_kwargs: dict[str, Any] | None = None,
+    ):
+        self.db_url = db_url
+        self.connect_kwargs = connect_kwargs or {}
+
+    @property
+    def engine(self) -> Engine:
+        if hasattr(self, "_engine"):
+            return cast(Engine, self._engine)  # type: ignore[has-type]
+        retries = 0
+        while True:
+            try:
+                engine = create_engine(self.db_url, connect_args=self.connect_kwargs)
+                # Attempt to create a connection or perform any necessary operations
+                engine.connect()
+                self._engine = engine
+                return engine  # Connection successful
+            except Exception as e:
+                if retries >= 10:
+                    raise RuntimeError(f"Could not create engine: {str(e)}")
+                retries += 1
+                time.sleep(60)
+
+    def create(self) -> None:
+        """Create database tables."""
+        Base.metadata.create_all(bind=self.engine)
+
+    def backup(self) -> bool:
+        """Backup the database. To be implemented by subclasses if supported."""
+        return False
+
+
+class SqliteDatabase(Database):
+    """SQLite database wrapper with optional S3 backup support."""
+
+    BACKUP_KEY = "userapi_sqlite_backup/backup.db.gz"
+
+    def __init__(
+        self,
+        db_url: str,
+        s3_client: S3Client | None = None,
+        s3_bucket: str | None = None,
+    ):
+        if not db_url.startswith("sqlite:///"):
+            raise ValueError(f"SQLiteRDSJobQueue requires SQLite DB URL, got: {db_url}")
+        if not ((s3_client is None) == (s3_bucket is None)):
+            raise ValueError(
+                "Both s3_client and s3_bucket must be provided for S3 backup/restore, or both must be None."
+            )
+        self.s3_client = s3_client
+        self.s3_bucket = s3_bucket
+        super().__init__(db_url, connect_kwargs={"check_same_thread": False})
+
+    def create(self) -> None:
+        self._restore_database()
+        super().create()
+
+    @property
+    def db_path(self) -> str:
+        return self.db_url[len("sqlite:///") :]
+
+    def backup(self) -> bool:
+        """Backup the SQLite database to S3."""
+        if not self.s3_bucket or not self.s3_client:
+            return False
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tmp_backup_path = os.path.join(temp_dir, "backup.db")
+            tmp_gzip_path = os.path.join(temp_dir, "backup.db.gz")
+            with sqlite3.connect(self.db_path) as source_conn:
+                with sqlite3.connect(tmp_backup_path) as backup_conn:
+                    source_conn.backup(backup_conn)
+
+            with open(tmp_backup_path, "rb") as f_in:
+                with gzip.open(tmp_gzip_path, "wb") as f_out:
+                    f_out.writelines(f_in)
+            self.s3_client.upload_file(tmp_gzip_path, self.s3_bucket, self.BACKUP_KEY)
+            return True
+
+    def _restore_database(self) -> bool:
+        """Restore the SQLite database from S3."""
+        if not self.s3_bucket or not self.s3_client:
+            return False
+
+        try:
+            self.s3_client.head_object(Bucket=self.s3_bucket, Key=self.BACKUP_KEY)
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                return False
+            raise
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tmp_gzip_path = os.path.join(temp_dir, "backup.db.gz")
+            tmp_backup_path = os.path.join(temp_dir, "backup.db")
+            self.s3_client.download_file(self.s3_bucket, self.BACKUP_KEY, tmp_gzip_path)
+            with gzip.open(tmp_gzip_path, "rb") as f_in:
+                with open(tmp_backup_path, "wb") as f_out:
+                    f_out.write(f_in.read())
+            os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
+            os.rename(tmp_backup_path, self.db_path)
+            return True

api/core/filesystem.py

@@ -7,11 +7,8 @@
 from pathlib import Path, PurePosixPath
 from typing import Any, BinaryIO, Callable, Generator
 
-import boto3
 import humanize
-from botocore.client import Config
 from botocore.response import StreamingBody
-from botocore.utils import fix_s3_host
 from fastapi import Request
 from fastapi.responses import FileResponse, StreamingResponse
 from mypy_boto3_s3 import S3Client
@@ -390,23 +387,16 @@ def download_url(
 def get_filesystem_with_root(
     root_path: str,
     filesystem: str,
-    s3_region: str,
-    s3_bucket: str | None,
+    s3_bucket: str | None = None,
+    s3_client: S3Client | None = None,
 ) -> FileSystem:
     """Get the filesystem to use."""
     predef_dirs = [e.value for e in models.UploadFileTypes] + [
         e.value for e in models.OutputEndpoints
     ]
     if filesystem == "s3":
+        assert s3_client is not None, "S3 client must be provided for S3 filesystem"
         assert s3_bucket is not None, "S3 bucket must be provided for S3 filesystem"
-        s3_client = boto3.client(
-            "s3",
-            region_name=s3_region,
-            endpoint_url=f"https://s3.{s3_region}.amazonaws.com",
-            config=Config(signature_version="v4", s3={"addressing_style": "path"}),
-        )
-        # this and config=... required to avoid DNS problems with new buckets
-        s3_client.meta.events.unregister("before-sign.s3", fix_s3_host)
         return S3Filesystem(root_path, s3_client, s3_bucket, predef_dirs=predef_dirs)
     elif filesystem == "local":
         return LocalFilesystem(root_path, predef_dirs=predef_dirs)
@@ -417,13 +407,13 @@ def get_filesystem_with_root(
 def user_filesystem_getter(
     user_data_root_path: str,
     filesystem: str,
-    s3_region: str,
-    s3_bucket: str | None,
+    s3_bucket: str | None = None,
+    s3_client: S3Client | None = None,
 ) -> Callable[[str], FileSystem]:
     """Get the filesystem to use for a user."""
     return lambda user_id: get_filesystem_with_root(
         str(Path(user_data_root_path) / user_id),
         filesystem=filesystem,
-        s3_region=s3_region,
         s3_bucket=s3_bucket,
+        s3_client=s3_client,
     )

api/database.py

@@ -1,42 +1,60 @@
-from typing import Any, Callable
+from typing import Any, Callable, Generator
 
+import boto3
 import requests
-from fastapi import Depends, Header, HTTPException, Request
+from botocore.config import Config
+from botocore.utils import fix_s3_host
+from fastapi import Depends, HTTPException, Request
 from fastapi.encoders import jsonable_encoder
-from fastapi.security import HTTPAuthorizationCredentials
-from fastapi_cloudauth.cognito import CognitoClaims, CognitoCurrentUser  # type: ignore
-from pydantic import BaseModel, Field
+from fastapi_cloudauth.cognito import CognitoClaims  # type: ignore
+from sqlalchemy.orm import Session, sessionmaker
 
 from api import settings
 from api.core import notifications
+from api.core.auth import APIKeyDependency, UserGroupCognitoCurrentUser
+from api.core.database import Database, SqliteDatabase
 from api.core.filesystem import FileSystem, user_filesystem_getter
 from api.schemas.job import QueueJob
 
+# S3 client setup
+s3_client = None
+if settings.s3_bucket:
+    s3_client = boto3.client(
+        "s3",
+        region_name=settings.s3_region,
+        endpoint_url=f"https://s3.{settings.s3_region}.amazonaws.com",
+        config=Config(signature_version="v4", s3={"addressing_style": "path"}),
+    )
+    # this and config=... required to avoid DNS problems with new buckets
+    s3_client.meta.events.unregister("before-sign.s3", fix_s3_host)
 
-class GroupClaims(CognitoClaims):  # type: ignore
-    """CognitoClaims with added groups claim."""
 
-    cognito_groups: list[str] | None = Field(alias="cognito:groups")
+# Database
+if settings.database_url.startswith("sqlite"):
+    db: Database = SqliteDatabase(
+        db_url=settings.database_url,
+        s3_client=s3_client,
+        s3_bucket=settings.s3_bucket,
+    )
+else:
+    db = Database(db_url=settings.database_url)
 
 
-class UserGroupCognitoCurrentUser(CognitoCurrentUser):  # type: ignore
-    """
-    Check membership in the 'users' group and add group membership information.
-    """
+async def db_dep() -> Database:
+    return db
 
-    user_info = GroupClaims
 
-    async def call(
-        self, http_auth: HTTPAuthorizationCredentials
-    ) -> BaseModel | dict[str, Any] | None:
-        user_info = await super().call(http_auth)
-        if "users" not in (getattr(user_info, "cognito_groups") or []):
-            raise HTTPException(
-                status_code=403, detail="Not a member of the 'users' group"
-            )
-        return user_info  # type: ignore
+def session_dep(db_dep: Database = Depends(db_dep)) -> Generator[Session, Any, None]:
+    """Get database session."""
+    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=db_dep.engine)
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
 
 
+# User authentication
 current_user_dep = UserGroupCognitoCurrentUser(
     region=settings.cognito_region,
     userPoolId=settings.cognito_user_pool_id,
@@ -52,12 +70,13 @@ async def current_user_global_dep(
     return current_user
 
 
+# Filesystem
 async def filesystem_getter_dep() -> Callable[[str], FileSystem]:
     """Get the user's filesystem getter."""
     return user_filesystem_getter(
         user_data_root_path=settings.user_data_root_path,
         filesystem=settings.filesystem,
-        s3_region=settings.s3_region,
+        s3_client=s3_client,
         s3_bucket=settings.s3_bucket,
     )
 
@@ -70,20 +89,11 @@ async def user_filesystem_dep(
     return filesystem_getter(current_user.username)
 
 
-class APIKeyDependency:
-    def __init__(self, key: str | None):
-        """Check API-internal key."""
-        self.key = key
-
-    def __call__(self, x_api_key: str | None = Header(...)) -> str | None:
-        if x_api_key != self.key:
-            raise HTTPException(status_code=401, detail="unauthorized")
-        return x_api_key
-
-
+# App-internal authentication (i.e. user-facing API <-> worker-facing API)
 workerfacing_api_auth_dep = APIKeyDependency(settings.internal_api_key_secret)
 
 
+# Notifications
 async def email_sender_dep() -> notifications.EmailSender:
     """Get the email sender."""
     service = settings.email_sender_service
@@ -110,6 +120,7 @@ async def email_sender_dep() -> notifications.EmailSender:
             )
 
 
+# Job enqueueing to worker-facing API
 async def enqueueing_function_dep() -> Callable[[QueueJob], None]:
     def enqueue(queue_item: QueueJob) -> None:
         resp = requests.post(

api/dependencies.py

@@ -5,8 +5,7 @@
 
 import api.core.notifications as notifications
 import api.crud.job as job_crud
-from api.database import get_db
-from api.dependencies import email_sender_dep, workerfacing_api_auth_dep
+from api.dependencies import email_sender_dep, session_dep, workerfacing_api_auth_dep
 from api.models import JobStates
 from api.schemas.job_update import JobUpdate
 
@@ -20,7 +19,7 @@
 )
 def update_job_status(
     update: JobUpdate,
-    db: Session = Depends(get_db),
+    db: Session = Depends(session_dep),
     email_sender: notifications.EmailSender = Depends(email_sender_dep),
 ) -> JobStates:
     db_job = job_crud.get_job(db, update.job_id)