Skip to content

Commit

Permalink
Allow running airflow against sqlite in-memory DB for tests (#37144)
Browse files Browse the repository at this point in the history
This makes no sense for production workloads, but it is useful for
writing tests of Airflow operators etc.

I also moved the warning from configure_vars to convigure_orm so that it
is checked after local settings have been applied in case that changes
anything.

(cherry picked from commit 7324400)
  • Loading branch information
ashb authored and potiuk committed Feb 13, 2024
1 parent d57042c commit 0dc8c28
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 14 deletions.
21 changes: 14 additions & 7 deletions airflow/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,13 +192,6 @@ def configure_vars():
global PLUGINS_FOLDER
global DONOT_MODIFY_HANDLERS
SQL_ALCHEMY_CONN = conf.get("database", "SQL_ALCHEMY_CONN")
if SQL_ALCHEMY_CONN.startswith("sqlite") and not SQL_ALCHEMY_CONN.startswith("sqlite:////"):
from airflow.exceptions import AirflowConfigException

raise AirflowConfigException(
f"Cannot use relative path: `{SQL_ALCHEMY_CONN}` to connect to sqlite. "
"Please use absolute path such as `sqlite:////tmp/airflow.db`."
)

DAGS_FOLDER = os.path.expanduser(conf.get("core", "DAGS_FOLDER"))

Expand Down Expand Up @@ -232,6 +225,20 @@ def configure_orm(disable_connection_pool=False, pool_class=None):
"""Configure ORM using SQLAlchemy."""
from airflow.utils.log.secrets_masker import mask_secret

if (
SQL_ALCHEMY_CONN
and SQL_ALCHEMY_CONN.startswith("sqlite")
and not SQL_ALCHEMY_CONN.startswith("sqlite:////")
# In memory is not useful for production, but useful for writing tests against Airflow for extensions
and SQL_ALCHEMY_CONN != "sqlite://"
):
from airflow.exceptions import AirflowConfigException

raise AirflowConfigException(
f"Cannot use relative path: `{SQL_ALCHEMY_CONN}` to connect to sqlite. "
"Please use absolute path such as `sqlite:////tmp/airflow.db`."
)

global Session
global engine
if os.environ.get("_AIRFLOW_SKIP_DB_TESTS") == "true":
Expand Down
24 changes: 17 additions & 7 deletions tests/core/test_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@
# under the License.
from __future__ import annotations

import contextlib
import os
import sys
import tempfile
from unittest import mock
from unittest.mock import MagicMock, call, patch
from unittest.mock import MagicMock, call

import pytest

Expand Down Expand Up @@ -216,11 +217,20 @@ def test_uses_updated_session_timeout_config_by_default(self):
assert session_lifetime_config == default_timeout_minutes


def test_sqlite_relative_path():
@pytest.mark.parametrize(
["value", "expectation"],
[
(
"sqlite:///./relative_path.db",
pytest.raises(AirflowConfigException, match=r"Cannot use relative path:"),
),
# Should not raise an exception
("sqlite://", contextlib.nullcontext()),
],
)
def test_sqlite_relative_path(monkeypatch, value, expectation):
from airflow import settings

with patch("airflow.settings.conf.get") as conf_get_mock:
conf_get_mock.return_value = "sqlite:///./relative_path.db"
with pytest.raises(AirflowConfigException) as exc:
settings.configure_vars()
assert "Cannot use relative path:" in str(exc.value)
monkeypatch.setattr(settings, "SQL_ALCHEMY_CONN", value)
with expectation:
settings.configure_orm()

0 comments on commit 0dc8c28

Please sign in to comment.