RehanAhmad25 commented on code in PR #74539:
URL: https://github.com/apache/airflow/pull/74539#discussion_r4308127347
##########
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:
Confirmed, tuple("dbo") gives ('d', 'b', 'o'). Only mssql_python Row objects
are converted now, everything else is returned as is. I read Row through
import_module so the import stays lazy. The existing test uses the Row class
patched in, and I added cases for a scalar, ints, a dict and lists of non-rows.
--
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]
