This is an automated email from the ASF dual-hosted git repository.

uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new e470d78462 Add support for configuring custom alembic file (#31415)
e470d78462 is described below

commit e470d784627502f171819fab072e0bbab4a05492
Author: Amogh Desai <[email protected]>
AuthorDate: Tue May 23 07:03:31 2023 +0530

    Add support for configuring custom alembic file (#31415)
    
    Co-authored-by: Hussein Awala <[email protected]>
    Co-authored-by: Amogh <[email protected]>
---
 airflow/config_templates/config.yml          |  8 ++++++++
 airflow/config_templates/default_airflow.cfg |  4 ++++
 airflow/utils/db.py                          |  6 +++++-
 tests/utils/test_db.py                       | 15 +++++++++++++++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/airflow/config_templates/config.yml 
b/airflow/config_templates/config.yml
index 6362e30db6..52b802a5c6 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -452,6 +452,14 @@ core:
 database:
   description: ~
   options:
+    alembic_ini_file_path:
+      description: |
+        Path to the ``alembic.ini`` file. You can either provide the file path 
relative
+        to the Airflow home directory or the absolute path if it is located 
elsewhere.
+      version_added: 2.7.0
+      type: string
+      example: ~
+      default: "alembic.ini"
     sql_alchemy_conn:
       description: |
         The SqlAlchemy connection string to the metadata database.
diff --git a/airflow/config_templates/default_airflow.cfg 
b/airflow/config_templates/default_airflow.cfg
index 93876c936b..ac350cf2a1 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -258,6 +258,10 @@ database_access_isolation = False
 # internal_api_url =
 
 [database]
+# Path to the ``alembic.ini`` file. You can either provide the file path 
relative
+# to the Airflow home directory or the absolute path if it is located 
elsewhere.
+alembic_ini_file_path = alembic.ini
+
 # The SqlAlchemy connection string to the metadata database.
 # SqlAlchemy supports many different database engines.
 # More information here:
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index d60655c7f0..015f535af0 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -736,7 +736,11 @@ def _get_alembic_config():
 
     package_dir = os.path.dirname(airflow.__file__)
     directory = os.path.join(package_dir, "migrations")
-    config = Config(os.path.join(package_dir, "alembic.ini"))
+    alembic_file = conf.get("database", "alembic_ini_file_path")
+    if os.path.isabs(alembic_file):
+        config = Config(alembic_file)
+    else:
+        config = Config(os.path.join(package_dir, alembic_file))
     config.set_main_option("script_location", directory.replace("%", "%%"))
     config.set_main_option("sqlalchemy.url", 
settings.SQL_ALCHEMY_CONN.replace("%", "%%"))
     return config
diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py
index 25a331a616..dee5118331 100644
--- a/tests/utils/test_db.py
+++ b/tests/utils/test_db.py
@@ -19,6 +19,7 @@ from __future__ import annotations
 
 import inspect
 import io
+import os
 import re
 from contextlib import redirect_stdout
 from unittest import mock
@@ -36,6 +37,7 @@ from airflow.exceptions import AirflowException
 from airflow.models import Base as airflow_base
 from airflow.settings import engine
 from airflow.utils.db import (
+    _get_alembic_config,
     check_migrations,
     compare_server_default,
     compare_type,
@@ -233,3 +235,16 @@ class TestDb:
             mock_init.assert_not_called()
         else:
             mock_init.assert_called_once_with(session=session_mock)
+
+    def test_alembic_configuration(self):
+        with mock.patch.dict(
+            os.environ, {"AIRFLOW__DATABASE__ALEMBIC_INI_FILE_PATH": 
"/tmp/alembic.ini"}, clear=True
+        ):
+            config = _get_alembic_config()
+            assert config.config_file_name == "/tmp/alembic.ini"
+
+        # default behaviour
+        config = _get_alembic_config()
+        import airflow
+
+        assert config.config_file_name == 
os.path.join(os.path.dirname(airflow.__file__), "alembic.ini")

Reply via email to