karenbraganz commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3991747024
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -1171,3 +1185,220 @@ 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 aget_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
get_async_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 _acreate_autocommit_connection(self, autocommit: bool = False):
+ conn = await self.aget_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)
+ if close:
+ await self._call(close)
+
+ @asynccontextmanager
+ async def _aget_cursor(self, conn):
+ cursor = getattr(conn, "cursor", None)
+ if cursor is None:
+ raise TypeError(
+ f"{type(conn).__name__} has no cursor() method. "
+ "Override _aget_cursor in the DB-specific hook."
+ )
+ cur_or_cm = cursor()
+ if inspect.isawaitable(cur_or_cm):
+ cur_or_cm = await cur_or_cm
+ if hasattr(cur_or_cm, "__aenter__") and hasattr(cur_or_cm,
"__aexit__"):
+ async with cur_or_cm as cur:
+ yield cur
+ return
+ execute = getattr(cur_or_cm, "execute", None)
+ if execute and inspect.iscoroutinefunction(execute):
+ try:
+ yield cur_or_cm
+ finally:
+ close = getattr(cur_or_cm, "aclose", None) or
getattr(cur_or_cm, "close", None)
+ if close:
+ if inspect.iscoroutinefunction(close):
+ await close()
+ else:
+ close()
+ return
+ raise RuntimeError(
+ f"Unsupported cursor type returned by {type(conn).__name__}. "
+ "Override _aget_cursor in the DB-specific hook."
+ )
+
+ def supports_async_execution(self) -> bool:
+ """Whether the DB-specific hook overrides :meth:`aget_conn` for a real
async driver."""
+ return type(self).aget_conn is not DbApiHook.aget_conn
+
+ def supports_readonly_execution(self) -> bool:
+ """Whether the DB-specific hook overrides :meth:`_aenter_read_only`."""
+ return type(self)._aenter_read_only is not DbApiHook._aenter_read_only
+
+ async def _aenter_read_only(self, conn):
+ """
+ Put the opened async connection in read-only mode.
+
+ Called by :meth:`arun` before any statement runs, whenever the operator
+ requests read-only execution. The base implementation raises an
exception
+ if this method is not implemented in the DB-specific hook.
+ """
+ raise NotImplementedError(
+ f"{type(self).__name__} does not implement read-only execution.
Override "
+ "_aenter_read_only to support deferrable read-only queries for
this database,"
Review Comment:
It is a typo.
--
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]
