sqlmesh.core.config.connection

   1from __future__ import annotations
   2
   3import abc
   4import base64
   5import importlib
   6import logging
   7import os
   8import pathlib
   9import re
  10import typing as t
  11from enum import Enum
  12from functools import partial
  13from sys import version_info
  14
  15import pydantic
  16from pydantic import Field, computed_field
  17from packaging import version
  18from pydantic_core import from_json
  19from sqlglot import exp
  20from sqlglot.errors import ParseError
  21from sqlglot.helper import subclasses
  22
  23from sqlmesh.core import engine_adapter
  24from sqlmesh.core.config.base import BaseConfig
  25from sqlmesh.core.config.common import (
  26    compile_regex_mapping,
  27    concurrent_tasks_validator,
  28    http_headers_validator,
  29)
  30from sqlmesh.core.engine_adapter import EngineAdapter
  31from sqlmesh.core.engine_adapter.shared import CatalogSupport
  32from sqlmesh.utils import debug_mode_enabled, str_to_bool
  33from sqlmesh.utils.aws import validate_s3_uri
  34from sqlmesh.utils.errors import ConfigError
  35from sqlmesh.utils.pydantic import (
  36    ValidationInfo,
  37    field_validator,
  38    get_concrete_types_from_typehint,
  39    model_validator,
  40    validation_data,
  41    validation_error_message,
  42)
  43
  44if t.TYPE_CHECKING:
  45    from sqlmesh.core._typing import Self
  46
  47logger = logging.getLogger(__name__)
  48
  49RECOMMENDED_STATE_SYNC_ENGINES = {
  50    "postgres",
  51    "gcp_postgres",
  52    "mysql",
  53    "mssql",
  54    "azuresql",
  55}
  56FORBIDDEN_STATE_SYNC_ENGINES = {
  57    # Do not support row-level operations
  58    "spark",
  59    "trino",
  60    # Nullable types are problematic
  61    "clickhouse",
  62    "starrocks",
  63}
  64MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
  65PASSWORD_REGEX = re.compile(r"(password=)(\S+)")
  66SUPPORTS_MSSQL_PYTHON_DRIVER = (version_info.major, version_info.minor) >= (3, 10)
  67
  68
  69def _get_engine_import_validator(
  70    import_name: str, engine_type: str, extra_name: t.Optional[str] = None, decorate: bool = True
  71) -> t.Callable:
  72    extra_name = extra_name or engine_type
  73
  74    def validate(cls: t.Any, data: t.Any) -> t.Any:
  75        check_import = (
  76            str_to_bool(str(data.pop("check_import", True))) if isinstance(data, dict) else True
  77        )
  78        if not check_import:
  79            return data
  80        try:
  81            importlib.import_module(import_name)
  82        except ImportError:
  83            if debug_mode_enabled():
  84                raise
  85
  86            logger.exception("Failed to import the engine library")
  87
  88            raise ConfigError(
  89                f"Failed to import the '{engine_type}' engine library. This may be due to a missing "
  90                "or incompatible installation. Please ensure the required dependency is installed by "
  91                f'running: `pip install "sqlmesh[{extra_name}]"`. For more details, check the logs '
  92                "in the 'logs/' folder, or rerun the command with the '--debug' flag."
  93            )
  94
  95        return data
  96
  97    return model_validator(mode="before")(validate) if decorate else validate
  98
  99
 100class ConnectionConfig(abc.ABC, BaseConfig):
 101    type_: str
 102    DIALECT: t.ClassVar[str]
 103    DISPLAY_NAME: t.ClassVar[str]
 104    DISPLAY_ORDER: t.ClassVar[int]
 105    concurrent_tasks: int
 106    register_comments: bool
 107    pre_ping: bool
 108    pretty_sql: bool = False
 109    schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None
 110    catalog_type_overrides: t.Optional[t.Dict[str, str]] = None
 111
 112    # Whether to share a  single connection across threads or create a new connection per thread.
 113    #
 114    # MyPy throws a "Decorators on top of @property are not supported" error despite this being a
 115    # valid decoration, and Pydantic recommend disabling the MyPy hint for this reason - see:
 116    # https://pydantic.dev/docs/validation/2.0/usage/computed_fields/
 117    @computed_field  # type: ignore[prop-decorator]
 118    @property
 119    def shared_connection(self) -> bool:
 120        return False
 121
 122    @property
 123    @abc.abstractmethod
 124    def _connection_kwargs_keys(self) -> t.Set[str]:
 125        """keywords that should be passed into the connection"""
 126
 127    @property
 128    @abc.abstractmethod
 129    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 130        """The engine adapter for this connection"""
 131
 132    @property
 133    @abc.abstractmethod
 134    def _connection_factory(self) -> t.Callable:
 135        """A function that is called to return a connection object for the given Engine Adapter"""
 136
 137    @property
 138    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 139        """The static connection kwargs for this connection"""
 140        return {}
 141
 142    @property
 143    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 144        """kwargs that are for execution config only"""
 145        return {}
 146
 147    @property
 148    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
 149        """A function that is called to initialize the cursor"""
 150        return None
 151
 152    @property
 153    def is_recommended_for_state_sync(self) -> bool:
 154        """Whether this engine is recommended for being used as a state sync for production state syncs"""
 155        return self.type_ in RECOMMENDED_STATE_SYNC_ENGINES
 156
 157    @property
 158    def is_forbidden_for_state_sync(self) -> bool:
 159        """Whether this engine is forbidden from being used as a state sync"""
 160        return self.type_ in FORBIDDEN_STATE_SYNC_ENGINES
 161
 162    @property
 163    def _connection_factory_with_kwargs(self) -> t.Callable[[], t.Any]:
 164        """A function that is called to return a connection object for the given Engine Adapter"""
 165        return partial(
 166            self._connection_factory,
 167            **{
 168                **self._static_connection_kwargs,
 169                **{k: v for k, v in self.dict().items() if k in self._connection_kwargs_keys},
 170            },
 171        )
 172
 173    def connection_validator(self) -> t.Callable[[], None]:
 174        """A function that validates the connection configuration"""
 175        return self.create_engine_adapter().ping
 176
 177    def create_engine_adapter(
 178        self, register_comments_override: bool = False, concurrent_tasks: t.Optional[int] = None
 179    ) -> EngineAdapter:
 180        """Returns a new instance of the Engine Adapter."""
 181
 182        concurrent_tasks = concurrent_tasks or self.concurrent_tasks
 183        return self._engine_adapter(
 184            self._connection_factory_with_kwargs,
 185            multithreaded=concurrent_tasks > 1,
 186            default_catalog=self.get_catalog(),
 187            cursor_init=self._cursor_init,
 188            register_comments=register_comments_override or self.register_comments,
 189            pre_ping=self.pre_ping,
 190            pretty_sql=self.pretty_sql,
 191            shared_connection=self.shared_connection,
 192            schema_differ_overrides=self.schema_differ_overrides,
 193            catalog_type_overrides=self.catalog_type_overrides,
 194            **self._extra_engine_config,
 195        )
 196
 197    def get_catalog(self) -> t.Optional[str]:
 198        """The catalog for this connection"""
 199        if hasattr(self, "catalog"):
 200            return self.catalog
 201        if hasattr(self, "database"):
 202            return self.database
 203        if hasattr(self, "db"):
 204            return self.db
 205        return None
 206
 207    @model_validator(mode="before")
 208    @classmethod
 209    def _expand_json_strings_to_concrete_types(cls, data: t.Any) -> t.Any:
 210        """
 211        There are situations where a connection config class has a field that is some kind of complex type
 212        (eg a list of strings or a dict) but the value is being supplied from a source such as an environment variable
 213
 214        When this happens, the value is supplied as a string rather than a Python object. We need some way
 215        of turning this string into the corresponding Python list or dict.
 216
 217        Rather than doing this piecemeal on every config subclass, this provides a generic implementatation
 218        to identify fields that may be be supplied as JSON strings and handle them transparently
 219        """
 220        if data and isinstance(data, dict):
 221            for maybe_json_field_name in cls._get_list_and_dict_field_names():
 222                if (value := data.get(maybe_json_field_name)) and isinstance(value, str):
 223                    # crude JSON check as we dont want to try and parse every string we get
 224                    value = value.strip()
 225                    if value.startswith("{") or value.startswith("["):
 226                        data[maybe_json_field_name] = from_json(value)
 227
 228        return data
 229
 230    @classmethod
 231    def _get_list_and_dict_field_names(cls) -> t.Set[str]:
 232        field_names = set()
 233        for name, field in cls.model_fields.items():
 234            if field.annotation:
 235                field_types = get_concrete_types_from_typehint(field.annotation)
 236
 237                # check if the field type is something that could concievably be supplied as a json string
 238                if any(ft is t for t in (list, tuple, set, dict) for ft in field_types):
 239                    field_names.add(name)
 240
 241        return field_names
 242
 243
 244class DuckDBAttachOptions(BaseConfig):
 245    type: str
 246    path: str
 247    read_only: bool = False
 248
 249    # DuckLake specific options
 250    data_path: t.Optional[str] = None
 251    override_data_path: t.Optional[bool] = False
 252    encrypted: bool = False
 253    data_inlining_row_limit: t.Optional[int] = None
 254    metadata_schema: t.Optional[str] = None
 255
 256    def to_sql(self, alias: str) -> str:
 257        options = []
 258        # 'duckdb' is actually not a supported type, but we'd like to allow it for
 259        # fully qualified attach options or integration testing, similar to duckdb-dbt
 260        if self.type not in ("duckdb", "ducklake", "motherduck"):
 261            options.append(f"TYPE {self.type.upper()}")
 262        if self.read_only:
 263            options.append("READ_ONLY")
 264
 265        # DuckLake specific options
 266        path = self.path
 267        if self.type == "ducklake":
 268            if not path.startswith("ducklake:"):
 269                path = f"ducklake:{path}"
 270            if self.data_path is not None:
 271                options.append(f"DATA_PATH '{self.data_path}'")
 272                if self.override_data_path:
 273                    options.append("OVERRIDE_DATA_PATH true")
 274            if self.encrypted:
 275                options.append("ENCRYPTED")
 276            if self.data_inlining_row_limit is not None:
 277                options.append(f"DATA_INLINING_ROW_LIMIT {self.data_inlining_row_limit}")
 278            if self.metadata_schema is not None:
 279                options.append(f"METADATA_SCHEMA '{self.metadata_schema}'")
 280
 281        options_sql = f" ({', '.join(options)})" if options else ""
 282        alias_sql = ""
 283        # TODO: Add support for Postgres schema. Currently adding it blocks access to the information_schema
 284
 285        # MotherDuck does not support aliasing
 286        alias_sql = (
 287            f" AS {alias}" if not (self.type == "motherduck" or self.path.startswith("md:")) else ""
 288        )
 289        return f"ATTACH IF NOT EXISTS '{path}'{alias_sql}{options_sql}"
 290
 291
 292class BaseDuckDBConnectionConfig(ConnectionConfig):
 293    """Common configuration for the DuckDB-based connections.
 294
 295    Args:
 296        database: The optional database name. If not specified, the in-memory database will be used.
 297        catalogs: Key is the name of the catalog and value is the path.
 298        extensions: A list of autoloadable extensions to load.
 299        connector_config: A dictionary of configuration to pass into the duckdb connector.
 300        secrets: A list of dictionaries used to generate DuckDB secrets for authenticating with external services (e.g. S3).
 301        filesystems: A list of dictionaries used to register `fsspec` filesystems to the DuckDB cursor.
 302        concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
 303        register_comments: Whether or not to register model comments with the SQL engine.
 304        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
 305        token: The optional MotherDuck token. If not specified and a MotherDuck path is in the catalog, the user will be prompted to login with their web browser.
 306    """
 307
 308    database: t.Optional[str] = None
 309    catalogs: t.Optional[t.Dict[str, t.Union[str, DuckDBAttachOptions]]] = None
 310    extensions: t.List[t.Union[str, t.Dict[str, t.Any]]] = []
 311    connector_config: t.Dict[str, t.Any] = {}
 312    secrets: t.Union[t.List[t.Dict[str, t.Any]], t.Dict[str, t.Dict[str, t.Any]]] = []
 313    filesystems: t.List[t.Dict[str, t.Any]] = []
 314
 315    concurrent_tasks: int = 1
 316    register_comments: bool = True
 317    pre_ping: t.Literal[False] = False
 318
 319    token: t.Optional[str] = None
 320
 321    @computed_field  # type: ignore[prop-decorator]
 322    @property
 323    def shared_connection(self) -> bool:
 324        return True
 325
 326    _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
 327
 328    @model_validator(mode="before")
 329    def _validate_database_catalogs(cls, data: t.Any) -> t.Any:
 330        if not isinstance(data, dict):
 331            return data
 332
 333        db_path = data.get("database")
 334        if db_path and data.get("catalogs"):
 335            raise ConfigError(
 336                "Cannot specify both `database` and `catalogs`. Define all your catalogs in `catalogs` and have the first entry be the default catalog"
 337            )
 338        if isinstance(db_path, str) and db_path.startswith("md:"):
 339            raise ConfigError(
 340                "Please use connection type 'motherduck' without the `md:` prefix if you want to use a MotherDuck database as the single `database`."
 341            )
 342
 343        return data
 344
 345    @property
 346    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 347        return engine_adapter.DuckDBEngineAdapter
 348
 349    @property
 350    def _connection_kwargs_keys(self) -> t.Set[str]:
 351        return {"database"}
 352
 353    @property
 354    def _connection_factory(self) -> t.Callable:
 355        import duckdb
 356
 357        return duckdb.connect
 358
 359    @property
 360    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
 361        """A function that is called to initialize the cursor"""
 362        import duckdb
 363        from duckdb import BinderException
 364
 365        def init(cursor: duckdb.DuckDBPyConnection) -> None:
 366            for extension in self.extensions:
 367                extension = extension if isinstance(extension, dict) else {"name": extension}
 368
 369                install_command = f"INSTALL {extension['name']}"
 370
 371                if extension.get("repository"):
 372                    install_command = f"{install_command} FROM {extension['repository']}"
 373
 374                if extension.get("force_install"):
 375                    install_command = f"FORCE {install_command}"
 376
 377                try:
 378                    cursor.execute(install_command)
 379                    cursor.execute(f"LOAD {extension['name']}")
 380                except Exception as e:
 381                    raise ConfigError(f"Failed to load extension {extension['name']}: {e}")
 382
 383            if self.connector_config:
 384                option_names = list(self.connector_config)
 385                in_part = ",".join("?" for _ in range(len(option_names)))
 386
 387                cursor.execute(
 388                    f"SELECT name, value FROM duckdb_settings() WHERE name IN ({in_part})",
 389                    option_names,
 390                )
 391
 392                existing_values = {field: setting for field, setting in cursor.fetchall()}
 393
 394                # only set connector_config items if the values differ from what is already set
 395                # trying to set options like 'temp_directory' even to the same value can throw errors like:
 396                # Not implemented Error: Cannot switch temporary directory after the current one has been used
 397                for field, setting in self.connector_config.items():
 398                    if existing_values.get(field) != setting:
 399                        try:
 400                            cursor.execute(f"SET {field} = '{setting}'")
 401                        except Exception as e:
 402                            raise ConfigError(
 403                                f"Failed to set connector config {field} to {setting}: {e}"
 404                            )
 405
 406            if self.secrets:
 407                duckdb_version = duckdb.__version__
 408                if version.parse(duckdb_version) < version.parse("0.10.0"):
 409                    from sqlmesh.core.console import get_console
 410
 411                    get_console().log_warning(
 412                        f"DuckDB version {duckdb_version} does not support secrets-based authentication (requires 0.10.0 or later).\n"
 413                        "To use secrets, please upgrade DuckDB. For older versions, configure legacy authentication via `connector_config`.\n"
 414                        "More info: https://duckdb.org/docs/stable/extensions/httpfs/s3api_legacy_authentication.html"
 415                    )
 416                else:
 417                    if isinstance(self.secrets, list):
 418                        secrets_items = [(secret_dict, "") for secret_dict in self.secrets]
 419                    else:
 420                        secrets_items = [
 421                            (secret_dict, secret_name)
 422                            for secret_name, secret_dict in self.secrets.items()
 423                        ]
 424
 425                    for secret_dict, secret_name in secrets_items:
 426                        secret_settings: t.List[str] = []
 427                        for field, setting in secret_dict.items():
 428                            secret_settings.append(f"{field} '{setting}'")
 429                        if secret_settings:
 430                            secret_clause = ", ".join(secret_settings)
 431                            try:
 432                                cursor.execute(
 433                                    f"CREATE OR REPLACE SECRET {secret_name} ({secret_clause});"
 434                                )
 435                            except Exception as e:
 436                                raise ConfigError(f"Failed to create secret: {e}")
 437
 438            if self.filesystems:
 439                from fsspec import filesystem  # type: ignore
 440
 441                for file_system in self.filesystems:
 442                    options = file_system.copy()
 443                    fs = options.pop("fs")
 444                    fs = filesystem(fs, **options)
 445                    cursor.register_filesystem(fs)
 446
 447            for i, (alias, path_options) in enumerate(
 448                (getattr(self, "catalogs", None) or {}).items()
 449            ):
 450                # we parse_identifier and generate to ensure that `alias` has exactly one set of quotes
 451                # regardless of whether it comes in quoted or not
 452                alias = exp.parse_identifier(alias, dialect="duckdb").sql(
 453                    identify=True, dialect="duckdb"
 454                )
 455                try:
 456                    if isinstance(path_options, DuckDBAttachOptions):
 457                        query = path_options.to_sql(alias)
 458                    else:
 459                        query = f"ATTACH IF NOT EXISTS '{path_options}'"
 460                        if not path_options.startswith("md:"):
 461                            query += f" AS {alias}"
 462                    cursor.execute(query)
 463                except BinderException as e:
 464                    # If a user tries to create a catalog pointing at `:memory:` and with the name `memory`
 465                    # then we don't want to raise since this happens by default. They are just doing this to
 466                    # set it as the default catalog.
 467                    # If a user tried to attach a MotherDuck database/share which has already by attached via
 468                    # `ATTACH 'md:'`, then we don't want to raise since this is expected.
 469                    if (
 470                        not (
 471                            'database with name "memory" already exists' in str(e)
 472                            and path_options == ":memory:"
 473                        )
 474                        and f"""database with name "{path_options.path.replace("md:", "")}" already exists"""
 475                        not in str(e)
 476                    ):
 477                        raise e
 478                if i == 0 and not getattr(self, "database", None):
 479                    cursor.execute(f"USE {alias}")
 480
 481        return init
 482
 483    def create_engine_adapter(
 484        self, register_comments_override: bool = False, concurrent_tasks: t.Optional[int] = None
 485    ) -> EngineAdapter:
 486        """Checks if another engine adapter has already been created that shares a catalog that points to the same data
 487        file. If so, it uses that same adapter instead of creating a new one. As a result, any additional configuration
 488        associated with the new adapter will be ignored."""
 489        data_files = set((self.catalogs or {}).values())
 490        if self.database:
 491            if isinstance(self, MotherDuckConnectionConfig):
 492                data_files.add(
 493                    f"md:{self.database}"
 494                    + (f"?motherduck_token={self.token}" if self.token else "")
 495                )
 496            else:
 497                data_files.add(self.database)
 498        data_files.discard(":memory:")
 499        for data_file in data_files:
 500            key = data_file if isinstance(data_file, str) else data_file.path
 501            adapter = BaseDuckDBConnectionConfig._data_file_to_adapter.get(key)
 502            if adapter is not None:
 503                logger.info(
 504                    f"Using existing DuckDB adapter due to overlapping data file: {self._mask_sensitive_data(key)}"
 505                )
 506                return adapter
 507
 508        if data_files:
 509            masked_files = {
 510                self._mask_sensitive_data(file if isinstance(file, str) else file.path)
 511                for file in data_files
 512            }
 513            logger.info(f"Creating new DuckDB adapter for data files: {masked_files}")
 514        else:
 515            logger.info("Creating new DuckDB adapter for in-memory database")
 516        adapter = super().create_engine_adapter(
 517            register_comments_override, concurrent_tasks=concurrent_tasks
 518        )
 519        for data_file in data_files:
 520            key = data_file if isinstance(data_file, str) else data_file.path
 521            BaseDuckDBConnectionConfig._data_file_to_adapter[key] = adapter
 522        return adapter
 523
 524    def get_catalog(self) -> t.Optional[str]:
 525        if self.database:
 526            # Remove `:` from the database name in order to handle if `:memory:` is passed in
 527            return pathlib.Path(self.database.replace(":memory:", "memory")).stem
 528        if self.catalogs:
 529            return list(self.catalogs)[0]
 530        return None
 531
 532    def _mask_sensitive_data(self, string: str) -> str:
 533        # Mask MotherDuck tokens with fixed number of asterisks
 534        result = MOTHERDUCK_TOKEN_REGEX.sub(
 535            lambda m: f"{m.group(1)}{m.group(2)}{'*' * 8 if m.group(3) else ''}", string
 536        )
 537        # Mask PostgreSQL/MySQL passwords with fixed number of asterisks
 538        result = PASSWORD_REGEX.sub(lambda m: f"{m.group(1)}{'*' * 8}", result)
 539        return result
 540
 541
 542class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
 543    """Configuration for the MotherDuck connection."""
 544
 545    type_: t.Literal["motherduck"] = Field(alias="type", default="motherduck")
 546    DIALECT: t.ClassVar[t.Literal["duckdb"]] = "duckdb"
 547    DISPLAY_NAME: t.ClassVar[t.Literal["MotherDuck"]] = "MotherDuck"
 548    DISPLAY_ORDER: t.ClassVar[t.Literal[5]] = 5
 549
 550    @property
 551    def _connection_kwargs_keys(self) -> t.Set[str]:
 552        return set()
 553
 554    @property
 555    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 556        """kwargs that are for execution config only"""
 557        from sqlmesh import __version__
 558
 559        custom_user_agent_config = {"custom_user_agent": f"SQLMesh/{__version__}"}
 560        connection_str = "md:"
 561        if self.database:
 562            # Attach single MD database instead of all databases on the account
 563            connection_str += f"{self.database}?attach_mode=single"
 564        if self.token:
 565            connection_str += f"{'&' if self.database else '?'}motherduck_token={self.token}"
 566        return {"database": connection_str, "config": custom_user_agent_config}
 567
 568    @property
 569    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 570        return {"is_motherduck": True}
 571
 572
 573class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
 574    """Configuration for the DuckDB connection."""
 575
 576    type_: t.Literal["duckdb"] = Field(alias="type", default="duckdb")
 577    DIALECT: t.ClassVar[t.Literal["duckdb"]] = "duckdb"
 578    DISPLAY_NAME: t.ClassVar[t.Literal["DuckDB"]] = "DuckDB"
 579    DISPLAY_ORDER: t.ClassVar[t.Literal[1]] = 1
 580
 581
 582class SnowflakeConnectionConfig(ConnectionConfig):
 583    """Configuration for the Snowflake connection.
 584
 585    Args:
 586        account: The Snowflake account name.
 587        user: The Snowflake username.
 588        password: The Snowflake password.
 589        warehouse: The optional warehouse name.
 590        database: The optional database name.
 591        role: The optional role name.
 592        concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
 593        authenticator: The optional authenticator name. Defaults to username/password authentication ("snowflake").
 594                       Options: https://github.com/snowflakedb/snowflake-connector-python/blob/e937591356c067a77f34a0a42328907fda792c23/src/snowflake/connector/network.py#L178-L183
 595        token: The optional oauth access token to use for authentication when authenticator is set to "oauth".
 596        private_key: The optional private key to use for authentication. Key can be Base64-encoded DER format (representing the key bytes), a plain-text PEM format, or bytes (Python config only). https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect#using-key-pair-authentication-key-pair-rotation
 597        private_key_path: The optional path to the private key to use for authentication. This would be used instead of `private_key`.
 598        private_key_passphrase: The optional passphrase to use to decrypt `private_key` or `private_key_path`. Keys can be created without encryption so only provide this if needed.
 599        register_comments: Whether or not to register model comments with the SQL engine.
 600        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
 601        session_parameters: The optional session parameters to set for the connection.
 602        host: Host address for the connection.
 603        port: Port for the connection.
 604    """
 605
 606    account: str
 607    user: t.Optional[str] = None
 608    password: t.Optional[str] = None
 609    warehouse: t.Optional[str] = None
 610    database: t.Optional[str] = None
 611    role: t.Optional[str] = None
 612    authenticator: t.Optional[str] = None
 613    token: t.Optional[str] = None
 614    host: t.Optional[str] = None
 615    port: t.Optional[int] = None
 616    application: t.Literal["Tobiko_SQLMesh"] = "Tobiko_SQLMesh"
 617
 618    # Private Key Auth
 619    private_key: t.Optional[t.Union[str, bytes]] = None
 620    private_key_path: t.Optional[str] = None
 621    private_key_passphrase: t.Optional[str] = None
 622
 623    concurrent_tasks: int = 4
 624    register_comments: bool = True
 625    pre_ping: bool = False
 626
 627    session_parameters: t.Optional[dict] = None
 628
 629    type_: t.Literal["snowflake"] = Field(alias="type", default="snowflake")
 630    DIALECT: t.ClassVar[t.Literal["snowflake"]] = "snowflake"
 631    DISPLAY_NAME: t.ClassVar[t.Literal["Snowflake"]] = "Snowflake"
 632    DISPLAY_ORDER: t.ClassVar[t.Literal[2]] = 2
 633
 634    _concurrent_tasks_validator = concurrent_tasks_validator
 635
 636    @model_validator(mode="before")
 637    def _validate_authenticator(cls, data: t.Any) -> t.Any:
 638        if not isinstance(data, dict):
 639            return data
 640
 641        from snowflake.connector.network import DEFAULT_AUTHENTICATOR, OAUTH_AUTHENTICATOR
 642
 643        auth = data.get("authenticator")
 644        auth = auth.upper() if auth else DEFAULT_AUTHENTICATOR
 645        user = data.get("user")
 646        password = data.get("password")
 647        data["private_key"] = cls._get_private_key(data, auth)  # type: ignore
 648
 649        if (
 650            auth == DEFAULT_AUTHENTICATOR
 651            and not data.get("private_key")
 652            and (not user or not password)
 653        ):
 654            raise ConfigError("User and password must be provided if using default authentication")
 655
 656        if auth == OAUTH_AUTHENTICATOR and not data.get("token"):
 657            raise ConfigError("Token must be provided if using oauth authentication")
 658
 659        return data
 660
 661    _engine_import_validator = _get_engine_import_validator(
 662        "snowflake.connector.network", "snowflake"
 663    )
 664
 665    @classmethod
 666    def _get_private_key(cls, values: t.Dict[str, t.Optional[str]], auth: str) -> t.Optional[bytes]:
 667        """
 668        source: https://github.com/dbt-labs/dbt-snowflake/blob/0374b4ec948982f2ac8ec0c95d53d672ad19e09c/dbt/adapters/snowflake/connections.py#L247C5-L285C1
 669
 670        Overall code change: Use local variables instead of class attributes + Validation
 671        """
 672        # Start custom code
 673        from cryptography.hazmat.backends import default_backend
 674        from cryptography.hazmat.primitives import serialization
 675        from snowflake.connector.network import (
 676            DEFAULT_AUTHENTICATOR,
 677            KEY_PAIR_AUTHENTICATOR,
 678        )
 679
 680        private_key = values.get("private_key")
 681        private_key_path = values.get("private_key_path")
 682        private_key_passphrase = values.get("private_key_passphrase")
 683        user = values.get("user")
 684        password = values.get("password")
 685        auth = auth if auth and auth != DEFAULT_AUTHENTICATOR else KEY_PAIR_AUTHENTICATOR
 686
 687        if not private_key and not private_key_path:
 688            return None
 689        if private_key and private_key_path:
 690            raise ConfigError("Cannot specify both `private_key` and `private_key_path`")
 691        if auth != KEY_PAIR_AUTHENTICATOR:
 692            raise ConfigError(
 693                f"Private key or private key path can only be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 694            )
 695        if not user:
 696            raise ConfigError(
 697                f"User must be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 698            )
 699        if password:
 700            raise ConfigError(
 701                f"Password cannot be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 702            )
 703
 704        if isinstance(private_key, bytes):
 705            return private_key
 706        # End Custom Code
 707
 708        if private_key_passphrase:
 709            encoded_passphrase = private_key_passphrase.encode()
 710        else:
 711            encoded_passphrase = None
 712
 713        if private_key:
 714            if private_key.startswith("-"):
 715                p_key = serialization.load_pem_private_key(
 716                    data=bytes(private_key, "utf-8"),
 717                    password=encoded_passphrase,
 718                    backend=default_backend(),
 719                )
 720
 721            else:
 722                p_key = serialization.load_der_private_key(
 723                    data=base64.b64decode(private_key),
 724                    password=encoded_passphrase,
 725                    backend=default_backend(),
 726                )
 727
 728        elif private_key_path:
 729            with open(private_key_path, "rb") as key:
 730                p_key = serialization.load_pem_private_key(
 731                    key.read(), password=encoded_passphrase, backend=default_backend()
 732                )
 733        else:
 734            return None
 735
 736        return p_key.private_bytes(
 737            encoding=serialization.Encoding.DER,
 738            format=serialization.PrivateFormat.PKCS8,
 739            encryption_algorithm=serialization.NoEncryption(),
 740        )
 741
 742    @property
 743    def _connection_kwargs_keys(self) -> t.Set[str]:
 744        return {
 745            "user",
 746            "password",
 747            "account",
 748            "warehouse",
 749            "database",
 750            "role",
 751            "authenticator",
 752            "token",
 753            "private_key",
 754            "session_parameters",
 755            "application",
 756            "host",
 757            "port",
 758        }
 759
 760    @property
 761    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 762        return engine_adapter.SnowflakeEngineAdapter
 763
 764    @property
 765    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 766        return {"autocommit": False}
 767
 768    @property
 769    def _connection_factory(self) -> t.Callable:
 770        from snowflake import connector
 771
 772        return connector.connect
 773
 774
 775class DatabricksConnectionConfig(ConnectionConfig):
 776    """
 777    Databricks connection that uses the SQL connector for SQL models and then Databricks Connect for Dataframe operations
 778
 779    Arg Source: https://github.com/databricks/databricks-sql-python/blob/main/src/databricks/sql/client.py#L39
 780    OAuth ref: https://docs.databricks.com/en/dev-tools/python-sql-connector.html#oauth-machine-to-machine-m2m-authentication
 781
 782    Args:
 783        server_hostname: Databricks instance host name.
 784        http_path: Http path either to a DBSQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef)
 785            or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)
 786        access_token: Http Bearer access token, e.g. Databricks Personal Access Token.
 787        auth_type: Set to 'databricks-oauth' or 'azure-oauth' to trigger OAuth (or dont set at all to use `access_token`)
 788        oauth_client_id: Client ID to use when auth_type is set to one of the 'oauth' types
 789        oauth_client_secret: Client Secret to use when auth_type is set to one of the 'oauth' types
 790        catalog: Default catalog to use for SQL models. Defaults to None which means it will use the default set in
 791            the Databricks cluster (most likely `hive_metastore`).
 792        http_headers: An optional list of (k, v) pairs that will be set as Http headers on every request
 793        session_configuration: An optional dictionary of Spark session parameters.
 794            Execute the SQL command `SET -v` to get a full list of available commands.
 795        databricks_connect_server_hostname: The hostname to use when establishing a connecting using Databricks Connect.
 796            Defaults to the `server_hostname` value.
 797        databricks_connect_access_token: The access token to use when establishing a connecting using Databricks Connect.
 798            Defaults to the `access_token` value.
 799        databricks_connect_cluster_id: The cluster id to use when establishing a connecting using Databricks Connect.
 800            Defaults to deriving the cluster id from the `http_path` value.
 801        force_databricks_connect: Force all queries to run using Databricks Connect instead of the SQL connector.
 802        disable_databricks_connect: Even if databricks connect is installed, do not use it.
 803        disable_spark_session: Do not use SparkSession if it is available (like when running in a notebook).
 804        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
 805    """
 806
 807    server_hostname: t.Optional[str] = None
 808    http_path: t.Optional[str] = None
 809    access_token: t.Optional[str] = None
 810    auth_type: t.Optional[str] = None
 811    oauth_client_id: t.Optional[str] = None
 812    oauth_client_secret: t.Optional[str] = None
 813    catalog: t.Optional[str] = None
 814    http_headers: t.Optional[t.List[t.Tuple[str, str]]] = None
 815    session_configuration: t.Optional[t.Dict[str, t.Any]] = None
 816    databricks_connect_server_hostname: t.Optional[str] = None
 817    databricks_connect_access_token: t.Optional[str] = None
 818    databricks_connect_cluster_id: t.Optional[str] = None
 819    databricks_connect_use_serverless: bool = False
 820    force_databricks_connect: bool = False
 821    disable_databricks_connect: bool = False
 822    disable_spark_session: bool = False
 823
 824    concurrent_tasks: int = 1
 825    register_comments: bool = True
 826    pre_ping: t.Literal[False] = False
 827
 828    type_: t.Literal["databricks"] = Field(alias="type", default="databricks")
 829    DIALECT: t.ClassVar[t.Literal["databricks"]] = "databricks"
 830    DISPLAY_NAME: t.ClassVar[t.Literal["Databricks"]] = "Databricks"
 831    DISPLAY_ORDER: t.ClassVar[t.Literal[3]] = 3
 832
 833    _concurrent_tasks_validator = concurrent_tasks_validator
 834    _http_headers_validator = http_headers_validator
 835
 836    @computed_field  # type: ignore[prop-decorator]
 837    @property
 838    def shared_connection(self) -> bool:
 839        """The connection should only be shared if U2M OAuth is being used"""
 840        return self.auth_type is not None and self.oauth_client_secret is None
 841
 842    @model_validator(mode="before")
 843    def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
 844        # SQLQueryContextLogger will output any error SQL queries even if they are in a try/except block.
 845        # Disabling this allows SQLMesh to determine what should be shown to the user.
 846        # Ex: We describe a table to see if it exists and therefore that execution can fail but we don't need to show
 847        # the user since it is expected if the table doesn't exist. Without this change the user would see the error.
 848        logging.getLogger("SQLQueryContextLogger").setLevel(logging.CRITICAL)
 849
 850        if not isinstance(data, dict):
 851            return data
 852
 853        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 854
 855        if DatabricksEngineAdapter.can_access_spark_session(
 856            bool(data.get("disable_spark_session"))
 857        ):
 858            return data
 859
 860        databricks_connect_use_serverless = data.get("databricks_connect_use_serverless")
 861        server_hostname, http_path, access_token, auth_type = (
 862            data.get("server_hostname"),
 863            data.get("http_path"),
 864            data.get("access_token"),
 865            data.get("auth_type"),
 866        )
 867
 868        if (not server_hostname or not http_path or not access_token) and (
 869            not databricks_connect_use_serverless and not auth_type
 870        ):
 871            raise ValueError(
 872                "`server_hostname`, `http_path`, and `access_token` are required for Databricks connections when not running in a notebook"
 873            )
 874        if (
 875            databricks_connect_use_serverless
 876            and not server_hostname
 877            and not data.get("databricks_connect_server_hostname")
 878        ):
 879            raise ValueError(
 880                "`server_hostname` or `databricks_connect_server_hostname` is required when `databricks_connect_use_serverless` is set"
 881            )
 882        if DatabricksEngineAdapter.can_access_databricks_connect(
 883            bool(data.get("disable_databricks_connect"))
 884        ):
 885            if not data.get("databricks_connect_access_token"):
 886                data["databricks_connect_access_token"] = access_token
 887            if not data.get("databricks_connect_server_hostname"):
 888                data["databricks_connect_server_hostname"] = f"https://{server_hostname}"
 889            if not databricks_connect_use_serverless and not data.get(
 890                "databricks_connect_cluster_id"
 891            ):
 892                if t.TYPE_CHECKING:
 893                    assert http_path is not None
 894                data["databricks_connect_cluster_id"] = http_path.split("/")[-1]
 895
 896        if auth_type:
 897            from databricks.sql.auth.auth import AuthType
 898
 899            all_data = [m.value for m in AuthType]
 900            if auth_type not in all_data:
 901                raise ValueError(
 902                    f"`auth_type` {auth_type} does not match a valid option: {all_data}"
 903                )
 904
 905            client_id = data.get("oauth_client_id")
 906            client_secret = data.get("oauth_client_secret")
 907
 908            if client_secret and not client_id:
 909                raise ValueError(
 910                    "`oauth_client_id` is required when `oauth_client_secret` is specified"
 911                )
 912
 913            if not http_path:
 914                raise ValueError("`http_path` is still required when using `auth_type`")
 915
 916        return data
 917
 918    _engine_import_validator = _get_engine_import_validator("databricks", "databricks")
 919
 920    @property
 921    def _connection_kwargs_keys(self) -> t.Set[str]:
 922        if self.use_spark_session_only:
 923            return set()
 924        return {
 925            "server_hostname",
 926            "http_path",
 927            "access_token",
 928            "http_headers",
 929            "session_configuration",
 930            "catalog",
 931        }
 932
 933    @property
 934    def _engine_adapter(self) -> t.Type[engine_adapter.DatabricksEngineAdapter]:
 935        return engine_adapter.DatabricksEngineAdapter
 936
 937    @property
 938    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 939        return {
 940            k: v
 941            for k, v in self.dict().items()
 942            if k.startswith("databricks_connect_")
 943            or k in ("catalog", "disable_databricks_connect", "disable_spark_session")
 944        }
 945
 946    @property
 947    def use_spark_session_only(self) -> bool:
 948        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 949
 950        return (
 951            DatabricksEngineAdapter.can_access_spark_session(self.disable_spark_session)
 952            or self.force_databricks_connect
 953        )
 954
 955    @property
 956    def _connection_factory(self) -> t.Callable:
 957        if self.use_spark_session_only:
 958            from sqlmesh.engines.spark.db_api.spark_session import connection
 959
 960            return connection
 961
 962        from databricks import sql  # type: ignore
 963
 964        return sql.connect
 965
 966    @property
 967    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 968        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 969
 970        if not self.use_spark_session_only:
 971            conn_kwargs: t.Dict[str, t.Any] = {
 972                "_user_agent_entry": "sqlmesh",
 973            }
 974
 975            if self.auth_type and "oauth" in self.auth_type:
 976                # there are two types of oauth: User-to-Machine (U2M) and Machine-to-Machine (M2M)
 977                if self.oauth_client_secret:
 978                    # if a client_secret exists, then a client_id also exists and we are using M2M
 979                    # ref: https://docs.databricks.com/en/dev-tools/python-sql-connector.html#oauth-machine-to-machine-m2m-authentication
 980                    # ref: https://github.com/databricks/databricks-sql-python/blob/main/examples/m2m_oauth.py
 981                    from databricks.sdk.core import Config, oauth_service_principal
 982
 983                    config = Config(
 984                        host=f"https://{self.server_hostname}",
 985                        client_id=self.oauth_client_id,
 986                        client_secret=self.oauth_client_secret,
 987                    )
 988                    conn_kwargs["credentials_provider"] = lambda: oauth_service_principal(config)
 989                else:
 990                    # if auth_type is set to an 'oauth' type but no client_id/secret are set, then we are using U2M
 991                    # ref: https://docs.databricks.com/en/dev-tools/python-sql-connector.html#oauth-user-to-machine-u2m-authentication
 992                    conn_kwargs["auth_type"] = self.auth_type
 993
 994            return conn_kwargs
 995
 996        if DatabricksEngineAdapter.can_access_spark_session(self.disable_spark_session):
 997            from pyspark.sql import SparkSession
 998
 999            return dict(
1000                spark=SparkSession.getActiveSession(),
1001                catalog=self.catalog,
1002            )
1003
1004        from databricks.connect import DatabricksSession
1005
1006        if t.TYPE_CHECKING:
1007            assert self.databricks_connect_server_hostname is not None
1008            assert self.databricks_connect_access_token is not None
1009
1010        if self.databricks_connect_use_serverless:
1011            builder = DatabricksSession.builder.remote(
1012                host=self.databricks_connect_server_hostname,
1013                token=self.databricks_connect_access_token,
1014                serverless=True,
1015            )
1016        else:
1017            if t.TYPE_CHECKING:
1018                assert self.databricks_connect_cluster_id is not None
1019            builder = DatabricksSession.builder.remote(
1020                host=self.databricks_connect_server_hostname,
1021                token=self.databricks_connect_access_token,
1022                cluster_id=self.databricks_connect_cluster_id,
1023            )
1024
1025        return dict(
1026            spark=builder.userAgent("sqlmesh").getOrCreate(),
1027            catalog=self.catalog,
1028        )
1029
1030
1031class BigQueryConnectionMethod(str, Enum):
1032    OAUTH = "oauth"
1033    OAUTH_SECRETS = "oauth-secrets"
1034    SERVICE_ACCOUNT = "service-account"
1035    SERVICE_ACCOUNT_JSON = "service-account-json"
1036
1037
1038class BigQueryPriority(str, Enum):
1039    BATCH = "batch"
1040    INTERACTIVE = "interactive"
1041
1042    @property
1043    def is_batch(self) -> bool:
1044        return self == self.BATCH
1045
1046    @property
1047    def is_interactive(self) -> bool:
1048        return self == self.INTERACTIVE
1049
1050    @property
1051    def bigquery_constant(self) -> str:
1052        from google.cloud.bigquery import QueryPriority
1053
1054        if self.is_batch:
1055            return QueryPriority.BATCH
1056        return QueryPriority.INTERACTIVE
1057
1058
1059class BigQueryConnectionConfig(ConnectionConfig):
1060    """
1061    BigQuery Connection Configuration.
1062    """
1063
1064    method: BigQueryConnectionMethod = BigQueryConnectionMethod.OAUTH
1065
1066    project: t.Optional[str] = None
1067    execution_project: t.Optional[str] = None
1068    quota_project: t.Optional[str] = None
1069    location: t.Optional[str] = None
1070    # Keyfile Auth
1071    keyfile: t.Optional[str] = None
1072    keyfile_json: t.Optional[t.Dict[str, t.Any]] = None
1073    # Oath Secret Auth
1074    token: t.Optional[str] = None
1075    refresh_token: t.Optional[str] = None
1076    client_id: t.Optional[str] = None
1077    client_secret: t.Optional[str] = None
1078    token_uri: t.Optional[str] = None
1079    scopes: t.Tuple[str, ...] = ("https://www.googleapis.com/auth/bigquery",)
1080    impersonated_service_account: t.Optional[str] = None
1081    # Extra Engine Config
1082    job_creation_timeout_seconds: t.Optional[int] = None
1083    job_execution_timeout_seconds: t.Optional[int] = None
1084    job_retries: t.Optional[int] = 1
1085    job_retry_deadline_seconds: t.Optional[int] = None
1086    priority: t.Optional[BigQueryPriority] = None
1087    maximum_bytes_billed: t.Optional[int] = None
1088    reservation: t.Optional[str] = None
1089
1090    concurrent_tasks: int = 1
1091    register_comments: bool = True
1092    pre_ping: t.Literal[False] = False
1093
1094    type_: t.Literal["bigquery"] = Field(alias="type", default="bigquery")
1095    DIALECT: t.ClassVar[t.Literal["bigquery"]] = "bigquery"
1096    DISPLAY_NAME: t.ClassVar[t.Literal["BigQuery"]] = "BigQuery"
1097    DISPLAY_ORDER: t.ClassVar[t.Literal[4]] = 4
1098
1099    _engine_import_validator = _get_engine_import_validator("google.cloud.bigquery", "bigquery")
1100
1101    @field_validator("execution_project")
1102    def validate_execution_project(
1103        cls,
1104        v: t.Optional[str],
1105        info: ValidationInfo,
1106    ) -> t.Optional[str]:
1107        if v and not validation_data(info).get("project"):
1108            raise ConfigError(
1109                "If the `execution_project` field is specified, you must also specify the `project` field to provide a default object location."
1110            )
1111        return v
1112
1113    @field_validator("quota_project")
1114    def validate_quota_project(
1115        cls,
1116        v: t.Optional[str],
1117        info: ValidationInfo,
1118    ) -> t.Optional[str]:
1119        if v and not validation_data(info).get("project"):
1120            raise ConfigError(
1121                "If the `quota_project` field is specified, you must also specify the `project` field to provide a default object location."
1122            )
1123        return v
1124
1125    @property
1126    def _connection_kwargs_keys(self) -> t.Set[str]:
1127        return set()
1128
1129    @property
1130    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1131        return engine_adapter.BigQueryEngineAdapter
1132
1133    @property
1134    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
1135        """The static connection kwargs for this connection"""
1136        import google.auth
1137        from google.api_core import client_info, client_options
1138        from google.auth import impersonated_credentials
1139        from google.oauth2 import credentials, service_account
1140
1141        if self.method == BigQueryConnectionMethod.OAUTH:
1142            creds, _ = google.auth.default(scopes=self.scopes)
1143        elif self.method == BigQueryConnectionMethod.SERVICE_ACCOUNT:
1144            creds = service_account.Credentials.from_service_account_file(
1145                self.keyfile, scopes=self.scopes
1146            )
1147        elif self.method == BigQueryConnectionMethod.SERVICE_ACCOUNT_JSON:
1148            creds = service_account.Credentials.from_service_account_info(
1149                self.keyfile_json, scopes=self.scopes
1150            )
1151        elif self.method == BigQueryConnectionMethod.OAUTH_SECRETS:
1152            creds = credentials.Credentials(
1153                token=self.token,
1154                refresh_token=self.refresh_token,
1155                client_id=self.client_id,
1156                client_secret=self.client_secret,
1157                token_uri=self.token_uri,
1158                scopes=self.scopes,
1159            )
1160        else:
1161            raise ConfigError("Invalid BigQuery Connection Method")
1162
1163        if self.impersonated_service_account:
1164            creds = impersonated_credentials.Credentials(
1165                source_credentials=creds,
1166                target_principal=self.impersonated_service_account,
1167                target_scopes=self.scopes,
1168            )
1169
1170        options = client_options.ClientOptions(quota_project_id=self.quota_project)
1171        project = self.execution_project or self.project or None
1172
1173        client = google.cloud.bigquery.Client(
1174            project=project and exp.parse_identifier(project, dialect="bigquery").name,
1175            credentials=creds,
1176            location=self.location,
1177            client_info=client_info.ClientInfo(user_agent="sqlmesh"),
1178            client_options=options,
1179        )
1180
1181        return {
1182            "client": client,
1183        }
1184
1185    @property
1186    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1187        return {
1188            k: v
1189            for k, v in self.dict().items()
1190            if k
1191            in {
1192                "job_creation_timeout_seconds",
1193                "job_execution_timeout_seconds",
1194                "job_retries",
1195                "job_retry_deadline_seconds",
1196                "priority",
1197                "maximum_bytes_billed",
1198                "reservation",
1199            }
1200        }
1201
1202    @property
1203    def _connection_factory(self) -> t.Callable:
1204        from google.cloud.bigquery.dbapi import connect
1205
1206        return connect
1207
1208    def get_catalog(self) -> t.Optional[str]:
1209        return self.project
1210
1211
1212class GCPPostgresConnectionConfig(ConnectionConfig):
1213    """
1214    Postgres Connection Configuration for GCP.
1215
1216    Args:
1217        instance_connection_string: Connection name for the postgres instance.
1218        user: Postgres or IAM user's name
1219        password: The postgres user's password. Only needed when the user is a postgres user.
1220        enable_iam_auth: Set to True when user is an IAM user.
1221        db: Name of the db to connect to.
1222        keyfile: string path to json service account credentials file
1223        keyfile_json: dict service account credentials info
1224        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
1225    """
1226
1227    instance_connection_string: str
1228    user: str
1229    password: t.Optional[str] = None
1230    enable_iam_auth: t.Optional[bool] = None
1231    db: str
1232    ip_type: t.Union[t.Literal["public"], t.Literal["private"], t.Literal["psc"]] = "public"
1233    # Keyfile Auth
1234    keyfile: t.Optional[str] = None
1235    keyfile_json: t.Optional[t.Dict[str, t.Any]] = None
1236    timeout: t.Optional[int] = None
1237    scopes: t.Tuple[str, ...] = ("https://www.googleapis.com/auth/sqlservice.admin",)
1238    driver: str = "pg8000"
1239
1240    type_: t.Literal["gcp_postgres"] = Field(alias="type", default="gcp_postgres")
1241    DIALECT: t.ClassVar[t.Literal["postgres"]] = "postgres"
1242    DISPLAY_NAME: t.ClassVar[t.Literal["GCP Postgres"]] = "GCP Postgres"
1243    DISPLAY_ORDER: t.ClassVar[t.Literal[13]] = 13
1244
1245    concurrent_tasks: int = 4
1246    register_comments: bool = True
1247    pre_ping: bool = True
1248
1249    _engine_import_validator = _get_engine_import_validator(
1250        "google.cloud.sql", "gcp_postgres", "gcppostgres"
1251    )
1252
1253    @model_validator(mode="before")
1254    def _validate_auth_method(cls, data: t.Any) -> t.Any:
1255        if not isinstance(data, dict):
1256            return data
1257
1258        password = data.get("password")
1259        enable_iam_auth = data.get("enable_iam_auth")
1260
1261        if not password and not enable_iam_auth:
1262            raise ConfigError(
1263                "GCP Postgres connection configuration requires either password set"
1264                " for a postgres user account or enable_iam_auth set to 'True'"
1265                " for an IAM user account."
1266            )
1267
1268        return data
1269
1270    @property
1271    def _connection_kwargs_keys(self) -> t.Set[str]:
1272        return {
1273            "instance_connection_string",
1274            "driver",
1275            "user",
1276            "password",
1277            "db",
1278            "enable_iam_auth",
1279            "timeout",
1280        }
1281
1282    @property
1283    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1284        return engine_adapter.PostgresEngineAdapter
1285
1286    @property
1287    def _connection_factory(self) -> t.Callable:
1288        from google.cloud.sql.connector import Connector
1289        from google.oauth2 import service_account
1290
1291        creds = None
1292        if self.keyfile:
1293            creds = service_account.Credentials.from_service_account_file(
1294                self.keyfile, scopes=self.scopes
1295            )
1296        elif self.keyfile_json:
1297            creds = service_account.Credentials.from_service_account_info(
1298                self.keyfile_json, scopes=self.scopes
1299            )
1300
1301        kwargs = {
1302            "credentials": creds,
1303            "ip_type": self.ip_type,
1304        }
1305
1306        if self.timeout:
1307            kwargs["timeout"] = self.timeout
1308
1309        return Connector(**kwargs).connect  # type: ignore
1310
1311
1312class RedshiftConnectionConfig(ConnectionConfig):
1313    """
1314    Redshift Connection Configuration.
1315
1316    Arg Source: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/__init__.py#L146
1317    Note: A subset of properties were selected. Please open an issue/PR if you want to see more supported.
1318
1319    Args:
1320        user: The username to use for authentication with the Amazon Redshift cluster.
1321        password: The password to use for authentication with the Amazon Redshift cluster.
1322        database: The name of the database instance to connect to.
1323        host: The hostname of the Amazon Redshift cluster.
1324        port: The port number of the Amazon Redshift cluster. Default value is 5439.
1325        source_address: No description provided
1326        unix_sock: No description provided
1327        ssl: Is SSL enabled. Default value is ``True``. SSL must be enabled when authenticating using IAM.
1328        sslmode: The security of the connection to the Amazon Redshift cluster. 'verify-ca' and 'verify-full' are supported.
1329        timeout: The number of seconds before the connection to the server will timeout. By default there is no timeout.
1330        tcp_keepalive: Is `TCP keepalive <https://en.wikipedia.org/wiki/Keepalive#TCP_keepalive>`_ used. The default value is ``True``.
1331        application_name: Sets the application name. The default value is None.
1332        preferred_role: The IAM role preferred for the current connection.
1333        principal_arn: The ARN of the IAM entity (user or role) for which you are generating a policy.
1334        credentials_provider: The class name of the IdP that will be used for authenticating with the Amazon Redshift cluster.
1335        region: The AWS region where the Amazon Redshift cluster is located.
1336        cluster_identifier: The cluster identifier of the Amazon Redshift cluster.
1337        iam: If IAM authentication is enabled. Default value is False. IAM must be True when authenticating using an IdP.
1338        db_user: The database user to authenticate as. Required when using IAM authentication.
1339        is_serverless: Redshift end-point is serverless or provisional. Default value false.
1340        serverless_acct_id: The account ID of the serverless. Default value None
1341        serverless_work_group: The name of work group for serverless end point. Default value None.
1342        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
1343        enable_merge: Whether to use the Redshift merge operation instead of the SQLMesh logical merge.
1344    """
1345
1346    user: t.Optional[str] = None
1347    password: t.Optional[str] = None
1348    database: t.Optional[str] = None
1349    host: t.Optional[str] = None
1350    port: t.Optional[int] = None
1351    source_address: t.Optional[str] = None
1352    unix_sock: t.Optional[str] = None
1353    ssl: t.Optional[bool] = None
1354    sslmode: t.Optional[str] = None
1355    timeout: t.Optional[int] = None
1356    tcp_keepalive: t.Optional[bool] = None
1357    application_name: t.Optional[str] = None
1358    preferred_role: t.Optional[str] = None
1359    principal_arn: t.Optional[str] = None
1360    credentials_provider: t.Optional[str] = None
1361    region: t.Optional[str] = None
1362    cluster_identifier: t.Optional[str] = None
1363    iam: t.Optional[bool] = None
1364    db_user: t.Optional[str] = None
1365    is_serverless: t.Optional[bool] = None
1366    serverless_acct_id: t.Optional[str] = None
1367    serverless_work_group: t.Optional[str] = None
1368    enable_merge: t.Optional[bool] = None
1369
1370    concurrent_tasks: int = 4
1371    register_comments: bool = True
1372    pre_ping: bool = False
1373
1374    type_: t.Literal["redshift"] = Field(alias="type", default="redshift")
1375    DIALECT: t.ClassVar[t.Literal["redshift"]] = "redshift"
1376    DISPLAY_NAME: t.ClassVar[t.Literal["Redshift"]] = "Redshift"
1377    DISPLAY_ORDER: t.ClassVar[t.Literal[7]] = 7
1378
1379    _engine_import_validator = _get_engine_import_validator("redshift_connector", "redshift")
1380
1381    @property
1382    def _connection_kwargs_keys(self) -> t.Set[str]:
1383        return {
1384            "user",
1385            "password",
1386            "database",
1387            "host",
1388            "port",
1389            "source_address",
1390            "unix_sock",
1391            "ssl",
1392            "sslmode",
1393            "timeout",
1394            "tcp_keepalive",
1395            "application_name",
1396            "preferred_role",
1397            "principal_arn",
1398            "credentials_provider",
1399            "region",
1400            "cluster_identifier",
1401            "iam",
1402            "db_user",
1403            "is_serverless",
1404            "serverless_acct_id",
1405            "serverless_work_group",
1406        }
1407
1408    @property
1409    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1410        return engine_adapter.RedshiftEngineAdapter
1411
1412    @property
1413    def _connection_factory(self) -> t.Callable:
1414        from redshift_connector import connect
1415
1416        return connect
1417
1418    @property
1419    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1420        return {"enable_merge": self.enable_merge}
1421
1422
1423class PostgresConnectionConfig(ConnectionConfig):
1424    host: str
1425    user: str
1426    password: str
1427    port: int
1428    database: str
1429    keepalives_idle: t.Optional[int] = None
1430    connect_timeout: int = 10
1431    role: t.Optional[str] = None
1432    sslmode: t.Optional[str] = None
1433    application_name: t.Optional[str] = None
1434
1435    concurrent_tasks: int = 4
1436    register_comments: bool = True
1437    pre_ping: bool = True
1438
1439    type_: t.Literal["postgres"] = Field(alias="type", default="postgres")
1440    DIALECT: t.ClassVar[t.Literal["postgres"]] = "postgres"
1441    DISPLAY_NAME: t.ClassVar[t.Literal["Postgres"]] = "Postgres"
1442    DISPLAY_ORDER: t.ClassVar[t.Literal[12]] = 12
1443
1444    _engine_import_validator = _get_engine_import_validator("psycopg2", "postgres")
1445
1446    @property
1447    def _connection_kwargs_keys(self) -> t.Set[str]:
1448        return {
1449            "host",
1450            "user",
1451            "password",
1452            "port",
1453            "database",
1454            "keepalives_idle",
1455            "connect_timeout",
1456            "sslmode",
1457            "application_name",
1458        }
1459
1460    @property
1461    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1462        return engine_adapter.PostgresEngineAdapter
1463
1464    @property
1465    def _connection_factory(self) -> t.Callable:
1466        from psycopg2 import connect
1467
1468        return connect
1469
1470    @property
1471    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
1472        if not self.role:
1473            return None
1474
1475        def init(cursor: t.Any) -> None:
1476            cursor.execute(f"SET ROLE {self.role}")
1477
1478        return init
1479
1480
1481class MySQLConnectionConfig(ConnectionConfig):
1482    host: str
1483    user: str
1484    password: str
1485    port: t.Optional[int] = None
1486    database: t.Optional[str] = None
1487    charset: t.Optional[str] = None
1488    collation: t.Optional[str] = None
1489    ssl_disabled: t.Optional[bool] = None
1490
1491    concurrent_tasks: int = 4
1492    register_comments: bool = True
1493    pre_ping: bool = True
1494
1495    type_: t.Literal["mysql"] = Field(alias="type", default="mysql")
1496    DIALECT: t.ClassVar[t.Literal["mysql"]] = "mysql"
1497    DISPLAY_NAME: t.ClassVar[t.Literal["MySQL"]] = "MySQL"
1498    DISPLAY_ORDER: t.ClassVar[t.Literal[14]] = 14
1499
1500    _engine_import_validator = _get_engine_import_validator("pymysql", "mysql")
1501
1502    @property
1503    def _connection_kwargs_keys(self) -> t.Set[str]:
1504        connection_keys = {
1505            "host",
1506            "user",
1507            "password",
1508        }
1509        if self.port is not None:
1510            connection_keys.add("port")
1511        if self.database is not None:
1512            connection_keys.add("database")
1513        if self.charset is not None:
1514            connection_keys.add("charset")
1515        if self.collation is not None:
1516            connection_keys.add("collation")
1517        if self.ssl_disabled is not None:
1518            connection_keys.add("ssl_disabled")
1519        return connection_keys
1520
1521    @property
1522    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1523        return engine_adapter.MySQLEngineAdapter
1524
1525    @property
1526    def _connection_factory(self) -> t.Callable:
1527        from pymysql import connect
1528
1529        return connect
1530
1531
1532class MSSQLConnectionConfig(ConnectionConfig):
1533    host: str
1534    user: t.Optional[str] = None
1535    password: t.Optional[str] = None
1536    database: t.Optional[str] = ""
1537    timeout: t.Optional[int] = 0
1538    login_timeout: t.Optional[int] = 60
1539    login_attempts: t.Optional[int] = 1
1540    charset: t.Optional[str] = "UTF-8"
1541    appname: t.Optional[str] = None
1542    port: t.Optional[int] = 1433
1543    conn_properties: t.Optional[t.Union[t.List[str], str]] = None
1544    autocommit: t.Optional[bool] = False
1545    tds_version: t.Optional[str] = None
1546
1547    # Driver options
1548    driver: t.Literal["pymssql", "pyodbc", "mssql-python"] = "pymssql"
1549    # PyODBC specific options
1550    driver_name: t.Optional[str] = None  # e.g. "ODBC Driver 18 for SQL Server"
1551    trust_server_certificate: t.Optional[bool] = None
1552    encrypt: t.Optional[bool] = None
1553    # Dictionary of arbitrary ODBC connection properties
1554    # See: https://learn.microsoft.com/en-us/sql/connect/odbc/dsn-connection-string-attribute
1555    odbc_properties: t.Optional[t.Dict[str, t.Any]] = None
1556
1557    concurrent_tasks: int = 4
1558    register_comments: bool = True
1559    pre_ping: bool = True
1560
1561    type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
1562    DIALECT: t.ClassVar[t.Literal["tsql"]] = "tsql"
1563    DISPLAY_NAME: t.ClassVar[t.Literal["MSSQL"]] = "MSSQL"
1564    DISPLAY_ORDER: t.ClassVar[t.Literal[11]] = 11
1565
1566    @model_validator(mode="before")
1567    @classmethod
1568    def _mssql_engine_import_validator(cls, data: t.Any) -> t.Any:
1569        if not isinstance(data, dict):
1570            return data
1571
1572        driver = data.get("driver", "pymssql")
1573
1574        # Define the mapping of driver to import module and extra name
1575        driver_configs = {
1576            "pymssql": ("pymssql", "mssql"),
1577            "pyodbc": ("pyodbc", "mssql-odbc"),
1578            "mssql-python": ("mssql_python", "mssql-python"),
1579        }
1580
1581        if driver not in driver_configs:
1582            raise ValueError(f"Unsupported driver: {driver}")
1583
1584        import_module, extra_name = driver_configs[driver]
1585
1586        # Use _get_engine_import_validator with decorate=False to get the raw validation function
1587        # This avoids the __wrapped__ issue in Python 3.9
1588        validator_func = _get_engine_import_validator(
1589            import_module, driver, extra_name, decorate=False
1590        )
1591
1592        # Call the raw validation function directly
1593        return validator_func(cls, data)
1594
1595    @property
1596    def _connection_kwargs_keys(self) -> t.Set[str]:
1597        base_keys = {
1598            "host",
1599            "user",
1600            "password",
1601            "database",
1602            "timeout",
1603            "login_timeout",
1604            "charset",
1605            "appname",
1606            "port",
1607            "conn_properties",
1608            "autocommit",
1609            "tds_version",
1610        }
1611
1612        if self.driver == "pyodbc":
1613            base_keys.update(
1614                {
1615                    "driver_name",
1616                    "trust_server_certificate",
1617                    "encrypt",
1618                    "odbc_properties",
1619                }
1620            )
1621            # Remove pymssql-specific parameters
1622            base_keys.discard("tds_version")
1623            base_keys.discard("conn_properties")
1624
1625        elif self.driver == "mssql-python":
1626            base_keys.update(
1627                {
1628                    "trust_server_certificate",
1629                    "encrypt",
1630                    "odbc_properties",
1631                    "login_attempts",
1632                }
1633            )
1634            # Remove pymssql-specific parameters
1635            base_keys.discard("tds_version")
1636            base_keys.discard("conn_properties")
1637
1638        return base_keys
1639
1640    @property
1641    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1642        return engine_adapter.MSSQLEngineAdapter
1643
1644    @property
1645    def _connection_factory(self) -> t.Callable:
1646        if self.driver == "pymssql":
1647            import pymssql
1648
1649            return pymssql.connect
1650
1651        if self.driver == "mssql-python":
1652            # The `mssql-python` implementation is API-compatible with
1653            # with the `pyodbc` equivalent for documented parameters.
1654
1655            if not SUPPORTS_MSSQL_PYTHON_DRIVER:
1656                raise ConfigError("The `mssql-python` driver requires Python 3.10 or higher.")
1657
1658            import mssql_python
1659
1660            def connect_mssql_python(**kwargs: t.Any) -> t.Callable:
1661                # Extract parameters for connection string
1662                host = kwargs.pop("host")
1663                port = kwargs.pop("port", 1433)
1664                database = kwargs.pop("database", "")
1665                user = kwargs.pop("user", None)
1666                password = kwargs.pop("password", None)
1667                authentication = kwargs.pop("authentication", None)
1668                trust_server_certificate = kwargs.pop("trust_server_certificate", False)
1669                encrypt = kwargs.pop("encrypt", True)
1670                timeout = kwargs.pop("timeout", 0)
1671                login_timeout = kwargs.pop("login_timeout", 59)
1672                login_attempts = kwargs.pop("login_attempts", 1)
1673
1674                # Build connection string
1675                conn_str_parts = [
1676                    f"Server={host},{port}",
1677                ]
1678
1679                if database:
1680                    conn_str_parts.append(f"Database={database}")
1681
1682                # Add security options
1683                conn_str_parts.append(f"Encrypt={'yes' if encrypt else 'no'}")
1684                if trust_server_certificate:
1685                    conn_str_parts.append("TrustServerCertificate=yes")
1686
1687                # `Connection Timeout=` is not a valid option so we leverage `ConnectRetry*`.
1688                # See the following:
1689                # - https://github.com/microsoft/mssql-python/issues/339
1690                # - https://github.com/microsoft/mssql-python/wiki/Connection-to-SQL-Database
1691                # - https://github.com/microsoft/mssql-python/wiki/Connection#timeout
1692                conn_str_parts.append(f"ConnectRetryCount={login_attempts}")
1693                conn_str_parts.append(f"ConnectRetryInterval={min(int(login_timeout), 60)}")
1694
1695                # Standard SQL Server authentication
1696                if user:
1697                    conn_str_parts.append(f"UID={user}")
1698                if password:
1699                    conn_str_parts.append(f"PWD={password}")
1700                if authentication:
1701                    conn_str_parts.append(f"Authentication={authentication}")
1702
1703                # Add any additional ODBC properties from the odbc_properties dictionary
1704                if self.odbc_properties:
1705                    for key, value in self.odbc_properties.items():
1706                        # Skip properties that we've already set above
1707                        if key.lower() in (
1708                            "driver",
1709                            "server",
1710                            "database",
1711                            "uid",
1712                            "pwd",
1713                            "encrypt",
1714                            "trustservercertificate",
1715                            "connectretrycount",
1716                            "connectretryinterval",
1717                            "connection timeout",
1718                        ):
1719                            continue
1720
1721                        # Handle boolean values properly
1722                        if isinstance(value, bool):
1723                            conn_str_parts.append(f"{key}={'yes' if value else 'no'}")
1724                        else:
1725                            conn_str_parts.append(f"{key}={value}")
1726
1727                # Create the connection
1728                conn_str = ";".join(conn_str_parts)
1729
1730                conn = mssql_python.connect(
1731                    conn_str,
1732                    autocommit=kwargs.get("autocommit", False),
1733                    timeout=timeout,
1734                )
1735
1736                # TODO: Remove this output converter as DATETIMEOFFSET
1737                # should be handled natively by `mssql-python`.
1738                # see "https://github.com/microsoft/mssql-python/issues/213"
1739
1740                def handle_datetimeoffset_mssql_python(dto_value: t.Any) -> t.Any:
1741                    import struct
1742                    from datetime import datetime, timedelta, timezone
1743
1744                    # Unpack the DATETIMEOFFSET binary format:
1745                    # Format: <6hI2h = (year, month, day, hour, minute, second, nanoseconds, tz_hour_offset, tz_minute_offset)
1746                    tup = struct.unpack("<6hI2h", dto_value)
1747                    return datetime(
1748                        tup[0],
1749                        tup[1],
1750                        tup[2],
1751                        tup[3],
1752                        tup[4],
1753                        tup[5],
1754                        tup[6] // 1000,
1755                        timezone(timedelta(hours=tup[7], minutes=tup[8])),
1756                    )
1757
1758                conn.add_output_converter(-155, handle_datetimeoffset_mssql_python)
1759
1760                return t.cast(t.Any, conn)
1761
1762            return connect_mssql_python
1763
1764        if self.driver == "pyodbc":
1765
1766            def connect_pyodbc(**kwargs: t.Any) -> t.Callable:
1767                # Extract parameters for connection string
1768                host = kwargs.pop("host")
1769                port = kwargs.pop("port", 1433)
1770                database = kwargs.pop("database", "")
1771                user = kwargs.pop("user", None)
1772                password = kwargs.pop("password", None)
1773                driver_name = kwargs.pop("driver_name", "ODBC Driver 18 for SQL Server")
1774                trust_server_certificate = kwargs.pop("trust_server_certificate", False)
1775                encrypt = kwargs.pop("encrypt", True)
1776                login_timeout = kwargs.pop("login_timeout", 60)
1777
1778                # Build connection string
1779                conn_str_parts = [
1780                    f"DRIVER={{{driver_name}}}",
1781                    f"SERVER={host},{port}",
1782                ]
1783
1784                if database:
1785                    conn_str_parts.append(f"DATABASE={database}")
1786
1787                # Add security options
1788                conn_str_parts.append(f"Encrypt={'YES' if encrypt else 'NO'}")
1789                if trust_server_certificate:
1790                    conn_str_parts.append("TrustServerCertificate=YES")
1791
1792                conn_str_parts.append(f"Connection Timeout={login_timeout}")
1793
1794                # Standard SQL Server authentication
1795                if user:
1796                    conn_str_parts.append(f"UID={user}")
1797                if password:
1798                    conn_str_parts.append(f"PWD={password}")
1799
1800                # Add any additional ODBC properties from the odbc_properties dictionary
1801                if self.odbc_properties:
1802                    for key, value in self.odbc_properties.items():
1803                        # Skip properties that we've already set above
1804                        if key.lower() in (
1805                            "driver",
1806                            "server",
1807                            "database",
1808                            "uid",
1809                            "pwd",
1810                            "encrypt",
1811                            "trustservercertificate",
1812                            "connection timeout",
1813                        ):
1814                            continue
1815
1816                        # Handle boolean values properly
1817                        if isinstance(value, bool):
1818                            conn_str_parts.append(f"{key}={'YES' if value else 'NO'}")
1819                        else:
1820                            conn_str_parts.append(f"{key}={value}")
1821
1822                # Create the connection
1823                conn_str = ";".join(conn_str_parts)
1824
1825                import pyodbc
1826
1827                conn = pyodbc.connect(conn_str, autocommit=kwargs.get("autocommit", False))
1828
1829                # Set up output converters for MSSQL-specific data types
1830                # Handle SQL type -155 (DATETIMEOFFSET) which is not yet supported by pyodbc
1831                # ref: https://github.com/mkleehammer/pyodbc/issues/134#issuecomment-281739794
1832                def handle_datetimeoffset_pyodbc(dto_value: t.Any) -> t.Any:
1833                    import struct
1834                    from datetime import datetime, timedelta, timezone
1835
1836                    # Unpack the DATETIMEOFFSET binary format:
1837                    # Format: <6hI2h = (year, month, day, hour, minute, second, nanoseconds, tz_hour_offset, tz_minute_offset)
1838                    tup = struct.unpack("<6hI2h", dto_value)
1839                    return datetime(
1840                        tup[0],
1841                        tup[1],
1842                        tup[2],
1843                        tup[3],
1844                        tup[4],
1845                        tup[5],
1846                        tup[6] // 1000,
1847                        timezone(timedelta(hours=tup[7], minutes=tup[8])),
1848                    )
1849
1850                conn.add_output_converter(-155, handle_datetimeoffset_pyodbc)
1851
1852                return t.cast(t.Any, conn)
1853
1854            return connect_pyodbc
1855
1856        raise ValueError(f"Unsupported driver: {self.driver}")
1857
1858    @property
1859    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1860        return {"catalog_support": CatalogSupport.REQUIRES_SET_CATALOG}
1861
1862
1863class AzureSQLConnectionConfig(MSSQLConnectionConfig):
1864    type_: t.Literal["azuresql"] = Field(alias="type", default="azuresql")  # type: ignore
1865    DISPLAY_NAME: t.ClassVar[t.Literal["Azure SQL"]] = "Azure SQL"  # type: ignore
1866    DISPLAY_ORDER: t.ClassVar[t.Literal[10]] = 10  # type: ignore
1867
1868    @property
1869    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1870        return {"catalog_support": CatalogSupport.SINGLE_CATALOG_ONLY}
1871
1872
1873class FabricConnectionConfig(MSSQLConnectionConfig):
1874    """
1875    Fabric Connection Configuration.
1876    Inherits most settings from MSSQLConnectionConfig and sets the type to 'fabric'.
1877    It is recommended to use the 'pyodbc' driver for Fabric.
1878    """
1879
1880    type_: t.Literal["fabric"] = Field(alias="type", default="fabric")  # type: ignore
1881    DIALECT: t.ClassVar[t.Literal["fabric"]] = "fabric"  # type: ignore
1882    DISPLAY_NAME: t.ClassVar[t.Literal["Fabric"]] = "Fabric"  # type: ignore
1883    DISPLAY_ORDER: t.ClassVar[t.Literal[17]] = 17  # type: ignore
1884    driver: t.Literal["pyodbc", "mssql-python"] = "pyodbc"
1885    workspace_id: str
1886    tenant_id: str
1887    autocommit: t.Optional[bool] = True
1888
1889    @property
1890    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1891        from sqlmesh.core.engine_adapter.fabric import FabricEngineAdapter
1892
1893        return FabricEngineAdapter
1894
1895    @property
1896    def _connection_factory(self) -> t.Callable:
1897        # Override to support catalog switching for Fabric
1898        base_factory = super()._connection_factory
1899
1900        def create_fabric_connection(
1901            target_catalog: t.Optional[str] = None, *args: t.Any, **kwargs: t.Any
1902        ) -> t.Callable:
1903            kwargs["database"] = target_catalog or self.database
1904            return base_factory(*args, **kwargs)
1905
1906        return create_fabric_connection
1907
1908    @property
1909    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1910        return {
1911            "database": self.database,
1912            # more operations than not require a specific catalog to be already active
1913            # in particular, create/drop view, create/drop schema and querying information_schema
1914            "catalog_support": CatalogSupport.REQUIRES_SET_CATALOG,
1915            "workspace_id": self.workspace_id,
1916            "tenant_id": self.tenant_id,
1917            "user": self.user,
1918            "password": self.password,
1919        }
1920
1921
1922class SparkConnectionConfig(ConnectionConfig):
1923    """
1924    Vanilla Spark Connection Configuration. Use `DatabricksConnectionConfig` for Databricks.
1925    """
1926
1927    config_dir: t.Optional[str] = None
1928    catalog: t.Optional[str] = None
1929    config: t.Dict[str, t.Any] = {}
1930    wap_enabled: bool = False
1931
1932    concurrent_tasks: int = 4
1933    register_comments: bool = True
1934    pre_ping: t.Literal[False] = False
1935
1936    type_: t.Literal["spark"] = Field(alias="type", default="spark")
1937    DIALECT: t.ClassVar[t.Literal["spark"]] = "spark"
1938    DISPLAY_NAME: t.ClassVar[t.Literal["Spark"]] = "Spark"
1939    DISPLAY_ORDER: t.ClassVar[t.Literal[8]] = 8
1940
1941    _engine_import_validator = _get_engine_import_validator("pyspark", "spark")
1942
1943    @property
1944    def _connection_kwargs_keys(self) -> t.Set[str]:
1945        return {
1946            "catalog",
1947        }
1948
1949    @property
1950    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1951        return engine_adapter.SparkEngineAdapter
1952
1953    @property
1954    def _connection_factory(self) -> t.Callable:
1955        from sqlmesh.engines.spark.db_api.spark_session import connection
1956
1957        return connection
1958
1959    @property
1960    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
1961        from pyspark.conf import SparkConf
1962        from pyspark.sql import SparkSession
1963
1964        spark_config = SparkConf()
1965        if self.config:
1966            for k, v in self.config.items():
1967                spark_config.set(k, v)
1968
1969        if self.config_dir:
1970            os.environ["SPARK_CONF_DIR"] = self.config_dir
1971        return {
1972            "spark": SparkSession.builder.config(conf=spark_config)
1973            .enableHiveSupport()
1974            .getOrCreate(),
1975        }
1976
1977    @property
1978    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
1979        return {"wap_enabled": self.wap_enabled}
1980
1981
1982class TrinoAuthenticationMethod(str, Enum):
1983    NO_AUTH = "no-auth"
1984    BASIC = "basic"
1985    LDAP = "ldap"
1986    KERBEROS = "kerberos"
1987    JWT = "jwt"
1988    CERTIFICATE = "certificate"
1989    OAUTH = "oauth"
1990
1991    @property
1992    def is_no_auth(self) -> bool:
1993        return self == self.NO_AUTH
1994
1995    @property
1996    def is_basic(self) -> bool:
1997        return self == self.BASIC
1998
1999    @property
2000    def is_ldap(self) -> bool:
2001        return self == self.LDAP
2002
2003    @property
2004    def is_kerberos(self) -> bool:
2005        return self == self.KERBEROS
2006
2007    @property
2008    def is_jwt(self) -> bool:
2009        return self == self.JWT
2010
2011    @property
2012    def is_certificate(self) -> bool:
2013        return self == self.CERTIFICATE
2014
2015    @property
2016    def is_oauth(self) -> bool:
2017        return self == self.OAUTH
2018
2019
2020class TrinoConnectionConfig(ConnectionConfig):
2021    method: TrinoAuthenticationMethod = TrinoAuthenticationMethod.NO_AUTH
2022    host: str
2023    user: str
2024    catalog: str
2025    port: t.Optional[int] = None
2026    http_scheme: t.Literal["http", "https"] = "https"
2027    # General Optional
2028    roles: t.Optional[t.Dict[str, str]] = None
2029    http_headers: t.Optional[t.Dict[str, str]] = None
2030    session_properties: t.Optional[t.Dict[str, str]] = None
2031    retries: int = 3
2032    timezone: t.Optional[str] = None
2033    # Basic/LDAP
2034    password: t.Optional[str] = None
2035    verify: t.Optional[bool] = None  # disable SSL verification (ignored if `cert` is provided)
2036    # LDAP
2037    impersonation_user: t.Optional[str] = None
2038    # Kerberos
2039    keytab: t.Optional[str] = None
2040    krb5_config: t.Optional[str] = None
2041    principal: t.Optional[str] = None
2042    service_name: str = "trino"
2043    hostname_override: t.Optional[str] = None
2044    mutual_authentication: bool = False
2045    force_preemptive: bool = False
2046    sanitize_mutual_error_response: bool = True
2047    delegate: bool = False
2048    # JWT
2049    jwt_token: t.Optional[str] = None
2050    # Certificate
2051    client_certificate: t.Optional[str] = None
2052    client_private_key: t.Optional[str] = None
2053    cert: t.Optional[str] = None
2054    source: str = "sqlmesh"
2055
2056    # SQLMesh options
2057    schema_location_mapping: t.Optional[dict[re.Pattern, str]] = None
2058    timestamp_mapping: t.Optional[dict[exp.DataType, exp.DataType]] = None
2059    concurrent_tasks: int = 4
2060    register_comments: bool = True
2061    pre_ping: t.Literal[False] = False
2062
2063    type_: t.Literal["trino"] = Field(alias="type", default="trino")
2064    DIALECT: t.ClassVar[t.Literal["trino"]] = "trino"
2065    DISPLAY_NAME: t.ClassVar[t.Literal["Trino"]] = "Trino"
2066    DISPLAY_ORDER: t.ClassVar[t.Literal[9]] = 9
2067
2068    _engine_import_validator = _get_engine_import_validator("trino", "trino")
2069
2070    @field_validator("schema_location_mapping", mode="before")
2071    @classmethod
2072    def _validate_regex_keys(
2073        cls, value: t.Dict[str | re.Pattern, str]
2074    ) -> t.Dict[re.Pattern, t.Any]:
2075        compiled = compile_regex_mapping(value)
2076        for replacement in compiled.values():
2077            if "@{schema_name}" not in replacement:
2078                raise ConfigError(
2079                    "schema_location_mapping needs to include the '@{schema_name}' placeholder in the value so SQLMesh knows where to substitute the schema name"
2080                )
2081        return compiled
2082
2083    @field_validator("timestamp_mapping", mode="before")
2084    @classmethod
2085    def _validate_timestamp_mapping(
2086        cls, value: t.Optional[dict[str, str]]
2087    ) -> t.Optional[dict[exp.DataType, exp.DataType]]:
2088        if value is None:
2089            return value
2090
2091        result: dict[exp.DataType, exp.DataType] = {}
2092        for source_type, target_type in value.items():
2093            try:
2094                source_datatype = exp.DataType.build(source_type)
2095            except ParseError:
2096                raise ConfigError(
2097                    f"Invalid SQL type string in timestamp_mapping: "
2098                    f"'{source_type}' is not a valid SQL data type."
2099                )
2100            try:
2101                target_datatype = exp.DataType.build(target_type)
2102            except ParseError:
2103                raise ConfigError(
2104                    f"Invalid SQL type string in timestamp_mapping: "
2105                    f"'{target_type}' is not a valid SQL data type."
2106                )
2107            result[source_datatype] = target_datatype
2108
2109        return result
2110
2111    @model_validator(mode="after")
2112    def _root_validator(self) -> Self:
2113        port = self.port
2114        if self.http_scheme == "http" and not self.method.is_no_auth and not self.method.is_basic:
2115            raise ConfigError("HTTP scheme can only be used with no-auth or basic method")
2116
2117        if port is None:
2118            self.port = 80 if self.http_scheme == "http" else 443
2119
2120        if (self.method.is_ldap or self.method.is_basic) and (not self.password or not self.user):
2121            raise ConfigError(
2122                f"Username and Password must be provided if using {self.method.value} authentication"
2123            )
2124
2125        if self.method.is_kerberos and (
2126            not self.principal or not self.keytab or not self.krb5_config
2127        ):
2128            raise ConfigError(
2129                "Kerberos requires the following fields: principal, keytab, and krb5_config"
2130            )
2131
2132        if self.method.is_jwt and not self.jwt_token:
2133            raise ConfigError("JWT requires `jwt_token` to be set")
2134
2135        if self.method.is_certificate and (
2136            not self.cert or not self.client_certificate or not self.client_private_key
2137        ):
2138            raise ConfigError(
2139                "Certificate requires the following fields: cert, client_certificate, and client_private_key"
2140            )
2141
2142        return self
2143
2144    @property
2145    def _connection_kwargs_keys(self) -> t.Set[str]:
2146        kwargs = {
2147            "host",
2148            "port",
2149            "catalog",
2150            "roles",
2151            "source",
2152            "http_scheme",
2153            "http_headers",
2154            "session_properties",
2155            "timezone",
2156        }
2157        return kwargs
2158
2159    @property
2160    def _engine_adapter(self) -> t.Type[EngineAdapter]:
2161        return engine_adapter.TrinoEngineAdapter
2162
2163    @property
2164    def _connection_factory(self) -> t.Callable:
2165        from trino.dbapi import connect
2166
2167        return connect
2168
2169    @property
2170    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
2171        from trino.auth import (
2172            BasicAuthentication,
2173            CertificateAuthentication,
2174            JWTAuthentication,
2175            KerberosAuthentication,
2176            OAuth2Authentication,
2177        )
2178
2179        auth: t.Optional[
2180            t.Union[
2181                BasicAuthentication,
2182                KerberosAuthentication,
2183                OAuth2Authentication,
2184                JWTAuthentication,
2185                CertificateAuthentication,
2186            ]
2187        ] = None
2188        if self.method.is_basic or self.method.is_ldap:
2189            assert self.password is not None  # for mypy since validator already checks this
2190            auth = BasicAuthentication(self.user, self.password)
2191        elif self.method.is_kerberos:
2192            if self.keytab:
2193                os.environ["KRB5_CLIENT_KTNAME"] = self.keytab
2194            auth = KerberosAuthentication(
2195                config=self.krb5_config,
2196                service_name=self.service_name,
2197                principal=self.principal,
2198                mutual_authentication=self.mutual_authentication,
2199                ca_bundle=self.cert,
2200                force_preemptive=self.force_preemptive,
2201                hostname_override=self.hostname_override,
2202                sanitize_mutual_error_response=self.sanitize_mutual_error_response,
2203                delegate=self.delegate,
2204            )
2205        elif self.method.is_oauth:
2206            auth = OAuth2Authentication()
2207        elif self.method.is_jwt:
2208            assert self.jwt_token is not None
2209            auth = JWTAuthentication(self.jwt_token)
2210        elif self.method.is_certificate:
2211            assert self.client_certificate is not None
2212            assert self.client_private_key is not None
2213            auth = CertificateAuthentication(self.client_certificate, self.client_private_key)
2214
2215        return {
2216            "auth": auth,
2217            "user": self.impersonation_user or self.user,
2218            "max_attempts": self.retries,
2219            "verify": self.cert if self.cert is not None else self.verify,
2220            "source": self.source,
2221        }
2222
2223    @property
2224    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
2225        return {
2226            "schema_location_mapping": self.schema_location_mapping,
2227            "timestamp_mapping": self.timestamp_mapping,
2228        }
2229
2230
2231class ClickhouseConnectionConfig(ConnectionConfig):
2232    """
2233    Clickhouse Connection Configuration.
2234
2235    Property reference: https://clickhouse.com/docs/en/integrations/python#client-initialization
2236    """
2237
2238    host: str
2239    username: str
2240    password: t.Optional[str] = None
2241    port: t.Optional[int] = None
2242    cluster: t.Optional[str] = None
2243    virtual_catalog: t.Optional[str] = None
2244    connect_timeout: int = 10
2245    send_receive_timeout: int = 300
2246    query_limit: int = 0
2247    use_compression: bool = True
2248    compression_method: t.Optional[str] = None
2249    connection_settings: t.Optional[t.Dict[str, t.Any]] = None
2250    http_proxy: t.Optional[str] = None
2251    # HTTPS/TLS settings
2252    verify: bool = True
2253    ca_cert: t.Optional[str] = None
2254    client_cert: t.Optional[str] = None
2255    client_cert_key: t.Optional[str] = None
2256    https_proxy: t.Optional[str] = None
2257    server_host_name: t.Optional[str] = None
2258    tls_mode: t.Optional[str] = None
2259    secure: bool = False
2260
2261    concurrent_tasks: int = 1
2262    register_comments: bool = True
2263    pre_ping: bool = False
2264
2265    # This object expects options from urllib3 and also from clickhouse-connect
2266    # See:
2267    # * https://urllib3.readthedocs.io/en/stable/advanced-usage.html
2268    # * https://clickhouse.com/docs/en/integrations/python#customizing-the-http-connection-pool
2269    connection_pool_options: t.Optional[t.Dict[str, t.Any]] = None
2270
2271    type_: t.Literal["clickhouse"] = Field(alias="type", default="clickhouse")
2272    DIALECT: t.ClassVar[t.Literal["clickhouse"]] = "clickhouse"
2273    DISPLAY_NAME: t.ClassVar[t.Literal["ClickHouse"]] = "ClickHouse"
2274    DISPLAY_ORDER: t.ClassVar[t.Literal[6]] = 6
2275
2276    _engine_import_validator = _get_engine_import_validator("clickhouse_connect", "clickhouse")
2277
2278    @field_validator("virtual_catalog")
2279    def validate_virtual_catalog(cls, v: t.Optional[str]) -> t.Optional[str]:
2280        if v is not None and not v.strip():
2281            raise ConfigError(
2282                "virtual_catalog cannot be an empty string. "
2283                "Omit the field to use the default synthetic prefix (__<gateway_name>__)."
2284            )
2285        if v is not None and "." in v:
2286            raise ConfigError(
2287                f"virtual_catalog must be a single identifier with no dots (got: {v!r})"
2288            )
2289        return v
2290
2291    @property
2292    def _connection_kwargs_keys(self) -> t.Set[str]:
2293        kwargs = {
2294            "host",
2295            "username",
2296            "port",
2297            "password",
2298            "connect_timeout",
2299            "send_receive_timeout",
2300            "query_limit",
2301            "http_proxy",
2302            "verify",
2303            "ca_cert",
2304            "client_cert",
2305            "client_cert_key",
2306            "https_proxy",
2307            "server_host_name",
2308            "tls_mode",
2309            "secure",
2310        }
2311        return kwargs
2312
2313    @property
2314    def _engine_adapter(self) -> t.Type[EngineAdapter]:
2315        return engine_adapter.ClickhouseEngineAdapter
2316
2317    @property
2318    def _connection_factory(self) -> t.Callable:
2319        from functools import partial
2320
2321        from clickhouse_connect.dbapi import connect  # type: ignore
2322        from clickhouse_connect.driver import httputil  # type: ignore
2323
2324        pool_manager_options: t.Dict[str, t.Any] = dict(
2325            # Match the maxsize to the number of concurrent tasks
2326            maxsize=self.concurrent_tasks,
2327            # Block if there are no free connections
2328            block=True,
2329            verify=self.verify,
2330            ca_cert=self.ca_cert,
2331            client_cert=self.client_cert,
2332            client_cert_key=self.client_cert_key,
2333            https_proxy=self.https_proxy,
2334        )
2335        # this doesn't happen automatically because we always supply our own pool manager to the connection
2336        # https://github.com/ClickHouse/clickhouse-connect/blob/3a7f4b04cad29c7c2536661b831fb744248e2ec0/clickhouse_connect/driver/httpclient.py#L109
2337        if self.server_host_name:
2338            pool_manager_options["server_hostname"] = self.server_host_name
2339            if self.verify:
2340                pool_manager_options["assert_hostname"] = self.server_host_name
2341        if self.connection_pool_options:
2342            pool_manager_options.update(self.connection_pool_options)
2343        pool_mgr = httputil.get_pool_manager(**pool_manager_options)
2344
2345        return partial(connect, pool_mgr=pool_mgr)
2346
2347    @property
2348    def cloud_mode(self) -> bool:
2349        return "clickhouse.cloud" in self.host
2350
2351    @property
2352    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
2353        return {
2354            "cluster": self.cluster,
2355            "cloud_mode": self.cloud_mode,
2356            "virtual_catalog": self.virtual_catalog,
2357        }
2358
2359    @property
2360    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
2361        from sqlmesh import __version__
2362
2363        # False = no compression
2364        # True = Clickhouse default compression method
2365        # string = specific compression method
2366        compress: bool | str = self.use_compression
2367        if compress and self.compression_method:
2368            compress = self.compression_method
2369
2370        # Clickhouse system settings passed to connection
2371        # https://clickhouse.com/docs/en/operations/settings/settings
2372        # - below are set to align with dbt-clickhouse
2373        # - https://github.com/ClickHouse/dbt-clickhouse/blob/44d26308ea6a3c8ead25c280164aa88191f05f47/dbt/adapters/clickhouse/dbclient.py#L77
2374        settings = self.connection_settings or {}
2375        #  mutations_sync = 2: "The query waits for all mutations [ALTER statements] to complete on all replicas (if they exist)"
2376        settings["mutations_sync"] = "2"
2377        #  insert_distributed_sync = 1: "INSERT operation succeeds only after all the data is saved on all shards"
2378        settings["insert_distributed_sync"] = "1"
2379        if self.cluster or self.cloud_mode:
2380            # database_replicated_enforce_synchronous_settings = 1:
2381            #   - "Enforces synchronous waiting for some queries"
2382            #   - https://github.com/ClickHouse/ClickHouse/blob/ccaa8d03a9351efc16625340268b9caffa8a22ba/src/Core/Settings.h#L709
2383            settings["database_replicated_enforce_synchronous_settings"] = "1"
2384            # insert_quorum = auto:
2385            #   - "INSERT succeeds only when ClickHouse manages to correctly write data to the insert_quorum of replicas during
2386            #       the insert_quorum_timeout"
2387            #   - "use majority number (number_of_replicas / 2 + 1) as quorum number"
2388            settings["insert_quorum"] = "auto"
2389
2390        return {
2391            "compress": compress,
2392            "client_name": f"SQLMesh/{__version__}",
2393            **settings,
2394        }
2395
2396
2397class AthenaConnectionConfig(ConnectionConfig):
2398    # PyAthena connection options
2399    aws_access_key_id: t.Optional[str] = None
2400    aws_secret_access_key: t.Optional[str] = None
2401    role_arn: t.Optional[str] = None
2402    role_session_name: t.Optional[str] = None
2403    region_name: t.Optional[str] = None
2404    work_group: t.Optional[str] = None
2405    s3_staging_dir: t.Optional[str] = None
2406    schema_name: t.Optional[str] = None
2407    catalog_name: t.Optional[str] = None
2408
2409    # SQLMesh options
2410    s3_warehouse_location: t.Optional[str] = None
2411    concurrent_tasks: int = 4
2412    register_comments: t.Literal[False] = (
2413        False  # because Athena doesnt support comments in most cases
2414    )
2415    pre_ping: t.Literal[False] = False
2416
2417    type_: t.Literal["athena"] = Field(alias="type", default="athena")
2418    DIALECT: t.ClassVar[t.Literal["athena"]] = "athena"
2419    DISPLAY_NAME: t.ClassVar[t.Literal["Athena"]] = "Athena"
2420    DISPLAY_ORDER: t.ClassVar[t.Literal[15]] = 15
2421
2422    _engine_import_validator = _get_engine_import_validator("pyathena", "athena")
2423
2424    @model_validator(mode="after")
2425    def _root_validator(self) -> Self:
2426        work_group = self.work_group
2427        s3_staging_dir = self.s3_staging_dir
2428        s3_warehouse_location = self.s3_warehouse_location
2429
2430        if not work_group and not s3_staging_dir:
2431            raise ConfigError("At least one of work_group or s3_staging_dir must be set")
2432
2433        if s3_staging_dir:
2434            self.s3_staging_dir = validate_s3_uri(s3_staging_dir, base=True, error_type=ConfigError)
2435
2436        if s3_warehouse_location:
2437            self.s3_warehouse_location = validate_s3_uri(
2438                s3_warehouse_location, base=True, error_type=ConfigError
2439            )
2440
2441        return self
2442
2443    @property
2444    def _connection_kwargs_keys(self) -> t.Set[str]:
2445        return {
2446            "aws_access_key_id",
2447            "aws_secret_access_key",
2448            "role_arn",
2449            "role_session_name",
2450            "region_name",
2451            "work_group",
2452            "s3_staging_dir",
2453            "schema_name",
2454            "catalog_name",
2455        }
2456
2457    @property
2458    def _engine_adapter(self) -> t.Type[EngineAdapter]:
2459        return engine_adapter.AthenaEngineAdapter
2460
2461    @property
2462    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
2463        return {"s3_warehouse_location": self.s3_warehouse_location}
2464
2465    @property
2466    def _connection_factory(self) -> t.Callable:
2467        from pyathena import connect  # type: ignore
2468
2469        return connect
2470
2471    def get_catalog(self) -> t.Optional[str]:
2472        return self.catalog_name
2473
2474
2475class RisingwaveConnectionConfig(ConnectionConfig):
2476    host: str
2477    user: str
2478    password: t.Optional[str] = None
2479    port: int
2480    database: str
2481    role: t.Optional[str] = None
2482    sslmode: t.Optional[str] = None
2483
2484    concurrent_tasks: int = 4
2485    register_comments: bool = True
2486    pre_ping: bool = True
2487
2488    type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
2489    DIALECT: t.ClassVar[t.Literal["risingwave"]] = "risingwave"
2490    DISPLAY_NAME: t.ClassVar[t.Literal["RisingWave"]] = "RisingWave"
2491    DISPLAY_ORDER: t.ClassVar[t.Literal[16]] = 16
2492
2493    _engine_import_validator = _get_engine_import_validator("psycopg2", "risingwave")
2494
2495    @property
2496    def _connection_kwargs_keys(self) -> t.Set[str]:
2497        return {
2498            "host",
2499            "user",
2500            "password",
2501            "port",
2502            "database",
2503            "role",
2504            "sslmode",
2505        }
2506
2507    @property
2508    def _engine_adapter(self) -> t.Type[EngineAdapter]:
2509        return engine_adapter.RisingwaveEngineAdapter
2510
2511    @property
2512    def _connection_factory(self) -> t.Callable:
2513        from psycopg2 import connect
2514
2515        return connect
2516
2517    @property
2518    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
2519        def init(cursor: t.Any) -> None:
2520            sql = "SET RW_IMPLICIT_FLUSH TO true;"
2521            cursor.execute(sql)
2522
2523        return init
2524
2525
2526class StarRocksConnectionConfig(ConnectionConfig):
2527    """Configuration for the StarRocks connection.
2528
2529    StarRocks uses MySQL network protocol and is compatible with MySQL ecosystem tools,
2530    JDBC/ODBC drivers, and various visualization tools.
2531
2532    Args:
2533        host: The hostname of the StarRocks FE (Frontend) node.
2534        user: The StarRocks username.
2535        password: The StarRocks password.
2536        port: The port number of the StarRocks FE node. Default is 9030.
2537        database: The optional database name.
2538        charset: The optional character set.  TODO: may be not supported yet.
2539        collation: The optional collation.  TODO: may be not supported yet.
2540        ssl_disabled: Whether to disable SSL connection.  TODO: need to check it.
2541        concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
2542        register_comments: Whether or not to register model comments with the SQL engine.
2543        local_infile: Whether or not to allow local file access.
2544        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
2545    """
2546
2547    host: str
2548    user: str
2549    password: str
2550    port: t.Optional[int] = 9030
2551    database: t.Optional[str] = None
2552    charset: t.Optional[str] = None
2553    collation: t.Optional[str] = None
2554    ssl_disabled: t.Optional[bool] = None
2555
2556    concurrent_tasks: int = 4
2557    register_comments: bool = True
2558    local_infile: bool = False
2559    pre_ping: bool = True
2560
2561    type_: t.Literal["starrocks"] = Field(alias="type", default="starrocks")
2562    DIALECT: t.ClassVar[t.Literal["starrocks"]] = "starrocks"
2563    DISPLAY_NAME: t.ClassVar[t.Literal["StarRocks"]] = "StarRocks"
2564    DISPLAY_ORDER: t.ClassVar[t.Literal[18]] = 18
2565
2566    _engine_import_validator = _get_engine_import_validator("pymysql", "starrocks")
2567
2568    @property
2569    def _connection_kwargs_keys(self) -> t.Set[str]:
2570        connection_keys = {
2571            "host",
2572            "user",
2573            "password",
2574        }
2575        if self.port is not None:
2576            connection_keys.add("port")
2577        if self.database is not None:
2578            connection_keys.add("database")
2579        if self.charset is not None:
2580            connection_keys.add("charset")
2581        if self.collation is not None:
2582            connection_keys.add("collation")
2583        if self.ssl_disabled is not None:
2584            connection_keys.add("ssl_disabled")
2585        if self.local_infile is not None:
2586            connection_keys.add("local_infile")
2587        return connection_keys
2588
2589    @property
2590    def _engine_adapter(self) -> t.Type[EngineAdapter]:
2591        return engine_adapter.StarRocksEngineAdapter
2592
2593    @property
2594    def _connection_factory(self) -> t.Callable:
2595        from pymysql import connect
2596
2597        return connect
2598
2599
2600_CONNECTION_CONFIG_EXCLUDE: t.Set[t.Type[ConnectionConfig]] = {
2601    ConnectionConfig,  # type: ignore[type-abstract]
2602    BaseDuckDBConnectionConfig,  # type: ignore[type-abstract]
2603}
2604
2605CONNECTION_CONFIG_TO_TYPE = {
2606    # Map all subclasses of ConnectionConfig to the value of their `type_` field.
2607    tpe.all_field_infos()["type_"].default: tpe
2608    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
2609}
2610
2611DIALECT_TO_TYPE = {
2612    tpe.all_field_infos()["type_"].default: tpe.DIALECT
2613    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
2614}
2615
2616INIT_DISPLAY_INFO_TO_TYPE = {
2617    tpe.all_field_infos()["type_"].default: (
2618        tpe.DISPLAY_ORDER,
2619        tpe.DISPLAY_NAME,
2620    )
2621    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
2622}
2623
2624
2625def parse_connection_config(v: t.Dict[str, t.Any]) -> ConnectionConfig:
2626    if "type" not in v:
2627        raise ConfigError("Missing connection type.")
2628
2629    connection_type = v["type"]
2630    if connection_type not in CONNECTION_CONFIG_TO_TYPE:
2631        raise ConfigError(f"Unknown connection type '{connection_type}'.")
2632
2633    return CONNECTION_CONFIG_TO_TYPE[connection_type](**v)
2634
2635
2636def _connection_config_validator(
2637    cls: t.Type, v: ConnectionConfig | t.Dict[str, t.Any] | None
2638) -> ConnectionConfig | None:
2639    if v is None or isinstance(v, ConnectionConfig):
2640        return v
2641
2642    check_config_and_vars_msg = "\n\nVerify your config.yaml and environment variables."
2643
2644    try:
2645        return parse_connection_config(v)
2646    except pydantic.ValidationError as e:
2647        raise ConfigError(
2648            validation_error_message(e, f"Invalid '{v['type']}' connection config:")
2649            + check_config_and_vars_msg
2650        )
2651    except ConfigError as e:
2652        raise ConfigError(str(e) + check_config_and_vars_msg)
2653
2654
2655connection_config_validator: t.Callable = field_validator(
2656    "connection",
2657    "state_connection",
2658    "test_connection",
2659    "default_connection",
2660    "default_test_connection",
2661    mode="before",
2662    check_fields=False,
2663)(_connection_config_validator)
2664
2665
2666if t.TYPE_CHECKING:
2667    # TypeAlias hasn't been introduced until Python 3.10 which means that we can't use it
2668    # outside the TYPE_CHECKING guard.
2669    SerializableConnectionConfig: t.TypeAlias = ConnectionConfig  # type: ignore
2670else:
2671    import pydantic
2672
2673    # Workaround for https://docs.pydantic.dev/latest/concepts/serialization/#serializing-with-duck-typing
2674    SerializableConnectionConfig = pydantic.SerializeAsAny[ConnectionConfig]  # type: ignore