dabla commented on code in PR #74539:
URL: https://github.com/apache/airflow/pull/74539#discussion_r4307639231
##########
providers/microsoft/mssql/src/airflow/providers/microsoft/mssql/hooks/mssql.py:
##########
@@ -110,16 +174,57 @@ def get_conn(self) -> PymssqlConnection:
**extra_conn_args,
)
- def set_autocommit(
- self,
- conn: PymssqlConnection,
- autocommit: bool,
- ) -> None:
- conn.autocommit(autocommit)
+ def _get_mssql_python_conn(self) -> Any:
+ """Return ``mssql_python`` connection object."""
+ try:
+ mssql_python = import_module("mssql_python")
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "The 'mssql_python' dbapi_driver needs the mssql-python
package. Install it with: "
+ "pip install
'apache-airflow-providers-microsoft-mssql[mssql-python]'"
+ ) from e
- def get_autocommit(self, conn: PymssqlConnection):
+ conn = self.connection
+ server = conn.host or ""
+ if server and conn.port:
+ server = f"{server},{conn.port}"
+ params: dict[str, Any] = {
+ "Server": server,
+ "Database": self.schema or conn.schema,
+ "UID": conn.login,
+ "PWD": conn.password,
+ }
+ params.update(
+ {
+ key: val
+ for key, val in conn.extra_dejson.items()
+ if key.lower() not in self._MSSQL_PYTHON_IGNORED_EXTRAS
+ }
+ )
+ # The keyword arguments are escaped by mssql-python when it builds the
connection string.
+ return mssql_python.connect(**{key: val for key, val in params.items()
if val not in (None, "")})
+
+ def set_autocommit(self, conn: Any, autocommit: bool) -> None:
+ if self.dbapi_driver == self.DBAPI_DRIVER_MSSQL_PYTHON:
+ super().set_autocommit(conn, autocommit)
+ else:
+ conn.autocommit(autocommit)
+
+ def get_autocommit(self, conn: Any) -> bool:
+ if self.dbapi_driver == self.DBAPI_DRIVER_MSSQL_PYTHON:
+ return super().get_autocommit(conn)
return conn.autocommit_state
+ def _make_common_data_structure(self, result: Any) -> tuple | list[tuple]
| None:
+ """Turn the ``Row`` objects of mssql_python into plain tuples."""
+ if self.dbapi_driver != self.DBAPI_DRIVER_MSSQL_PYTHON:
+ return super()._make_common_data_structure(result)
+ if result is None:
+ return None
+ if isinstance(result, list):
Review Comment:
**[warning]** Every handler result is converted, not only the driver's `Row`
objects.
`DbApiHook.run()` passes whatever the handler returns to
`_make_common_data_structure`. The pymssql path and the base class return
results that are not rows unchanged. This override calls `tuple()` on anything
that is not `None` or a list:
- `hook.run("SELECT SCHEMA_NAME()", handler=lambda cur: cur.fetchone()[0])`
returns `('d', 'b', 'o')` instead of `'dbo'`.
- A handler that returns `cur.rowcount` (an `int`) raises `TypeError: 'int'
object is not iterable`.
- A handler that returns a list of dicts gets a list of key tuples.
mssql-python exports its row class (`from mssql_python import Row`), so the
conversion can be limited to rows and leave anything else alone:
```python
from mssql_python import Row
if isinstance(result, Row):
return tuple(result)
if isinstance(result, list):
return [tuple(row) if isinstance(row, Row) else row for row in result]
return result
```
The test currently uses a hand-written `Row` stand-in. Add a case with a
scalar handler result to pin this down.
##########
providers/microsoft/mssql/src/airflow/providers/microsoft/mssql/hooks/mssql.py:
##########
@@ -97,10 +151,20 @@ def get_sqlalchemy_connection(
engine = self.get_sqlalchemy_engine(engine_kwargs=engine_kwargs)
return engine.connect(**(connect_kwargs or {}))
- def get_conn(self) -> PymssqlConnection:
+ def get_conn(self) -> Any:
+ """Return a ``pymssql`` or ``mssql_python`` connection object,
depending on ``dbapi_driver``."""
+ if self.dbapi_driver == self.DBAPI_DRIVER_MSSQL_PYTHON:
Review Comment:
**[nit]** Six methods branch on the same `dbapi_driver` value.
`placeholder`, `_default_sqlalchemy_scheme`, `get_conn`, `set_autocommit`,
`get_autocommit` and `_make_common_data_structure` each repeat `if
self.dbapi_driver == self.DBAPI_DRIVER_MSSQL_PYTHON`. Supporting a third driver
later would mean editing all six, and a missed branch fails quietly ([2] is one
example of a branch that grew wider than its driver). One small class per
driver, holding the connect, the autocommit pair, the row conversion and the
defaults, and chosen once from `dbapi_driver`, keeps each driver's behaviour in
one place ([Replace Conditional with
Polymorphism](https://refactoring.com/catalog/replaceConditionalWithPolymorphism.html)).
This is fine as a follow-up if you'd rather keep this PR small.
##########
providers/microsoft/mssql/src/airflow/providers/microsoft/mssql/hooks/mssql.py:
##########
@@ -80,12 +134,12 @@ def get_uri(self) -> str:
from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit
r = list(urlsplit(super().get_uri()))
- # change pymssql driver:
+ # change the sqlalchemy driver:
r[0] = self.sqlalchemy_scheme
- # remove query string 'sqlalchemy_scheme' like parameters:
+ # remove query string parameters that are meant for the hook, not for
the sqlalchemy dialect:
qs = parse_qs(r[3], keep_blank_values=True)
for k in list(qs.keys()):
- if k.lower() == "sqlalchemy_scheme":
+ if k.lower() in self._AIRFLOW_ONLY_EXTRAS:
Review Comment:
**[warning]** With `mssql_python`, `get_uri()` keeps the hook's own extras,
and the SQLAlchemy route then fails to connect.
`get_conn()` drops `_MSSQL_PYTHON_IGNORED_EXTRAS` (`placeholder`,
`escape_column_names`, `insert_statement_format`, ...), but `get_uri()` only
drops `_AIRFLOW_ONLY_EXTRAS`. On the SQLAlchemy route the two drivers handle
unknown keys in opposite ways:
- `mssql_python.connect(**kwargs)` logs a warning and skips an unknown
keyword (`Connection._construct_connection_string`, "Ignoring unknown
connection parameter from kwargs").
- SQLAlchemy 2.1's `mssql+mssqlpython` dialect puts every URL query key into
the connection string (`MSDialect_mssqlpython.create_connect_args`).
mssql-python parses that string with
`_ConnectionStringParser(validate_keywords=True)`, which raises
`ConnectionStringParseError("Unknown keyword 'placeholder' is not recognized")`.
So a connection with `{"dbapi_driver": "mssql_python", "placeholder": "?"}`
works with `get_conn()` but fails with `get_sqlalchemy_engine()`,
`get_sqlalchemy_connection()` and `hook.inspector`. `inspector` is not only
reached directly: `insert_rows()` without `target_fields` calls
`dialect.get_target_fields()`, which uses it.
The comment at lines 64-65 has the same mix-up. The kwargs route does not
need the filter to work, but the URI route does. Suggestion: in `get_uri()`,
strip `_MSSQL_PYTHON_IGNORED_EXTRAS` when the driver is `mssql_python`. Add a
test that sets `placeholder` and asserts the key is absent from the URI. The
current `test_get_uri_strips_dbapi_driver` only covers `dbapi_driver`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
