Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Also use ODBC connection for sqlalchemy engine in OdbcHook like JdbcHook #43145

Merged
merged 10 commits into from
Oct 23, 2024
13 changes: 13 additions & 0 deletions providers/src/airflow/providers/odbc/hooks/odbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,19 @@ def connect_kwargs(self) -> dict:

return merged_connect_kwargs

def get_sqlalchemy_engine(self, engine_kwargs=None):
"""
Get an sqlalchemy_engine object.

:param engine_kwargs: Kwargs used in :func:`~sqlalchemy.create_engine`.
:return: the created engine.
"""
if engine_kwargs is None:
engine_kwargs = {}
engine_kwargs["creator"] = self.get_conn

return super().get_sqlalchemy_engine(engine_kwargs)

def get_conn(self) -> Connection:
"""Return ``pyodbc`` connection object."""
conn = connect(self.odbc_connection_string, **self.connect_kwargs)
Expand Down
9 changes: 9 additions & 0 deletions providers/tests/odbc/hooks/test_odbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import json
import logging
import sqlite3
from dataclasses import dataclass
from unittest import mock
from unittest.mock import patch
Expand Down Expand Up @@ -340,3 +341,11 @@ def test_query_no_handler_return_none(self):
hook = mock_hook(OdbcHook)
result = hook.run("SQL")
assert result is None

def test_get_sqlalchemy_engine_verify_creator_is_being_used(self):
hook = mock_hook(OdbcHook, conn_params={"extra": {"sqlalchemy_scheme": "sqlite"}})

with sqlite3.connect(":memory:") as connection:
hook.get_conn = lambda: connection
engine = hook.get_sqlalchemy_engine()
assert engine.connect().connection.connection == connection