diff --git a/airflow-core/src/airflow/settings.py b/airflow-core/src/airflow/settings.py index 54150dd8bfd5b..148254f86c51c 100644 --- a/airflow-core/src/airflow/settings.py +++ b/airflow-core/src/airflow/settings.py @@ -241,13 +241,34 @@ def load_policy_plugins(pm: pluggy.PluginManager): def _get_async_conn_uri_from_sync(sync_uri): - # Mapping of backend to async driver: + """ + Convert a sync SQLAlchemy URI to its async equivalent. + + Args: + sync_uri: The synchronous SQLAlchemy connection URI + + Returns: + The async URI with appropriate async driver (e.g., aiosqlite, asyncpg, aiomysql), + or the original URI if no async driver mapping exists + + Raises: + ValueError: If the URI is malformed (missing ':' scheme separator) + """ + # Mapping of backend to async driver. AIO_LIBS_MAPPING = { "sqlite": "aiosqlite", "postgresql": "psycopg_async" if _USE_PSYCOPG3 else "asyncpg", "mysql": "aiomysql", } + if not sync_uri or ":" not in sync_uri: + raise ValueError( + f"Invalid SQLAlchemy connection URI: {sync_uri!r}. " + "The URI must be in the format 'scheme://host/path' or similar with a ':' separator. " + "Check that AIRFLOW__DATABASE__SQL_ALCHEMY_CONN environment variable or " + "sql_alchemy_conn in airflow.cfg contains a valid database URI." + ) + scheme, rest = sync_uri.split(":", maxsplit=1) scheme = scheme.split("+", maxsplit=1)[0] aiolib = AIO_LIBS_MAPPING.get(scheme) diff --git a/airflow-core/tests/unit/core/test_settings.py b/airflow-core/tests/unit/core/test_settings.py index c703ee05ab941..7c0e48e825775 100644 --- a/airflow-core/tests/unit/core/test_settings.py +++ b/airflow-core/tests/unit/core/test_settings.py @@ -533,3 +533,62 @@ def test_early_return_when_all_none(self): settings.dispose_orm(do_log=False) mock_close.assert_not_called() + + +class TestGetAsyncConnUriFromSync: + """Tests for _get_async_conn_uri_from_sync function.""" + + @pytest.mark.parametrize( + ("sync_uri", "expected"), + [ + ("sqlite:///path/to/db.sqlite", "sqlite+aiosqlite:///path/to/db.sqlite"), + ("mysql://user:pass@localhost/dbname", "mysql+aiomysql://user:pass@localhost/dbname"), + ], + ) + def test_supported_scheme_conversion(self, sync_uri, expected): + """Test conversion of supported sync SQLAlchemy URIs to async driver variants.""" + result = settings._get_async_conn_uri_from_sync(sync_uri) + assert result == expected + + @pytest.mark.parametrize( + ("sync_uri", "expected_template"), + [ + ("postgresql://user:pass@localhost/dbname", "postgresql+{driver}://user:pass@localhost/dbname"), + ("postgresql+psycopg2://user@localhost/db", "postgresql+{driver}://user@localhost/db"), + ], + ) + def test_postgresql_scheme_conversion(self, sync_uri, expected_template): + """Test conversion of PostgreSQL sync URIs to the configured async driver.""" + postgres_async_driver = "psycopg_async" if settings._USE_PSYCOPG3 else "asyncpg" + assert settings._get_async_conn_uri_from_sync(sync_uri) == expected_template.format( + driver=postgres_async_driver + ) + + def test_unsupported_scheme_returns_original_uri(self): + """Test that unsupported schemes return the original URI unchanged.""" + uri = "oracle://user:pass@localhost:1521/dbname" + result = settings._get_async_conn_uri_from_sync(uri) + assert result == uri + + @pytest.mark.parametrize( + ("invalid_uri", "error_match"), + [ + ("", "Invalid SQLAlchemy connection URI"), + (None, "Invalid SQLAlchemy connection URI"), + ("notavaliduri", "Invalid SQLAlchemy connection URI.*':' separator"), + ], + ) + def test_invalid_uri_raises_value_error(self, invalid_uri, error_match): + """Test that invalid URIs raise ValueError with actionable guidance.""" + with pytest.raises(ValueError, match=error_match): + settings._get_async_conn_uri_from_sync(invalid_uri) + + def test_error_message_is_helpful(self): + """Test that error message contains helpful guidance.""" + with pytest.raises(ValueError, match="Invalid SQLAlchemy connection URI") as exc_info: + settings._get_async_conn_uri_from_sync("invalid_value") + + error_msg = str(exc_info.value) + assert "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" in error_msg + assert "sql_alchemy_conn" in error_msg + assert "airflow.cfg" in error_msg