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