From c85e11aa48c7fb16eb0c07e941198330e1aa9065 Mon Sep 17 00:00:00 2001 From: Kevin Glisson Date: Mon, 24 Jun 2024 14:58:47 -0700 Subject: [PATCH] Adding linting to migrations --- .github/workflows/lint-migrations.yml | 73 +++++++++++++++++++ src/dispatch/cli.py | 9 ++- src/dispatch/database/revisions/core/env.py | 38 +++++++++- src/dispatch/database/revisions/tenant/env.py | 39 +++++++++- 4 files changed, 154 insertions(+), 5 deletions(-) create mode 100644 .github/workflows/lint-migrations.yml diff --git a/.github/workflows/lint-migrations.yml b/.github/workflows/lint-migrations.yml new file mode 100644 index 000000000000..cbe9abc6a719 --- /dev/null +++ b/.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" diff --git a/src/dispatch/cli.py b/src/dispatch/cli.py index 02551ba34115..5361700e5043 100644 --- a/src/dispatch/cli.py +++ b/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) diff --git a/src/dispatch/database/revisions/core/env.py b/src/dispatch/database/revisions/core/env.py index 0e6967292ec3..c001ddfbf068 100644 --- a/src/dispatch/database/revisions/core/env.py +++ b/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,6 +61,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, ) @@ -65,7 +69,37 @@ def process_revision_directives(context, revision, directives): 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() diff --git a/src/dispatch/database/revisions/tenant/env.py b/src/dispatch/database/revisions/tenant/env.py index 0aba8c3ee437..8e3d98a96cad 100644 --- a/src/dispatch/database/revisions/tenant/env.py +++ b/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()