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")