Adding linting to migrations

.github/workflows/lint-migrations.yml

@@ -0,0 +1,73 @@
+name: Eugene CI check
+
+env:
+  EUGENE_VERSION: "0.6.1"
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  lint-migration:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.11.2
+      - uses: actions/cache@v2
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+      - name: Install python dependencies
+        run: |
+          export DISPATCH_LIGHT_BUILD=1
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+      - name: Download eugene
+        run: |
+          curl -L  https://github.com/kaaveland/eugene/releases/download/$EUGENE_VERSION/eugene-x86_64-unknown-linux-musl -o eugene
+          chmod +x eugene
+      - name: Generate sql files
+        run: |
+          # Navigate to the git repository root, if not already there
+          cd "$(git rev-parse --show-toplevel)" || exit
+
+          # Perform git diff to check for changes in the specified directory
+          changed_files=$(git diff --name-only -- src/dispatch/database/revisions/tenant/versions)
+
+          # Filter for newly added files (*.py)
+          new_files=$(echo "$changed_files" | grep -E "^src/dispatch/database/revisions/tenant/versions/2024-06-12_.*\.py$")
+
+          # Check if there is exactly one new file and handle errors for multiple or no files
+          if [[ $(echo "$new_files" | wc -l) -ne 1 ]]; then
+            echo "Error: There should be exactly one new revision file. Found: $(echo "$new_files" | wc -l)"
+            exit 1
+          fi
+
+          new_file_path=$(echo "$new_files" | head -n 1)
+          new_revision_id=$(basename "$new_file_path" .py | cut -d'_' -f2)
+
+          # Get the most recent revision ID excluding the newly added one
+          old_file=$(git ls-files src/dispatch/database/revisions/tenant/versions | grep -oP '2024-06-12_\K[^.]+(?=\.py)' | grep -v "$new_revision_id" | sort -rV | head -n 1)
+
+          if [[ -z "$old_file" ]]; then
+            echo "Error: No old revision ID found"
+            exit 1
+          fi
+
+          old_revision_id=$(echo "$old_file")
+          dispatch database upgrade --revision-type tenant --revision="${old_revision_id}:${new_revision_id}" --sql
+      - name: Lint files
+        run: ./eugene lint --git-diff origin/main migration-scripts -f md --accept-failures > lint.md
+      - name: Post Comment
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          COMMENT=$(cat lint.md)
+          gh pr comment ${{ github.event.pull_request.number }} --body "$COMMENT"

src/dispatch/cli.py

@@ -396,6 +396,8 @@ def drop_database(yes):
 def upgrade_database(tag, sql, revision, revision_type):
     """Upgrades database schema to newest version."""
     import sqlalchemy
+    from contextlib import redirect_stdout
+
     from alembic import command as alembic_command
     from alembic.config import Config as AlembicConfig
     from sqlalchemy import inspect
@@ -426,7 +428,12 @@ def upgrade_database(tag, sql, revision, revision_type):
                 path = config.ALEMBIC_TENANT_REVISION_PATH
 
             alembic_cfg.set_main_option("script_location", path)
-            alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
+            if sql:
+                with open("alembic_output.sql", "w") as sql_file:
+                    with redirect_stdout(sql_file):
+                        alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
+            else:
+                alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
         else:
             for path in [config.ALEMBIC_CORE_REVISION_PATH, config.ALEMBIC_TENANT_REVISION_PATH]:
                 alembic_cfg.set_main_option("script_location", path)

src/dispatch/database/revisions/core/env.py

@@ -45,7 +45,10 @@ def process_revision_directives(context, revision, directives):
             log.info("No changes found skipping revision creation.")
 
     connectable = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool
+        config.get_section(config.config_ini_section),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+        connect_args={'options': '-c lock_timeout=4000 -c statement_timeout=5000'}
     )
 
     log.info("Migrating dispatch core schema...")
@@ -58,14 +61,45 @@ def process_revision_directives(context, revision, directives):
             target_metadata=target_metadata,
             include_schemas=True,
             include_object=include_object,
+            transaction_per_migration=True,
             process_revision_directives=process_revision_directives,
         )
 
         with context.begin_transaction():
             context.run_migrations()
 
 
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well. By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+
+    # Get the URL from the config
+    url = config.get_main_option("sqlalchemy.url")
+
+    # Set the migration options
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        include_schemas=True,
+        include_object=include_object,
+        literal_binds=True,  # Binds parameters with their string values
+    )
+
+    # Start a transaction and run migrations
+    with context.begin_transaction():
+        context.run_migrations()
+
+
 if context.is_offline_mode():
-    log.info("Can't run migrations offline")
+    run_migrations_offline()
 else:
     run_migrations_online()

src/dispatch/database/revisions/tenant/env.py

@@ -1,3 +1,4 @@
+import os
 from alembic import context
 from sqlalchemy import engine_from_config, pool, inspect
 
@@ -51,7 +52,10 @@ def process_revision_directives(context, revision, directives):
             log.info("No changes found skipping revision creation.")
 
     connectable = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool
+        config.get_section(config.config_ini_section),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+        connect_args={'options': '-c lock_timeout=4000 -c statement_timeout=5000'}
     )
 
     with connectable.connect() as connection:
@@ -66,6 +70,7 @@ def process_revision_directives(context, revision, directives):
                 target_metadata=target_metadata,
                 include_schemas=True,
                 include_object=include_object,
+                transaction_per_migration=True,
                 process_revision_directives=process_revision_directives,
             )
 
@@ -77,7 +82,37 @@ def process_revision_directives(context, revision, directives):
                     break
 
 
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well. By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+
+    # Get the URL from the config
+    url = config.get_main_option("sqlalchemy.url")
+
+    # Set the migration options
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        include_schemas=True,
+        include_object=include_object,
+        literal_binds=True,  # Binds parameters with their string values
+    )
+
+    # Start a transaction and run migrations
+    with context.begin_transaction():
+        context.run_migrations()
+
+
 if context.is_offline_mode():
-    log.info("Can't run migrations offline")
+    run_migrations_offline()
 else:
     run_migrations_online()