atlan-application-sdk-api 3.51.0__py3-none-any.whl

application_sdk/clients/sql.py

@@ -0,0 +1,1218 @@
+"""
+SQL client implementation for database connections.
+
+This module provides SQL client classes for both synchronous and asynchronous
+database operations, supporting batch processing and server-side cursors.
+"""
+
+import asyncio
+import concurrent.futures
+import contextlib
+import contextvars
+import datetime
+import functools
+import hashlib
+import math
+import re
+import threading
+from collections.abc import AsyncIterator, Callable, Iterator, Mapping
+from dataclasses import dataclass
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any, TypeVar, Union, cast
+from urllib.parse import quote, quote_plus
+
+from application_sdk._logging import get_logger
+from application_sdk._runtime.offload import (
+    CancelHandle,
+    run_in_thread,
+    submit_in_thread,
+)
+from application_sdk.clients._interface import ClientInterface
+from application_sdk.clients._sql_failure_classifier import (
+    SqlFailureKind,
+    classify_sql_exception,
+)
+from application_sdk.clients.models import DatabaseConfig
+from application_sdk.clients.sql_errors import (
+    EngineNotInitializedError,
+    InvalidSqlEngineTypeError,
+    MissingSqlParamError,
+    SqlAwsCredentialsError,
+    SqlBindParameterError,
+    SqlClientAuthFailedError,
+    SqlClientConfigError,
+    SqlClientConnectTimeoutError,
+    SqlClientConnectUnclassifiedError,
+    SqlClientSourceUnavailableError,
+    SqlCredentialsParseError,
+    SqlPandasResultError,
+    UnsupportedSqlCursorError,
+)
+from application_sdk.clients.sql_typecasters import install_tolerant_text_decoder_hook
+from application_sdk.common.aws_utils import (
+    _sts_deadline,
+    generate_aws_rds_token_with_iam_role,
+    generate_aws_rds_token_with_iam_user,
+)
+from application_sdk.constants import (
+    AWS_SESSION_NAME,
+    SQL_CLIENT_CONNECT_TIMEOUT_SECONDS,
+    USE_SERVER_SIDE_CURSOR,
+)
+from application_sdk.credentials.utils import parse_credentials_extra
+from application_sdk.errors import AppError, sanitize_cause_repr
+
+logger = get_logger(__name__)
+
+T = TypeVar("T")
+
+
+SqlBindValue = (
+    str
+    | int
+    | float
+    | bool
+    | Decimal
+    | datetime.datetime
+    | datetime.date
+    | datetime.time
+    | bytes
+    | None
+)
+"""A value ``run_query`` can bind to a ``:name`` placeholder: a scalar every DB-API driver accepts."""
+
+# SQL read just far enough to find placeholders: quoted text and comments are skipped whole, so a
+# ``:name`` inside them stays literal; outside them, a ``:name`` placeholder by SQLAlchemy
+# ``text()``'s own rule (not preceded by a colon, word character or backslash), or any other colon.
+# Escaping and validating ``params`` both read the query through this pattern, so they cannot
+# disagree on what is a bind. Quoting is standard SQL ('' doubles a quote; a backslash is literal)
+# plus MySQL backtick identifiers and Postgres dollar quotes.
+_SQL_COLON_SCAN = re.compile(
+    r"""
+      (?P<quoted>
+          '[^']*(?:''[^']*)*'
+        | "[^"]*(?:""[^"]*)*"
+        | `[^`]*`
+        | (?<!\w)\$(?P<tag>(?:[A-Za-z_]\w*)?)\$.*?\$(?P=tag)\$
+        | --[^\n]*
+        | /\*.*?\*/
+      )
+    | (?<![:\w\\]):(?P<name>\w+)
+    | :
+    """,
+    re.DOTALL | re.VERBOSE,
+)
+
+
+def _escape_colons_for_text(
+    query: str, bind_names: frozenset[str] = frozenset()
+) -> str:
+    """Backslash-escape literal colons so SQLAlchemy ``text()`` does not treat them as bind params.
+
+    ``text()`` reads any ``:name`` as a required bind parameter and raises at compile time on
+    colon-bearing SQL -- a regex ``(?:...)``, a ``::text`` cast, a ``'12:30:00'`` literal.
+    SQLAlchemy strips the backslash before the driver, so the DB receives the literal colon.
+
+    Every colon is escaped except a ``:name`` placeholder whose name is in ``bind_names``,
+    the parameters the caller is binding. A colon inside quoted text or a comment is always
+    escaped, so ``':name'`` stays a literal. Both colons of ``::`` are escaped, so ``x::name``
+    is never taken for the bind ``:name`` while ``:name::type`` binds ``name`` and keeps the cast.
+    """
+
+    def escape(match: re.Match[str]) -> str:
+        token = match.group(0)
+        if match.group("name") in bind_names:
+            return token
+        return token.replace(":", "\\:")
+
+    return _SQL_COLON_SCAN.sub(escape, query)
+
+
+def _bound_params(
+    query: str, params: Mapping[str, SqlBindValue] | None
+) -> dict[str, SqlBindValue] | None:
+    """``params`` checked against ``query``, or None when nothing is bound.
+
+    Raises :class:`SqlBindParameterError` for a name that ``query`` never references as a
+    ``:name`` placeholder outside quoted text and comments: executing anyway would run the query without that value, so a
+    filter would silently not apply.
+    """
+    if not params:
+        return None
+    referenced = {
+        m.group("name") for m in _SQL_COLON_SCAN.finditer(query) if m.group("name")
+    }
+    for name in params:
+        if name not in referenced:
+            raise SqlBindParameterError(
+                message=f"Bind parameter {name!r} is not referenced as a :name placeholder in the query"
+            )
+    return dict(params)
+
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from sqlalchemy.engine import URL, Connection
+    from sqlalchemy.engine.interfaces import DBAPICursor, ExecutionContext
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+    from sqlalchemy.orm import Session
+
+
+@dataclass(frozen=True)
+class _CancellableRead:
+    """One cancellable read, as seen from the worker thread running it."""
+
+    handle: CancelHandle
+    cancel_cursor: "Callable[[DBAPICursor], None]"
+
+
+#: The read the current worker thread is running, if it is cancellable. Set
+#: only inside the thread (run_in_thread runs each call in its own copy of the
+#: context), so it never leaks into the task or across calls.
+_ACTIVE_READ: contextvars.ContextVar[_CancellableRead | None] = contextvars.ContextVar(
+    "_ACTIVE_READ", default=None
+)
+
+
+@contextlib.contextmanager
+def _reading(read: _CancellableRead) -> Iterator[None]:
+    token = _ACTIVE_READ.set(read)
+    try:
+        yield
+    finally:
+        _ACTIVE_READ.reset(token)
+
+
+def _register_cursor_for_cancel(
+    conn: "Connection",
+    cursor: "DBAPICursor",
+    statement: str,
+    parameters: object,
+    context: "ExecutionContext | None",
+    executemany: bool,
+) -> None:
+    """``before_cursor_execute`` listener: hand the cursor to the read's handle.
+
+    Runs on the worker thread, at the first moment the DBAPI cursor exists.
+    A statement on any other thread — or outside a cancellable read — is left
+    alone.
+    """
+    read = _ACTIVE_READ.get()
+    if read is None:
+        return
+    read.handle.set(functools.partial(read.cancel_cursor, cursor))
+    if read.handle.requested:
+        # The caller is already gone; the action fired on set. Don't start a
+        # statement nobody is waiting for.
+        raise concurrent.futures.CancelledError
+
+
+def _install_cursor_cancel_listener(engine: object) -> None:
+    """Attach the cursor listener to ``engine`` once; skip anything not an Engine.
+
+    Done lazily on each read rather than in ``load()``, so a subclass that
+    builds its own engine is covered too. An async engine's reads don't run
+    on a worker thread, and a test double has no event surface.
+    """
+    from sqlalchemy import event  # noqa: PLC0415 — optional dep: sqlalchemy
+    from sqlalchemy.engine import Engine  # noqa: PLC0415 — optional dep: sqlalchemy
+
+    if not isinstance(engine, Engine):
+        return
+    if not event.contains(engine, "before_cursor_execute", _register_cursor_for_cancel):
+        event.listen(engine, "before_cursor_execute", _register_cursor_for_cancel)
+
+
+def _invalidate_if_cancelled(connection: "Connection", handle: CancelHandle) -> None:
+    """Drop a connection whose statement was cancelled instead of pooling it.
+
+    A cancelled connection may be mid-protocol, or carry server state the next
+    borrower does not expect; and a pool that pings or re-runs session setup on
+    checkout would pay a round trip to find out.
+    """
+    if handle.requested:
+        connection.invalidate()
+
+
+#: The connect-timeout keyword of each driver that has no connect timeout by
+#: default, so a black-holed host holds the connect (and its worker thread)
+#: until the OS gives up. Drivers that already bound their connect (pymysql,
+#: asyncpg, snowflake, pymssql, oracledb) are left alone so their bound is never
+#: loosened, and so are drivers whose ``timeout`` bounds every socket read
+#: (pg8000, redshift_connector): that would cut long queries short too.
+_DRIVER_CONNECT_TIMEOUT_ARG: dict[str, str] = {
+    "psycopg2": "connect_timeout",
+    "psycopg": "connect_timeout",
+    "mysqldb": "connect_timeout",
+}
+
+#: SQLAlchemy's default driver for a URL that names only the dialect.
+_DIALECT_DEFAULT_DRIVER: dict[str, str] = {
+    "postgresql": "psycopg2",
+    "mysql": "mysqldb",
+    "mariadb": "mysqldb",
+}
+
+
+def _connect_timeout_arg(connection_string: object) -> tuple[str, "URL"] | None:
+    """The URL's driver connect-timeout keyword and the parsed URL, if it has one.
+
+    None for a URL that does not parse or names a driver not in
+    ``_DRIVER_CONNECT_TIMEOUT_ARG``.
+    """
+    from sqlalchemy.engine import make_url  # noqa: PLC0415 — optional dep: sqlalchemy
+    from sqlalchemy.exc import ArgumentError  # noqa: PLC0415 — optional dep: sqlalchemy
+
+    try:
+        url = make_url(cast(Any, connection_string))
+    except (ArgumentError, ValueError):
+        return None
+    # Read from the URL rather than ``get_driver_name()``, which imports the
+    # dialect (and fails for one that is not installed).
+    dialect, _, driver = url.drivername.lower().partition("+")
+    driver = driver or _DIALECT_DEFAULT_DRIVER.get(dialect, "")
+    arg = _DRIVER_CONNECT_TIMEOUT_ARG.get(driver)
+    return None if arg is None else (arg, url)
+
+
+def _with_connect_timeout(
+    connection_string: object, connect_args: dict[str, Any], seconds: float
+) -> dict[str, Any]:
+    """``connect_args`` plus the driver's connect timeout, when it has none.
+
+    Never overrides a value already set in ``connect_args`` or in the URL's
+    query string. A URL that does not parse, or names a driver not in
+    ``_DRIVER_CONNECT_TIMEOUT_ARG``, gets ``connect_args`` unchanged.
+    """
+    found = _connect_timeout_arg(connection_string)
+    if found is None:
+        return connect_args
+    arg, url = found
+    if arg in connect_args or arg in url.query:
+        return connect_args
+    return {**connect_args, arg: max(1, math.ceil(seconds))}
+
+
+def _cap_next_connect_timeout(
+    engine: object, connection_string: object, seconds: float
+) -> None:
+    """Cap the driver connect timeout of ``engine``'s next connect at ``seconds``.
+
+    ``load()``'s ping starts after the connection string is built, which may
+    have spent much of the budget on STS. With the full timeout from
+    ``connect_args`` a black-holed connect would outlive ``load()`` and hold
+    its worker thread past the deadline; capped at what is left, the thread is
+    freed when ``load()`` gives up. Only the next connect (the ping) is capped:
+    later connects keep the full value. A value of any size — ours, from
+    ``connect_args`` or from the URL — is capped. Anything not an Engine (a
+    test double), or a driver not in ``_DRIVER_CONNECT_TIMEOUT_ARG``, is left
+    alone.
+    """
+    from sqlalchemy import event  # noqa: PLC0415 — optional dep: sqlalchemy
+    from sqlalchemy.engine import Engine  # noqa: PLC0415 — optional dep: sqlalchemy
+
+    found = _connect_timeout_arg(connection_string)
+    if not isinstance(engine, Engine) or found is None:
+        return
+    arg, _ = found
+    cap = max(1, math.ceil(seconds))
+
+    def _cap(
+        dialect: object, conn_rec: object, cargs: object, cparams: dict[str, Any]
+    ) -> None:
+        try:
+            current = int(float(cparams.get(arg, cap)))
+        except (TypeError, ValueError):
+            return  # not a number: the driver will reject it on its own
+        cparams[arg] = min(current, cap)
+
+    event.listen(engine, "do_connect", _cap, once=True)
+
+
+def _budget_left(budget: asyncio.Timeout) -> float:
+    """Seconds until ``budget`` expires, measured on the running loop's clock."""
+    return cast(float, budget.when()) - asyncio.get_running_loop().time()
+
+
+def _load_budget_expired() -> SqlClientConnectTimeoutError:
+    """The error for a ``load()`` that outran ``SQL_CLIENT_CONNECT_TIMEOUT_SECONDS``.
+
+    Raised in place of the bare ``TimeoutError`` the expired budget leaves
+    behind, which says nothing about what ran out of time.
+    """
+    return SqlClientConnectTimeoutError(
+        message=(
+            "SQL client did not connect to the source within "
+            f"{SQL_CLIENT_CONNECT_TIMEOUT_SECONDS:g}s"
+        )
+    )
+
+
+def _load_failure(exc: Exception) -> AppError:
+    """Type a failed ``load()`` by its root cause instead of calling it auth.
+
+    Classified by the same reader ``SqlApp.prime_sql_auth`` uses. An
+    unrecognised failure is reported as pending classification, not as a
+    rejected password.
+
+    No message carries driver text: SQLAlchemy embeds the connection string,
+    password included. The driver error rides ``cause``, which serialises
+    sanitised.
+    """
+    kind = classify_sql_exception(exc)
+    if kind is SqlFailureKind.AUTH:
+        return SqlClientAuthFailedError(cause=exc)
+    if kind is SqlFailureKind.TIMEOUT:
+        return SqlClientConnectTimeoutError(cause=exc)
+    if kind is SqlFailureKind.SOURCE_UNAVAILABLE:
+        return SqlClientSourceUnavailableError(cause=exc)
+    return SqlClientConnectUnclassifiedError(cause=exc)
+
+
+class BaseSQLClient(ClientInterface):
+    """SQL client for database operations.
+
+    This class provides functionality for connecting to and querying SQL databases,
+    with support for batch processing and server-side cursors.
+
+    Attributes:
+        connection: Database connection instance.
+        engine: SQLAlchemy engine instance.
+        credentials (Dict[str, Any]): Database credentials.
+        resolved_credentials (Dict[str, Any]): Resolved credentials after reading from secret manager.
+        use_server_side_cursor (bool): Whether to use server-side cursors.
+    """
+
+    connection = None
+    engine = None
+    credentials: dict[str, Any]
+    resolved_credentials: dict[str, Any]
+    use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR
+    DB_CONFIG: DatabaseConfig | None = None
+
+    def __init__(
+        self,
+        use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR,
+        credentials: dict[str, Any] | None = None,
+        chunk_size: int = 5000,
+    ):
+        """
+        Initialize the SQL client.
+
+        Args:
+            use_server_side_cursor (bool, optional): Whether to use server-side cursors.
+                Defaults to USE_SERVER_SIDE_CURSOR.
+            credentials (Optional[Dict[str, Any]], optional): Database credentials.
+                Defaults to None, which is treated as an empty dict.
+        """
+        self.use_server_side_cursor = use_server_side_cursor
+        self.credentials = credentials if credentials is not None else {}
+        self.resolved_credentials = {}
+        self.chunk_size = chunk_size
+
+    async def load(self, credentials: dict[str, Any]) -> None:
+        """Load credentials and prepare engine for lazy connections.
+
+        This method now only stores credentials and creates the engine without
+        establishing a persistent connection. Connections are created on-demand.
+
+        Args:
+            credentials (Dict[str, Any]): Database connection credentials.
+
+        Raises:
+            SqlClientAuthFailedError: If the source rejected the credentials.
+            SqlClientSourceUnavailableError: If the source could not be reached
+                (DNS, refused or filtered connection, TLS).
+            SqlClientConnectTimeoutError: If the connect or login timed out, or
+                ``load()`` as a whole outran ``ATLAN_SQL_CLIENT_CONNECT_TIMEOUT_SECONDS``.
+            SqlClientConnectUnclassifiedError: If the failure matched none of
+                the above.
+            AppError: Any typed error raised while connecting, unchanged — e.g.
+                ``MissingSqlParamError`` or ``SqlCredentialsParseError``.
+        """
+        if not self.DB_CONFIG:
+            raise SqlClientConfigError()
+
+        self.credentials = credentials  # Update the instance credentials
+        budget = asyncio.timeout(SQL_CLIENT_CONNECT_TIMEOUT_SECONDS)
+        try:
+            from sqlalchemy import (  # noqa: PLC0415 — optional dep: sqlalchemy
+                create_engine,
+            )
+
+            async with budget:
+                # Offloaded: for IAM-role auth this runs an STS AssumeRole,
+                # a network round trip that would otherwise block the loop.
+                # Cancelling the await does not stop the thread, so the STS
+                # call is fitted into the budget too.
+                with _sts_deadline(_budget_left(budget)):
+                    connection_string = await run_in_thread(
+                        self.get_sqlalchemy_connection_string
+                    )
+
+                # Create engine but no persistent connection
+                self.engine = create_engine(
+                    connection_string,
+                    connect_args=_with_connect_timeout(
+                        connection_string,
+                        self.DB_CONFIG.connect_args,
+                        SQL_CLIENT_CONNECT_TIMEOUT_SECONDS,
+                    ),
+                    pool_pre_ping=self.DB_CONFIG.pool_pre_ping,
+                )
+
+                # Install a tolerant UTF-8 text decoder for psycopg2/psycopg3 so
+                # query-history rows containing non-UTF-8 bytes (e.g. 0x96 from
+                # Windows-1252 paste) yield a logged U+FFFD instead of raising
+                # UnicodeDecodeError. psycopg3 is wired per connection; psycopg2 has
+                # to be remapped once per process, before any connection is opened —
+                # see sql_typecasters for why. Tracking: WARE-970, CONAT-767.
+                install_tolerant_text_decoder_hook(self.engine)
+
+                # Test connection briefly to validate credentials.
+                # Wrapped in run_in_thread because SQLAlchemy's synchronous
+                # engine.connect() blocks the event loop — critical for Temporal
+                # activities where blocking starves the auto-heartbeat. run_in_thread
+                # dispatches onto the SDK's dedicated pool rather than asyncio's
+                # shared default executor, which Temporal's own scheduling also uses.
+                # Capture engine in a local variable so the closure doesn't need to
+                # re-read self.engine (which is typed Optional) and pyright can narrow it.
+                _engine = self.engine
+
+                def _ping() -> None:
+                    with _engine.connect() as _:
+                        pass  # Connection test successful
+
+                # Cancelling the await does not stop the thread, so the driver
+                # must give up by the deadline on its own.
+                _cap_next_connect_timeout(
+                    _engine, connection_string, _budget_left(budget)
+                )
+                await run_in_thread(_ping)
+
+            # Don't store persistent connection
+            self.connection = None
+
+        # conformance: ignore[E004] exc_info omitted intentionally; SQLAlchemy embeds password in traceback; cause is sanitized and re-raised as typed error
+        except Exception as e:
+            # No exc_info here: SQLAlchemy errors embed the full connection
+            # string (including the password) in their message, and the
+            # traceback would print it verbatim into logs.
+            logger.error(  # conformance: ignore[E005,L004,L009] exc_info would expose SQLAlchemy password; logs sanitized cause at failure site — raised typed error wraps only the raw unsanitized cause
+                "Error loading SQL client: %s", sanitize_cause_repr(e)
+            )
+            if self.engine:
+                self.engine.dispose()
+                self.engine = None
+            if budget.expired():
+                raise _load_budget_expired() from e
+            # Already typed (a MissingSqlParamError from building the
+            # connection string, a connector's own leaf): keep its verdict.
+            if isinstance(e, AppError):
+                raise
+            raise _load_failure(e) from e
+
+    async def close(self) -> None:
+        """Close the database connection."""
+        if self.engine:
+            self.engine.dispose()
+            self.engine = None
+        self.connection = None  # Should already be None, but ensure cleanup
+
+    def cancel_cursor(self, dbapi_cursor: "DBAPICursor") -> None:
+        """Stop the statement running on ``dbapi_cursor``. Default: do nothing.
+
+        Called when the task awaiting a read (``get_results``,
+        ``get_batched_results``, ``run_query``) is cancelled while the
+        statement is still running on its worker thread. It runs on the SDK's
+        ``sdk-cancel-`` pool — a third thread, concurrently with the blocked
+        driver call — so override it only with a driver call documented as
+        safe from another thread, e.g. ``dbapi_cursor.cancel()`` or
+        ``dbapi_cursor.connection.cancel()``.
+
+        DBAPI has no standard cancel, so the default is a no-op: the read
+        still hands control back at once, and the statement runs to
+        completion on its thread. Either way the connection is invalidated
+        rather than returned to the pool.
+
+        A failure is logged at WARNING and otherwise ignored. Never stop or
+        suspend a warehouse here: that is not this statement's to stop.
+
+        Args:
+            dbapi_cursor: The driver's own cursor, as passed to SQLAlchemy's
+                ``before_cursor_execute`` event.
+        """
+
+    def get_iam_user_token(self):
+        """Get an IAM user token for AWS RDS database authentication.
+
+        This method generates a temporary authentication token for IAM user-based
+        authentication with AWS RDS databases. It requires AWS access credentials
+        and database connection details.
+
+        Returns:
+            str: A temporary authentication token for database access.
+
+        Raises:
+            SqlCredentialsParseError: If required credentials (username or database) are missing.
+        """
+        extra = parse_credentials_extra(self.credentials)
+        aws_access_key_id = self.credentials.get("username")
+        aws_secret_access_key = self.credentials.get("password")
+        host = self.credentials.get("host")
+        user = extra.get("username")
+        database = extra.get("database")
+        if not user:
+            raise SqlCredentialsParseError(
+                field="username",
+                message="username is required for IAM user authentication",
+            )
+        if not database:
+            raise SqlCredentialsParseError(
+                field="database",
+                message="database is required for IAM user authentication",
+            )
+
+        port = self.credentials.get("port")
+        region = self.credentials.get("region")
+        token = generate_aws_rds_token_with_iam_user(
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            host=host,
+            user=user,
+            port=port,
+            region=region,
+        )
+
+        return token
+
+    def get_iam_role_token(self):
+        """Get an IAM role token for AWS RDS database authentication.
+
+        This method generates a temporary authentication token for IAM role-based
+        authentication with AWS RDS databases. It requires an AWS role ARN and
+        database connection details.
+
+        Returns:
+            str: A temporary authentication token for database access.
+
+        Raises:
+            SqlAwsCredentialsError: If required credentials (aws_role_arn or database) are missing.
+        """
+        extra = parse_credentials_extra(self.credentials)
+        aws_role_arn = extra.get("aws_role_arn")
+        database = extra.get("database")
+        external_id = extra.get("aws_external_id")
+
+        if not aws_role_arn:
+            raise SqlAwsCredentialsError(
+                message="aws_role_arn is required for IAM role authentication",
+            )
+        if not database:
+            raise SqlAwsCredentialsError(
+                message="database is required for IAM role authentication",
+            )
+
+        session_name = AWS_SESSION_NAME
+        username = self.credentials.get("username")
+        host = self.credentials.get("host")
+        port = self.credentials.get("port")
+        region = self.credentials.get("region")
+
+        token = generate_aws_rds_token_with_iam_role(
+            role_arn=aws_role_arn,
+            host=host,
+            user=username,
+            external_id=external_id,
+            session_name=session_name,
+            port=port,
+            region=region,
+        )
+        return token
+
+    def get_auth_token(self) -> str:
+        """Get the appropriate authentication token based on auth type.
+
+        This method determines the authentication type from credentials and returns
+        the corresponding token. Supports basic auth, IAM user, and IAM role
+        authentication methods.
+
+        Returns:
+            str: URL-encoded authentication token.
+
+        Raises:
+            SqlCredentialsParseError: If an invalid authentication type is specified.
+        """
+        authType = self.credentials.get("authType", "basic")  # Default to basic auth
+        token = None
+
+        match authType:
+            case "iam_user":
+                token = self.get_iam_user_token()
+            case "iam_role":
+                token = self.get_iam_role_token()
+            case "basic":
+                token = self.credentials.get("password")
+            case _:
+                raise SqlCredentialsParseError(field="authType", value_summary=authType)
+
+        # Handle None values and ensure token is a string before encoding.
+        # This token is interpolated into the userinfo of a SQLAlchemy URL,
+        # which SQLAlchemy decodes with ``urllib.parse.unquote`` (percent-only).
+        # ``quote_plus`` encodes space as ``+``, and ``unquote`` never turns
+        # ``+`` back into a space, so a password containing a space would reach
+        # the driver corrupted (CONNECT-361). ``quote`` encodes space as ``%20``,
+        # which round-trips correctly. ``safe=""`` is required so that ``/`` in
+        # the credential is also encoded (an unencoded ``/`` in userinfo would
+        # otherwise break URL parsing). Query-string params below still use
+        # ``quote_plus`` — SQLAlchemy decodes query values with ``+`` -> space.
+        encoded_token = quote(str(token or ""), safe="")
+        return encoded_token
+
+    def add_connection_params(
+        self, connection_string: str, source_connection_params: dict[str, Any]
+    ) -> str:
+        """Add additional connection parameters to a SQLAlchemy connection string.
+
+        Args:
+            connection_string (str): Base SQLAlchemy connection string.
+            source_connection_params (Dict[str, Any]): Additional connection parameters
+                to append to the connection string. Keys and values must be
+                **raw, not pre-encoded** — both are URL-encoded here (via
+                ``quote_plus``) to prevent ``&``/``=`` injection, so a
+                pre-encoded value would be double-encoded.
+
+        Returns:
+            str: Connection string with additional parameters appended.
+        """
+        for key, value in source_connection_params.items():
+            if "?" not in connection_string:
+                connection_string += "?"
+            else:
+                connection_string += "&"
+            # URL-encode so values containing &, =, or spaces can't inject
+            # extra connection parameters (e.g. sslmode overrides).
+            connection_string += f"{quote_plus(str(key))}={quote_plus(str(value))}"
+
+        return connection_string
+
+    def get_sqlalchemy_connection_string(self) -> str:
+        """Generate a SQLAlchemy connection string for database connection.
+
+        This method constructs a connection string using the configured database
+        parameters and credentials. It handles different authentication methods
+        and includes necessary connection parameters.
+
+        Returns:
+            str: Complete SQLAlchemy connection string.
+
+        Raises:
+            SqlClientConfigError: If DB_CONFIG is not set.
+            MissingSqlParamError: If required connection parameters are missing.
+        """
+        if not self.DB_CONFIG:
+            raise SqlClientConfigError()
+
+        extra = parse_credentials_extra(self.credentials)
+
+        auth_token = self.get_auth_token()
+
+        # Prepare parameters
+        param_values = {}
+        for param in self.DB_CONFIG.required:
+            if param == "password":
+                param_values[param] = auth_token
+            else:
+                value = self.credentials.get(param) or extra.get(param)
+                if value is None:
+                    raise MissingSqlParamError(field=param)
+                param_values[param] = value
+
+        # Fill in base template
+        conn_str = self.DB_CONFIG.template.format(**param_values)
+
+        # Append defaults if not already in the template
+        if self.DB_CONFIG.defaults:
+            conn_str = self.add_connection_params(conn_str, self.DB_CONFIG.defaults)
+
+        if self.DB_CONFIG.parameters:
+            parameter_keys = self.DB_CONFIG.parameters
+            parameter_values = {
+                key: value
+                for key in parameter_keys
+                if (value := self.credentials.get(key) or extra.get(key)) is not None
+            }
+            if parameter_values:
+                conn_str = self.add_connection_params(conn_str, parameter_values)
+
+        return conn_str
+
+    async def run_query(
+        self,
+        query: str,
+        batch_size: int = 100000,
+        params: Mapping[str, SqlBindValue] | None = None,
+    ):
+        """Execute a SQL query and return results in batches using lazy connections.
+
+        This method creates a connection on-demand, executes the query in batches,
+        and automatically closes the connection when done. This prevents memory
+        leaks from persistent connections.
+
+        Args:
+            query (str): SQL query to execute.
+            batch_size (int, optional): Number of records to fetch in each batch.
+                Defaults to 100000.
+            params (Mapping[str, SqlBindValue], optional): Values to bind to ``:name``
+                placeholders in ``query``. The driver receives them as bound
+                values, never as SQL text, so a user-supplied filter value cannot
+                change the statement. Every name must appear in ``query``.
+
+        Yields:
+            List[Dict[str, Any]]: Batches of query results, where each result is
+                a dictionary mapping column names to values.
+
+        Raises:
+            EngineNotInitializedError: If engine is not initialized.
+            SqlBindParameterError: If ``params`` does not match the query.
+        """
+        if not self.engine:
+            raise EngineNotInitializedError()
+        bound = _bound_params(query, params)
+
+        logger.debug(
+            "Running query (sha=%s, len=%d, bound=%d)",
+            hashlib.sha256(query.encode("utf-8", errors="replace")).hexdigest()[:16],
+            len(query),
+            len(bound or {}),
+        )
+
+        _install_cursor_cancel_listener(self.engine)
+        engine = self.engine
+        # Driver calls, connect included, go through run_in_thread, so a cancelled caller stops
+        # waiting at once — but the thread keeps using the connection until the
+        # call returns. The lock keeps the close (below) from running on the
+        # connection while an abandoned call still holds it. The cancel handle
+        # is what makes that call return early: the listener registers the
+        # cursor on it, and cancel_cursor() fires off both of those threads.
+        in_use = threading.Lock()
+        handle = CancelHandle()
+        read = _CancellableRead(handle, self.cancel_cursor)
+        opened: list["Connection"] = []
+
+        def _holding_connection(func: Callable[..., T]) -> Callable[..., T]:
+            @functools.wraps(func)  # keeps the callee's name for run_in_thread's label
+            def call(*args: Any) -> T:
+                with in_use, _reading(read):
+                    return func(*args)
+
+            return call
+
+        def _connect() -> "Connection":
+            connection = engine.connect()
+            if handle.requested:
+                connection.close()
+                raise concurrent.futures.CancelledError
+            opened.append(connection)
+            return connection
+
+        def _close() -> None:
+            for connection in opened:
+                _invalidate_if_cancelled(connection, handle)
+                connection.close()
+
+        try:
+            connection = await run_in_thread(handle.bind(_holding_connection(_connect)))
+            if self.use_server_side_cursor:
+                connection = connection.execution_options(yield_per=batch_size)
+
+            from sqlalchemy import text  # noqa: PLC0415 — optional dep: sqlalchemy
+
+            statement = text(_escape_colons_for_text(query, frozenset(bound or ())))
+            execute = handle.bind(_holding_connection(connection.execute))
+            cursor = await run_in_thread(execute, statement, bound)
+            if not cursor or not cursor.cursor:
+                raise UnsupportedSqlCursorError()
+            column_names: list[str] = [
+                description.name.lower() for description in cursor.cursor.description
+            ]
+
+            fetchmany = handle.bind(_holding_connection(cursor.fetchmany))
+            while True:
+                rows = await run_in_thread(fetchmany, batch_size)
+                if not rows:
+                    break
+
+                results = [dict(zip(column_names, row)) for row in rows]
+                yield results
+        except BaseException:
+            # Cancellation, an early aclose(), or a failure. A driver call may
+            # still be running, so closing here on the loop would block it on
+            # that call; hand the close to the pool to run once the call ends.
+            submit_in_thread(_holding_connection(_close))
+            raise
+        await run_in_thread(_holding_connection(_close))
+
+        logger.info("Query execution completed")
+
+    def _execute_pandas_query(
+        self, conn, query, chunksize: int | None
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Helper function to execute SQL query using pandas.
+           The function is responsible for using import_optional_dependency method of the pandas library to import sqlalchemy
+           This function helps pandas in determining weather to use the sqlalchemy connection object and constructs like text()
+           or use the underlying database connection object. This has been done to make sure connectors like the Redshift connector,
+           which do not support the sqlalchemy connection object, can be made compatible with the application-sdk.
+
+        Args:
+            conn: Database connection object.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        import pandas as pd  # noqa: PLC0415 — optional dep: pandas
+        from pandas.compat._optional import (  # noqa: PLC0415 — optional dep: pandas
+            import_optional_dependency,
+        )
+        from sqlalchemy import text  # noqa: PLC0415 — optional dep: sqlalchemy
+
+        if import_optional_dependency("sqlalchemy", errors="ignore"):
+            return pd.read_sql_query(
+                text(_escape_colons_for_text(query)), conn, chunksize=chunksize
+            )
+        else:
+            dbapi_conn = getattr(conn, "connection", None)
+            return pd.read_sql_query(query, dbapi_conn, chunksize=chunksize)
+
+    def _read_sql_query(
+        self, session: "Session", query: str, chunksize: int | None
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Execute SQL query using the provided session.
+
+        Args:
+            session: SQLAlchemy session for database operations.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        conn = session.connection()
+        return self._execute_pandas_query(conn, query, chunksize=chunksize)
+
+    def _execute_query(
+        self, query: str, chunksize: int | None
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Execute SQL query using the provided engine and pandas.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        if not self.engine:
+            raise EngineNotInitializedError()
+
+        with self.engine.connect() as conn:
+            try:
+                return self._execute_pandas_query(conn, query, chunksize)
+            finally:
+                read = _ACTIVE_READ.get()
+                if read is not None:
+                    _invalidate_if_cancelled(conn, read.handle)
+
+    async def _execute_async_read_operation(
+        self, query: str, chunksize: int | None
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Helper to execute async read operation with either async session or thread executor."""
+        if isinstance(self.engine, str):
+            raise InvalidSqlEngineTypeError()
+
+        from sqlalchemy.ext.asyncio import (  # noqa: PLC0415 — optional dep: sqlalchemy
+            AsyncEngine,
+            AsyncSession,
+        )
+
+        async_session = None
+        if self.engine and isinstance(self.engine, AsyncEngine):
+            from sqlalchemy.orm import (  # noqa: PLC0415 — optional dep: sqlalchemy
+                sessionmaker,
+            )
+
+            async_session = sessionmaker(
+                self.engine, expire_on_commit=False, class_=AsyncSession
+            )
+
+        if async_session:
+            async with async_session() as session:
+                return await session.run_sync(
+                    self._read_sql_query, query, chunksize=chunksize
+                )
+        else:
+            # run_in_thread, not a per-call ThreadPoolExecutor: leaving that
+            # executor's `with` block on cancellation calls shutdown(wait=True)
+            # on the loop thread, which freezes the loop until the query ends.
+            # The shared pool lets a cancelled read return at once while its
+            # thread drains on its own — and the handle lets cancel_cursor()
+            # cut that drain short at the driver.
+            _install_cursor_cancel_listener(self.engine)
+            handle = CancelHandle()
+            read = _CancellableRead(handle, self.cancel_cursor)
+            execute_query = self._execute_query
+
+            @functools.wraps(execute_query)  # keeps the name for the hold label
+            def _cancellable_query(
+                query: str, chunksize: int | None
+            ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+                with _reading(read):
+                    return execute_query(query, chunksize)
+
+            return await run_in_thread(
+                handle.bind(_cancellable_query), query, chunksize
+            )
+
+    async def get_batched_results(
+        self,
+        query: str,
+    ) -> AsyncIterator["pd.DataFrame"] | Iterator["pd.DataFrame"]:  # type: ignore
+        """Get query results as batched pandas DataFrames asynchronously.
+
+        Returns:
+            AsyncIterator["pd.DataFrame"]: Async iterator yielding batches of query results.
+
+        Raises:
+            InvalidSqlEngineTypeError: If engine is a string instead of a SQLAlchemy engine.
+            SqlPandasResultError: If there's an error executing the query.
+        """
+        try:
+            # We cast to Iterator because passing chunk_size guarantees an Iterator return
+            result = await self._execute_async_read_operation(query, self.chunk_size)
+            return cast(Iterator["pd.DataFrame"], result)
+        except AppError:
+            raise
+        # conformance: ignore[E004] exception re-raised immediately as typed SqlPandasResultError; cause chain preserved via `from e`
+        except Exception as e:
+            # retryable=True: unclassified driver exception — do not override the
+            # activity's declared retry policy. See SqlPandasResultError docstring.
+            raise SqlPandasResultError(
+                message="Error reading batched data from SQL", cause=e, retryable=True
+            ) from e
+
+    async def get_results(self, query: str) -> "pd.DataFrame":
+        """Get all query results as a single pandas DataFrame asynchronously.
+
+        Returns:
+            pd.DataFrame: Query results as a DataFrame.
+
+        Raises:
+            InvalidSqlEngineTypeError: If engine is a string instead of a SQLAlchemy engine.
+            SqlPandasResultError: If there's an error executing the query.
+        """
+        try:
+            result = await self._execute_async_read_operation(query, None)
+            import pandas as pd  # noqa: PLC0415 — optional dep: pandas
+
+            if isinstance(result, pd.DataFrame):
+                return result
+            raise SqlPandasResultError(
+                invariant="_execute_async_read_operation must return a pd.DataFrame"
+            )
+        except AppError:
+            raise
+        # conformance: ignore[E004] exception re-raised immediately as typed SqlPandasResultError; cause chain preserved via `from e`
+        except Exception as e:
+            # retryable=True: unclassified driver exception — do not override the
+            # activity's declared retry policy. See SqlPandasResultError docstring.
+            raise SqlPandasResultError(cause=e, retryable=True) from e
+
+
+class AsyncBaseSQLClient(BaseSQLClient):
+    """Asynchronous SQL client for database operations.
+
+    This class extends BaseSQLClient to provide asynchronous database operations,
+    with support for batch processing and server-side cursors. It uses SQLAlchemy's
+    async engine and connection interfaces for non-blocking database operations.
+
+    Attributes:
+        connection (AsyncConnection): Async database connection instance.
+        engine (AsyncEngine): Async SQLAlchemy engine instance.
+        credentials (Dict[str, Any]): Database credentials.
+        use_server_side_cursor (bool): Whether to use server-side cursors.
+    """
+
+    connection: "AsyncConnection"
+    engine: "AsyncEngine"
+
+    async def load(self, credentials: dict[str, Any]) -> None:
+        """Load credentials and prepare async engine for lazy connections.
+
+        This method stores credentials and creates an async engine without establishing
+        a persistent connection. Connections are created on-demand for better memory efficiency.
+
+        Args:
+            credentials (Dict[str, Any]): Database connection credentials including
+                host, port, username, password, and other connection parameters.
+
+        Raises:
+            SqlClientConfigError: If DB_CONFIG is not set.
+            SqlClientAuthFailedError: If the source rejected the credentials.
+            SqlClientSourceUnavailableError: If the source could not be reached
+                (DNS, refused or filtered connection, TLS).
+            SqlClientConnectTimeoutError: If the connect or login timed out, or
+                ``load()`` as a whole outran ``ATLAN_SQL_CLIENT_CONNECT_TIMEOUT_SECONDS``.
+            SqlClientConnectUnclassifiedError: If the failure matched none of
+                the above.
+            AppError: Any typed error raised while connecting, unchanged — e.g.
+                ``MissingSqlParamError`` or ``SqlCredentialsParseError``.
+        """
+        self.credentials = credentials
+        if not self.DB_CONFIG:
+            raise SqlClientConfigError()
+
+        budget = asyncio.timeout(SQL_CLIENT_CONNECT_TIMEOUT_SECONDS)
+        try:
+            from sqlalchemy.ext.asyncio import (  # noqa: PLC0415 — optional dep: sqlalchemy
+                create_async_engine,
+            )
+
+            async with budget:
+                # Offloaded: for IAM-role auth this runs an STS AssumeRole,
+                # a network round trip that would otherwise block the loop.
+                # Cancelling the await does not stop the thread, so the STS
+                # call is fitted into the budget too.
+                with _sts_deadline(_budget_left(budget)):
+                    connection_string = await run_in_thread(
+                        self.get_sqlalchemy_connection_string
+                    )
+
+                # Create async engine but no persistent connection
+                self.engine = create_async_engine(
+                    connection_string,
+                    connect_args=_with_connect_timeout(
+                        connection_string,
+                        self.DB_CONFIG.connect_args,
+                        SQL_CLIENT_CONNECT_TIMEOUT_SECONDS,
+                    ),
+                    pool_pre_ping=self.DB_CONFIG.pool_pre_ping,
+                )
+
+                # Same WARE-970 tolerant-decoder hook as the sync client. Async
+                # engines surface DBAPI events on the underlying sync_engine.
+                install_tolerant_text_decoder_hook(self.engine.sync_engine)
+
+                # Test connection briefly to validate credentials
+                async with self.engine.connect() as _:
+                    pass  # Connection test successful
+
+            # Don't store persistent connection
+            self.connection = None
+
+        # conformance: ignore[E004] exc_info omitted intentionally; SQLAlchemy embeds password in traceback; cause is sanitized and re-raised as typed error
+        except Exception as e:
+            # No exc_info here: SQLAlchemy errors embed the full connection
+            # string (including the password) in their message, and the
+            # traceback would print it verbatim into logs.
+            logger.error(  # conformance: ignore[E005,L004,L009] exc_info would expose SQLAlchemy password; logs sanitized cause at failure site — raised typed error wraps only the raw unsanitized cause
+                "Error establishing database connection: %s", sanitize_cause_repr(e)
+            )
+            if self.engine:
+                await self.engine.dispose()
+                self.engine = None
+            if budget.expired():
+                raise _load_budget_expired() from e
+            # Already typed (a MissingSqlParamError from building the
+            # connection string, a connector's own leaf): keep its verdict.
+            if isinstance(e, AppError):
+                raise
+            raise _load_failure(e) from e
+
+    async def close(self) -> None:
+        """Close the async database connection and dispose of the engine."""
+        if self.engine:
+            await self.engine.dispose()
+            self.engine = None
+        self.connection = None
+
+    async def run_query(
+        self,
+        query: str,
+        batch_size: int = 100000,
+        params: Mapping[str, SqlBindValue] | None = None,
+    ):
+        """Execute a SQL query asynchronously and return results in batches using lazy connections.
+
+        This method creates an async connection on-demand, executes the query in batches,
+        and automatically closes the connection when done. This prevents memory leaks
+        from persistent connections.
+
+        Args:
+            query (str): SQL query to execute.
+            batch_size (int, optional): Number of records to fetch in each batch.
+                Defaults to 100000.
+            params (Mapping[str, SqlBindValue], optional): Values to bind to ``:name``
+                placeholders in ``query``, as for :meth:`BaseSQLClient.run_query`.
+
+        Yields:
+            List[Dict[str, Any]]: Batches of query results, where each result is
+                a dictionary mapping column names to values.
+
+        Raises:
+            EngineNotInitializedError: If engine is not initialized.
+        """
+        if not self.engine:
+            raise EngineNotInitializedError()
+        bound = _bound_params(query, params)
+
+        logger.debug(
+            "Running query (sha=%s, len=%d, bound=%d)",
+            hashlib.sha256(query.encode("utf-8", errors="replace")).hexdigest()[:16],
+            len(query),
+            len(bound or {}),
+        )
+        use_server_side_cursor = self.use_server_side_cursor
+
+        # Use async context manager for automatic connection cleanup
+        async with self.engine.connect() as connection:
+            try:
+                from sqlalchemy import text  # noqa: PLC0415 — optional dep: sqlalchemy
+
+                if use_server_side_cursor:
+                    connection = await connection.execution_options(
+                        yield_per=batch_size
+                    )
+
+                statement = text(_escape_colons_for_text(query, frozenset(bound or ())))
+                result = (
+                    await connection.stream(statement, bound)
+                    if use_server_side_cursor
+                    else await connection.execute(statement, bound)
+                )
+
+                column_names = list(result.keys())
+
+                while True:
+                    rows = (
+                        await result.fetchmany(batch_size)
+                        if use_server_side_cursor
+                        else result.cursor.fetchmany(batch_size)
+                        if result.cursor
+                        else None
+                    )
+                    if not rows:
+                        break
+                    yield [dict(zip(column_names, row)) for row in rows]
+
+            except AppError:
+                raise
+            # conformance: ignore[E004] exception re-raised immediately as typed SqlPandasResultError; cause chain preserved via `from e`
+            except Exception as e:
+                # retryable=True: unclassified driver exception — do not override
+                # the activity's declared retry policy. A transient source-side
+                # condition (concurrent DDL, lock contention, dropped connection)
+                # is indistinguishable from a permanent one here.
+                # See SqlPandasResultError docstring.
+                raise SqlPandasResultError(
+                    message="Error executing SQL query", cause=e, retryable=True
+                ) from e
+            # Async connection automatically closed by context manager
+
+        logger.info("Query execution completed")