Fix airflow db check-migrations

[ashb]

This command broke after Define datetime and StringID column types centrally in migrations #19408 was merged, but we didn't notice as Github styles timedout checks badly.

---

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code change, Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in UPDATING.md.

[jedcunningham]

This fixes the check in my environment.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[potiuk]

Shoudl we add some test for that one? It might be similarly broken in the future?

[ashb]

Shoudl we add some test for that one? It might be similarly broken in the future?

Yes good call. Added a simple happy-path test for now.

[potiuk]

Yes good call. Added a simple happy-path test for now.

That's exactly what I had in mind :)

[potiuk]

Looks like prod image building was broken by a failed runner or smth (and it was needed to test the Helm Chart). Closed/reopened to re-run it

[uranusjr]

Helm Chart tests have been failing consistently for a while now.

[ashb]

Right -- I didn't think it was likely.

[potiuk]

I am able to reproduce it easily locally with breeze:

1. ./breeze kind-cluster start
2. ./breeze kind-cluster deploy
   it will fail eventually while waiting for webserver
3. ./breeze kind-cluster k9s:

[potiuk]

(wait - it was without the latest PR merged).

[potiuk]

Retrying

[ashb]

My kind cluster is now working (https://kind.sigs.k8s.io/docs/user/known-issues/#failure-to-create-cluster-with-cgroups-v2 was my problem), also trying.

[potiuk]

Yeah. Reproduced it with latest PR merged too (same sequence as above):

[potiuk]

Confirmed I have the PR in the airflow installed in the image, so this is a "legit" error

[potiuk]

AAAAAAHHHHH

{{ define "wait-for-migrations-command" }}
  {{/* From Airflow 2.0.0 this can become [airflow, db, check-migrations] */}}
  - python
  - -c
  - |
        import airflow
        import logging
        import os
        import time

        from alembic.config import Config
        from alembic.runtime.migration import MigrationContext
        from alembic.script import ScriptDirectory

        from airflow import settings

        package_dir = os.path.abspath(os.path.dirname(airflow.__file__))
        directory = os.path.join(package_dir, 'migrations')
        config = Config(os.path.join(package_dir, 'alembic.ini'))
        config.set_main_option('script_location', directory)
        config.set_main_option('sqlalchemy.url', settings.SQL_ALCHEMY_CONN.replace('%', '%%'))
        script_ = ScriptDirectory.from_config(config)

        timeout=60

        with settings.engine.connect() as connection:
            context = MigrationContext.configure(connection)
            ticker = 0
            while True:
                source_heads = set(script_.get_heads())

                db_heads = set(context.get_current_heads())
                if source_heads == db_heads:
                    break

                if ticker >= timeout:
                    raise TimeoutError("There are still unapplied migrations after {} seconds.".format(ticker))
                ticker += 1
                time.sleep(1)
                logging.info('Waiting for migrations... %s second(s)', ticker)
{{- end }}

[potiuk]

Not the best idea to have it in the helm chart.

[kaxil]

Whhhoooopsss yeah -- I think it was in Helm Chart as previous Airflow versions didn't have support for "airflow db wait-for-migrations" command

[ashb]

OH!

[ashb]

      initContainers:
      - args:
        - python
        - -c
        - |
          import airflow
          import logging
          import os
          import time

          from alembic.config import Config
          from alembic.runtime.migration import MigrationContext
          from alembic.script import ScriptDirectory

          from airflow import settings

          package_dir = os.path.abspath(os.path.dirname(airflow.__file__))
          directory = os.path.join(package_dir, 'migrations')
          config = Config(os.path.join(package_dir, 'alembic.ini'))
          config.set_main_option('script_location', directory)
          config.set_main_option('sqlalchemy.url', settings.SQL_ALCHEMY_CONN.replace('%', '%%'))
          script_ = ScriptDirectory.from_config(config)

          timeout=60

          with settings.engine.connect() as connection:
              context = MigrationContext.configure(connection)
              ticker = 0
              while True:
                  source_heads = set(script_.get_heads())

                  db_heads = set(context.get_current_heads())
                  if source_heads == db_heads:
                      break

                  if ticker >= timeout:
                      raise TimeoutError("There are still unapplied migrations after {} seconds.".format(ticker))
                  ticker += 1
                  time.sleep(1)
                  logging.info('Waiting for migrations... %s second(s)', ticker)

[ashb]

Yeah, that's the problem. The helm chart isn't using that command.

5 mins too slow.

[potiuk]

So eventually the Helm errors ARE real errors :). Happy we have the tests.

[potiuk]

I think for now reverting the change is the only choice and it will be a breaking change for anyone using Helm < 1.4.0 even if we fix it there (and run airflow db wait-for-migrations for 2+)

[ashb]

Something like this will do the job:

try:
   from airflow.cli.commands.db_command import check_migrations
except ImportError:
   # existing code

[potiuk]

Yeah. But the problem is for those who already use the chart

[ashb]

We could (somehow) say the Airflow 2.3 will need Helm Chart >= 1.4 I guess? Far from ideal.

[kaxil]

Not possible to make it backwards compatible somehow 😬 ? #19408 is mainly just to centralize the IDs and types that are used in Migration

[potiuk]

Another option is smart workaround. For example what we could do:

in settings (from airlfow import settings) we could check if this import is run from "wait-for-db-migrations" container and if so - we could run the migration using the "check_migrations" and sys.exit(0).

Not the "cleanest" solution but should work.

[potiuk]

or even in "init.py" of "airflow" package - it is imported in the first line

[potiuk]

Should not be too difficult - we could check what was the command python was executed with :)

[potiuk]

And as of chart 1.4.0 we would do it "properly".

[ashb]

Here is the "proper" fix anyway which should make tests go green #19676

[ashb]

And back-compat fix incoming.

[ashb]

#19677

[ashb]

This is fixed in main now, and 2.3 won't need a particular version of the Helm chart either.