dabla commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3507320902
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -1136,3 +1149,194 @@ def get_db_log_messages(self, conn) -> None:
:param conn: Connection object
"""
+
+ def _translate_sql(self, sql: str) -> str:
+ """
+ Translate SQL to driver-specific paramstyle.
+
+ DB-specific hooks may override this to translate from a canonical style
+ to their driver's paramstyle if you want a unified SQL authoring style.
+ """
+ return sql
+
+ def _prepare_parameters(self, parameters: Iterable | Mapping[str, Any] |
None):
+ """DB hooks may override to adapt parameter style."""
+ return parameters
+
+ def _build_conn_kwargs_from_airflow_connection(self, db) -> dict:
+ """Build a DB-agnostic kwargs dict."""
+ extra = {}
+ extra_dejson = getattr(db, "extra_dejson", None)
+ if isinstance(extra_dejson, dict):
+ extra = extra_dejson
+ try:
+ port = int(db.port) if db.port else None
+ except (TypeError, ValueError):
+ port = None
+ return {
+ "host": db.host or "",
+ "port": port,
+ "username": db.login or "",
+ "password": db.password or "",
+ "database": extra.get("database") or extra.get("dbname") or "",
+ "schema": db.schema or "",
+ "conn_type": db.conn_type,
+ "extra": extra,
+ "raw_connection": db,
+ }
+
+ async def async_get_conn(self) -> Any:
+ if not hasattr(self, "_conn_lock"):
+ self._conn_lock = asyncio.Lock()
+ async with self._conn_lock:
+ if not self._connection:
+ self._connection = await
sync_to_async(self.get_connection)(self.get_conn_id())
+ db = self._connection
+ if self.connector is None:
+ raise RuntimeError(f"{type(self).__name__} didn't have
`self.connector` set!")
+ conn_kwargs = self._build_conn_kwargs_from_airflow_connection(db)
+ return await self.connector.connect(**conn_kwargs)
+
+ async def _call(self, func: Callable, *args, **kwargs) -> Any:
+ if inspect.iscoroutinefunction(func):
+ return await func(*args, **kwargs)
+ return await sync_to_async(func)(*args, **kwargs)
+
+ @asynccontextmanager
+ async def _async_create_autocommit_connection(self, autocommit: bool =
False):
+ conn = await self.async_get_conn()
+ try:
+ if self.supports_autocommit:
+ set_autocommit = getattr(conn, "set_autocommit", None)
+ if set_autocommit and
inspect.iscoroutinefunction(set_autocommit):
+ await set_autocommit(autocommit)
+ else:
+ self.set_autocommit(conn, autocommit)
+ yield conn
+ finally:
+ close = getattr(conn, "aclose", None) or getattr(conn, "close",
None)
Review Comment:
Should the second
--
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]
