sqlmesh.core.config.connection

   1from __future__ import annotations
   2
   3import abc
   4import base64
   5import logging
   6import os
   7import pathlib
   8import sys
   9import typing as t
  10from enum import Enum
  11from functools import partial
  12
  13from pydantic import Field
  14from sqlglot import exp
  15from sqlglot.helper import subclasses
  16
  17from sqlmesh.core import engine_adapter
  18from sqlmesh.core.config.base import BaseConfig
  19from sqlmesh.core.config.common import (
  20    concurrent_tasks_validator,
  21    http_headers_validator,
  22)
  23from sqlmesh.core.engine_adapter import EngineAdapter
  24from sqlmesh.utils.errors import ConfigError
  25from sqlmesh.utils.pydantic import (
  26    PYDANTIC_MAJOR_VERSION,
  27    field_validator,
  28    model_validator,
  29    model_validator_v1_args,
  30)
  31
  32if sys.version_info >= (3, 9):
  33    from typing import Literal
  34else:
  35    from typing_extensions import Literal
  36
  37
  38logger = logging.getLogger(__name__)
  39
  40RECOMMENDED_STATE_SYNC_ENGINES = {"postgres", "gcp_postgres", "mysql", "duckdb"}
  41
  42
  43class ConnectionConfig(abc.ABC, BaseConfig):
  44    type_: str
  45    concurrent_tasks: int
  46    register_comments: bool
  47
  48    @property
  49    @abc.abstractmethod
  50    def _connection_kwargs_keys(self) -> t.Set[str]:
  51        """keywords that should be passed into the connection"""
  52
  53    @property
  54    @abc.abstractmethod
  55    def _engine_adapter(self) -> t.Type[EngineAdapter]:
  56        """The engine adapter for this connection"""
  57
  58    @property
  59    @abc.abstractmethod
  60    def _connection_factory(self) -> t.Callable:
  61        """A function that is called to return a connection object for the given Engine Adapter"""
  62
  63    @property
  64    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
  65        """The static connection kwargs for this connection"""
  66        return {}
  67
  68    @property
  69    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
  70        """kwargs that are for execution config only"""
  71        return {}
  72
  73    @property
  74    def _cursor_kwargs(self) -> t.Optional[t.Dict[str, t.Any]]:
  75        """Key-value arguments that will be passed during cursor construction."""
  76        return None
  77
  78    @property
  79    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
  80        """A function that is called to initialize the cursor"""
  81        return None
  82
  83    @property
  84    def is_recommended_for_state_sync(self) -> bool:
  85        """Whether this connection is recommended for being used as a state sync for production state syncs"""
  86        return self.type_ in RECOMMENDED_STATE_SYNC_ENGINES
  87
  88    @property
  89    def _connection_factory_with_kwargs(self) -> t.Callable[[], t.Any]:
  90        """A function that is called to return a connection object for the given Engine Adapter"""
  91        return partial(
  92            self._connection_factory,
  93            **{
  94                **self._static_connection_kwargs,
  95                **{k: v for k, v in self.dict().items() if k in self._connection_kwargs_keys},
  96            },
  97        )
  98
  99    def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
 100        """Returns a new instance of the Engine Adapter."""
 101        return self._engine_adapter(
 102            self._connection_factory_with_kwargs,
 103            multithreaded=self.concurrent_tasks > 1,
 104            cursor_kwargs=self._cursor_kwargs,
 105            default_catalog=self.get_catalog(),
 106            cursor_init=self._cursor_init,
 107            register_comments=register_comments_override or self.register_comments,
 108            **self._extra_engine_config,
 109        )
 110
 111    def get_catalog(self) -> t.Optional[str]:
 112        """The catalog for this connection"""
 113        if hasattr(self, "catalog"):
 114            return self.catalog
 115        if hasattr(self, "database"):
 116            return self.database
 117        if hasattr(self, "db"):
 118            return self.db
 119        return None
 120
 121
 122class BaseDuckDBConnectionConfig(ConnectionConfig):
 123    """Common configuration for the DuckDB-based connections.
 124
 125    Args:
 126        extensions: A list of autoloadable extensions to load.
 127        connector_config: A dictionary of configuration to pass into the duckdb connector.
 128        concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
 129        register_comments: Whether or not to register model comments with the SQL engine.
 130    """
 131
 132    extensions: t.List[str] = []
 133    connector_config: t.Dict[str, t.Any] = {}
 134
 135    concurrent_tasks: Literal[1] = 1
 136    register_comments: bool = True
 137
 138    @property
 139    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 140        return engine_adapter.DuckDBEngineAdapter
 141
 142    @property
 143    def _connection_factory(self) -> t.Callable:
 144        import duckdb
 145
 146        return duckdb.connect
 147
 148    @property
 149    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
 150        """A function that is called to initialize the cursor"""
 151        import duckdb
 152        from duckdb import BinderException
 153
 154        def init(cursor: duckdb.DuckDBPyConnection) -> None:
 155            for extension in self.extensions:
 156                try:
 157                    cursor.execute(f"INSTALL {extension}")
 158                    cursor.execute(f"LOAD {extension}")
 159                except Exception as e:
 160                    raise ConfigError(f"Failed to load extension {extension}: {e}")
 161
 162            for field, setting in self.connector_config.items():
 163                try:
 164                    cursor.execute(f"SET {field} = '{setting}'")
 165                except Exception as e:
 166                    raise ConfigError(f"Failed to set connector config {field} to {setting}: {e}")
 167
 168            for i, (alias, path_options) in enumerate(
 169                (getattr(self, "catalogs", None) or {}).items()
 170            ):
 171                # we parse_identifier and generate to ensure that `alias` has exactly one set of quotes
 172                # regardless of whether it comes in quoted or not
 173                alias = exp.parse_identifier(alias, dialect="duckdb").sql(
 174                    identify=True, dialect="duckdb"
 175                )
 176                try:
 177                    query = (
 178                        path_options.to_sql(alias)
 179                        if isinstance(path_options, DuckDBAttachOptions)
 180                        else f"ATTACH '{path_options}' AS {alias}"
 181                    )
 182                    cursor.execute(query)
 183                except BinderException as e:
 184                    # If a user tries to create a catalog pointing at `:memory:` and with the name `memory`
 185                    # then we don't want to raise since this happens by default. They are just doing this to
 186                    # set it as the default catalog.
 187                    if not (
 188                        'database with name "memory" already exists' in str(e)
 189                        and path_options == ":memory:"
 190                    ):
 191                        raise e
 192                if i == 0 and not getattr(self, "database", None):
 193                    cursor.execute(f"USE {alias}")
 194
 195        return init
 196
 197
 198class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
 199    """Configuration for the MotherDuck connection.
 200
 201    Args:
 202        database: The database name.
 203        token: The optional MotherDuck token. If not specified, the user will be prompted to login with their web browser.
 204    """
 205
 206    database: str
 207    token: t.Optional[str] = None
 208
 209    type_: Literal["motherduck"] = Field(alias="type", default="motherduck")
 210
 211    @property
 212    def _connection_kwargs_keys(self) -> t.Set[str]:
 213        return set()
 214
 215    @property
 216    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 217        """kwargs that are for execution config only"""
 218        connection_str = f"md:{self.database}"
 219        if self.token:
 220            connection_str += f"?motherduck_token={self.token}"
 221        return {"database": connection_str}
 222
 223
 224class DuckDBAttachOptions(BaseConfig):
 225    type: str
 226    path: str
 227    read_only: bool = False
 228
 229    def to_sql(self, alias: str) -> str:
 230        options = []
 231        # 'duckdb' is actually not a supported type, but we'd like to allow it for
 232        # fully qualified attach options or integration testing, similar to duckdb-dbt
 233        if self.type != "duckdb":
 234            options.append(f"TYPE {self.type.upper()}")
 235        if self.read_only:
 236            options.append("READ_ONLY")
 237        options_sql = f" ({', '.join(options)})" if options else ""
 238        return f"ATTACH '{self.path}' AS {alias}{options_sql}"
 239
 240
 241class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
 242    """Configuration for the DuckDB connection.
 243
 244    Args:
 245        database: The optional database name. If not specified, the in-memory database will be used.
 246        catalogs: Key is the name of the catalog and value is the path.
 247    """
 248
 249    database: t.Optional[str] = None
 250    catalogs: t.Optional[t.Dict[str, t.Union[str, DuckDBAttachOptions]]] = None
 251
 252    type_: Literal["duckdb"] = Field(alias="type", default="duckdb")
 253
 254    _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
 255
 256    @model_validator(mode="before")
 257    @model_validator_v1_args
 258    def _validate_database_catalogs(
 259        cls, values: t.Dict[str, t.Optional[str]]
 260    ) -> t.Dict[str, t.Optional[str]]:
 261        if values.get("database") and values.get("catalogs"):
 262            raise ConfigError(
 263                "Cannot specify both `database` and `catalogs`. Define all your catalogs in `catalogs` and have the first entry be the default catalog"
 264            )
 265        return values
 266
 267    @property
 268    def _connection_kwargs_keys(self) -> t.Set[str]:
 269        return {"database"}
 270
 271    def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
 272        """Checks if another engine adapter has already been created that shares a catalog that points to the same data
 273        file. If so, it uses that same adapter instead of creating a new one. As a result, any additional configuration
 274        associated with the new adapter will be ignored."""
 275        data_files = set((self.catalogs or {}).values())
 276        if self.database:
 277            data_files.add(self.database)
 278        data_files.discard(":memory:")
 279        for data_file in data_files:
 280            key = data_file if isinstance(data_file, str) else data_file.path
 281            if adapter := DuckDBConnectionConfig._data_file_to_adapter.get(key):
 282                logger.info(f"Using existing DuckDB adapter due to overlapping data file: {key}")
 283                return adapter
 284
 285        if data_files:
 286            logger.info(f"Creating new DuckDB adapter for data files: {data_files}")
 287        else:
 288            logger.info("Creating new DuckDB adapter for in-memory database")
 289        adapter = super().create_engine_adapter(register_comments_override)
 290        for data_file in data_files:
 291            key = data_file if isinstance(data_file, str) else data_file.path
 292            DuckDBConnectionConfig._data_file_to_adapter[key] = adapter
 293        return adapter
 294
 295    def get_catalog(self) -> t.Optional[str]:
 296        if self.database:
 297            # Remove `:` from the database name in order to handle if `:memory:` is passed in
 298            return pathlib.Path(self.database.replace(":", "")).stem
 299        if self.catalogs:
 300            return list(self.catalogs)[0]
 301        return None
 302
 303
 304class SnowflakeConnectionConfig(ConnectionConfig):
 305    """Configuration for the Snowflake connection.
 306
 307    Args:
 308        account: The Snowflake account name.
 309        user: The Snowflake username.
 310        password: The Snowflake password.
 311        warehouse: The optional warehouse name.
 312        database: The optional database name.
 313        role: The optional role name.
 314        concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
 315        authenticator: The optional authenticator name. Defaults to username/password authentication ("snowflake").
 316                       Options: https://github.com/snowflakedb/snowflake-connector-python/blob/e937591356c067a77f34a0a42328907fda792c23/src/snowflake/connector/network.py#L178-L183
 317        token: The optional oauth access token to use for authentication when authenticator is set to "oauth".
 318        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
 319        private_key_path: The optional path to the private key to use for authentication. This would be used instead of `private_key`.
 320        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.
 321        register_comments: Whether or not to register model comments with the SQL engine.
 322    """
 323
 324    account: str
 325    user: t.Optional[str] = None
 326    password: t.Optional[str] = None
 327    warehouse: t.Optional[str] = None
 328    database: t.Optional[str] = None
 329    role: t.Optional[str] = None
 330    authenticator: t.Optional[str] = None
 331    token: t.Optional[str] = None
 332
 333    # Private Key Auth
 334    private_key: t.Optional[t.Union[str, bytes]] = None
 335    private_key_path: t.Optional[str] = None
 336    private_key_passphrase: t.Optional[str] = None
 337
 338    concurrent_tasks: int = 4
 339    register_comments: bool = True
 340
 341    type_: Literal["snowflake"] = Field(alias="type", default="snowflake")
 342
 343    _concurrent_tasks_validator = concurrent_tasks_validator
 344
 345    @model_validator(mode="before")
 346    @model_validator_v1_args
 347    def _validate_authenticator(
 348        cls, values: t.Dict[str, t.Optional[str]]
 349    ) -> t.Dict[str, t.Optional[str]]:
 350        from snowflake.connector.network import (
 351            DEFAULT_AUTHENTICATOR,
 352            OAUTH_AUTHENTICATOR,
 353        )
 354
 355        auth = values.get("authenticator")
 356        auth = auth.upper() if auth else DEFAULT_AUTHENTICATOR
 357        user = values.get("user")
 358        password = values.get("password")
 359        values["private_key"] = cls._get_private_key(values, auth)  # type: ignore
 360        if (
 361            auth == DEFAULT_AUTHENTICATOR
 362            and not values.get("private_key")
 363            and (not user or not password)
 364        ):
 365            raise ConfigError("User and password must be provided if using default authentication")
 366        if auth == OAUTH_AUTHENTICATOR and not values.get("token"):
 367            raise ConfigError("Token must be provided if using oauth authentication")
 368        return values
 369
 370    @classmethod
 371    def _get_private_key(cls, values: t.Dict[str, t.Optional[str]], auth: str) -> t.Optional[bytes]:
 372        """
 373        source: https://github.com/dbt-labs/dbt-snowflake/blob/0374b4ec948982f2ac8ec0c95d53d672ad19e09c/dbt/adapters/snowflake/connections.py#L247C5-L285C1
 374
 375        Overall code change: Use local variables instead of class attributes + Validation
 376        """
 377        # Start custom code
 378        from cryptography.hazmat.backends import default_backend
 379        from cryptography.hazmat.primitives import serialization
 380        from snowflake.connector.network import (
 381            DEFAULT_AUTHENTICATOR,
 382            KEY_PAIR_AUTHENTICATOR,
 383        )
 384
 385        private_key = values.get("private_key")
 386        private_key_path = values.get("private_key_path")
 387        private_key_passphrase = values.get("private_key_passphrase")
 388        user = values.get("user")
 389        password = values.get("password")
 390        auth = auth if auth and auth != DEFAULT_AUTHENTICATOR else KEY_PAIR_AUTHENTICATOR
 391
 392        if not private_key and not private_key_path:
 393            return None
 394        if private_key and private_key_path:
 395            raise ConfigError("Cannot specify both `private_key` and `private_key_path`")
 396        if auth != KEY_PAIR_AUTHENTICATOR:
 397            raise ConfigError(
 398                f"Private key or private key path can only be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 399            )
 400        if not user:
 401            raise ConfigError(
 402                f"User must be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 403            )
 404        if password:
 405            raise ConfigError(
 406                f"Password cannot be provided when using {KEY_PAIR_AUTHENTICATOR} authentication"
 407            )
 408
 409        if isinstance(private_key, bytes):
 410            return private_key
 411        # End Custom Code
 412
 413        if private_key_passphrase:
 414            encoded_passphrase = private_key_passphrase.encode()
 415        else:
 416            encoded_passphrase = None
 417
 418        if private_key:
 419            if private_key.startswith("-"):
 420                p_key = serialization.load_pem_private_key(
 421                    data=bytes(private_key, "utf-8"),
 422                    password=encoded_passphrase,
 423                    backend=default_backend(),
 424                )
 425
 426            else:
 427                p_key = serialization.load_der_private_key(
 428                    data=base64.b64decode(private_key),
 429                    password=encoded_passphrase,
 430                    backend=default_backend(),
 431                )
 432
 433        elif private_key_path:
 434            with open(private_key_path, "rb") as key:
 435                p_key = serialization.load_pem_private_key(
 436                    key.read(), password=encoded_passphrase, backend=default_backend()
 437                )
 438        else:
 439            return None
 440
 441        return p_key.private_bytes(
 442            encoding=serialization.Encoding.DER,
 443            format=serialization.PrivateFormat.PKCS8,
 444            encryption_algorithm=serialization.NoEncryption(),
 445        )
 446
 447    @property
 448    def _connection_kwargs_keys(self) -> t.Set[str]:
 449        return {
 450            "user",
 451            "password",
 452            "account",
 453            "warehouse",
 454            "database",
 455            "role",
 456            "authenticator",
 457            "token",
 458            "private_key",
 459        }
 460
 461    @property
 462    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 463        return engine_adapter.SnowflakeEngineAdapter
 464
 465    @property
 466    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 467        return {"autocommit": False}
 468
 469    @property
 470    def _connection_factory(self) -> t.Callable:
 471        from snowflake import connector
 472
 473        return connector.connect
 474
 475
 476class DatabricksConnectionConfig(ConnectionConfig):
 477    """
 478    Databricks connection that uses the SQL connector for SQL models and then Databricks Connect for Dataframe operations
 479
 480    Arg Source: https://github.com/databricks/databricks-sql-python/blob/main/src/databricks/sql/client.py#L39
 481    Args:
 482        server_hostname: Databricks instance host name.
 483        http_path: Http path either to a DBSQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef)
 484            or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)
 485        access_token: Http Bearer access token, e.g. Databricks Personal Access Token.
 486        catalog: Default catalog to use for SQL models. Defaults to None which means it will use the default set in
 487            the Databricks cluster (most likely `hive_metastore`).
 488        http_headers: An optional list of (k, v) pairs that will be set as Http headers on every request
 489        session_configuration: An optional dictionary of Spark session parameters.
 490            Execute the SQL command `SET -v` to get a full list of available commands.
 491        databricks_connect_server_hostname: The hostname to use when establishing a connecting using Databricks Connect.
 492            Defaults to the `server_hostname` value.
 493        databricks_connect_access_token: The access token to use when establishing a connecting using Databricks Connect.
 494            Defaults to the `access_token` value.
 495        databricks_connect_cluster_id: The cluster id to use when establishing a connecting using Databricks Connect.
 496            Defaults to deriving the cluster id from the `http_path` value.
 497        force_databricks_connect: Force all queries to run using Databricks Connect instead of the SQL connector.
 498        disable_databricks_connect: Even if databricks connect is installed, do not use it.
 499    """
 500
 501    server_hostname: t.Optional[str] = None
 502    http_path: t.Optional[str] = None
 503    access_token: t.Optional[str] = None
 504    catalog: t.Optional[str] = None
 505    http_headers: t.Optional[t.List[t.Tuple[str, str]]] = None
 506    session_configuration: t.Optional[t.Dict[str, t.Any]] = None
 507    databricks_connect_server_hostname: t.Optional[str] = None
 508    databricks_connect_access_token: t.Optional[str] = None
 509    databricks_connect_cluster_id: t.Optional[str] = None
 510    force_databricks_connect: bool = False
 511    disable_databricks_connect: bool = False
 512
 513    concurrent_tasks: int = 1
 514    register_comments: bool = True
 515
 516    type_: Literal["databricks"] = Field(alias="type", default="databricks")
 517
 518    _concurrent_tasks_validator = concurrent_tasks_validator
 519    _http_headers_validator = http_headers_validator
 520
 521    @model_validator(mode="before")
 522    @model_validator_v1_args
 523    def _databricks_connect_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
 524        from sqlmesh import RuntimeEnv
 525        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 526
 527        runtime_env = RuntimeEnv.get()
 528
 529        if runtime_env.is_databricks:
 530            return values
 531        server_hostname, http_path, access_token = (
 532            values.get("server_hostname"),
 533            values.get("http_path"),
 534            values.get("access_token"),
 535        )
 536        if not server_hostname or not http_path or not access_token:
 537            raise ValueError(
 538                "`server_hostname`, `http_path`, and `access_token` are required for Databricks connections when not running in a notebook"
 539            )
 540        if DatabricksEngineAdapter.can_access_spark_session:
 541            if not values.get("databricks_connect_server_hostname"):
 542                values["databricks_connect_server_hostname"] = f"https://{server_hostname}"
 543            if not values.get("databricks_connect_access_token"):
 544                values["databricks_connect_access_token"] = access_token
 545            if not values.get("databricks_connect_cluster_id"):
 546                values["databricks_connect_cluster_id"] = http_path.split("/")[-1]
 547        if not values.get("session_configuration"):
 548            values["session_configuration"] = {}
 549        values["session_configuration"]["spark.sql.sources.partitionOverwriteMode"] = "dynamic"
 550        return values
 551
 552    @property
 553    def _connection_kwargs_keys(self) -> t.Set[str]:
 554        if self.use_spark_session_only:
 555            return set()
 556        return {
 557            "server_hostname",
 558            "http_path",
 559            "access_token",
 560            "http_headers",
 561            "session_configuration",
 562            "catalog",
 563        }
 564
 565    @property
 566    def _engine_adapter(self) -> t.Type[engine_adapter.DatabricksEngineAdapter]:
 567        return engine_adapter.DatabricksEngineAdapter
 568
 569    @property
 570    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 571        return {
 572            k: v
 573            for k, v in self.dict().items()
 574            if k.startswith("databricks_connect_") or k in ("catalog", "disable_databricks_connect")
 575        }
 576
 577    @property
 578    def use_spark_session_only(self) -> bool:
 579        from sqlmesh import RuntimeEnv
 580
 581        return RuntimeEnv.get().is_databricks or self.force_databricks_connect
 582
 583    @property
 584    def _connection_factory(self) -> t.Callable:
 585        if self.use_spark_session_only:
 586            from sqlmesh.engines.spark.db_api.spark_session import connection
 587
 588            return connection
 589
 590        from databricks import sql
 591
 592        return sql.connect
 593
 594    @property
 595    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 596        from sqlmesh import RuntimeEnv
 597
 598        if not self.use_spark_session_only:
 599            return {}
 600
 601        if RuntimeEnv.get().is_databricks:
 602            from pyspark.sql import SparkSession
 603
 604            return dict(
 605                spark=SparkSession.getActiveSession(),
 606                catalog=self.catalog,
 607            )
 608
 609        from databricks.connect import DatabricksSession
 610
 611        return dict(
 612            spark=DatabricksSession.builder.remote(
 613                host=self.databricks_connect_server_hostname,
 614                token=self.databricks_connect_access_token,
 615                cluster_id=self.databricks_connect_cluster_id,
 616            ).getOrCreate(),
 617            catalog=self.catalog,
 618        )
 619
 620
 621class BigQueryConnectionMethod(str, Enum):
 622    OAUTH = "oauth"
 623    OAUTH_SECRETS = "oauth-secrets"
 624    SERVICE_ACCOUNT = "service-account"
 625    SERVICE_ACCOUNT_JSON = "service-account-json"
 626
 627
 628class BigQueryPriority(str, Enum):
 629    BATCH = "batch"
 630    INTERACTIVE = "interactive"
 631
 632    @property
 633    def is_batch(self) -> bool:
 634        return self == self.BATCH
 635
 636    @property
 637    def is_interactive(self) -> bool:
 638        return self == self.INTERACTIVE
 639
 640    @property
 641    def bigquery_constant(self) -> str:
 642        from google.cloud.bigquery import QueryPriority
 643
 644        if self.is_batch:
 645            return QueryPriority.BATCH
 646        return QueryPriority.INTERACTIVE
 647
 648
 649class BigQueryConnectionConfig(ConnectionConfig):
 650    """
 651    BigQuery Connection Configuration.
 652    """
 653
 654    method: BigQueryConnectionMethod = BigQueryConnectionMethod.OAUTH
 655
 656    project: t.Optional[str] = None
 657    execution_project: t.Optional[str] = None
 658    location: t.Optional[str] = None
 659    # Keyfile Auth
 660    keyfile: t.Optional[str] = None
 661    keyfile_json: t.Optional[t.Dict[str, t.Any]] = None
 662    # Oath Secret Auth
 663    token: t.Optional[str] = None
 664    refresh_token: t.Optional[str] = None
 665    client_id: t.Optional[str] = None
 666    client_secret: t.Optional[str] = None
 667    token_uri: t.Optional[str] = None
 668    scopes: t.Tuple[str, ...] = ("https://www.googleapis.com/auth/bigquery",)
 669    job_creation_timeout_seconds: t.Optional[int] = None
 670    # Extra Engine Config
 671    job_execution_timeout_seconds: t.Optional[int] = None
 672    job_retries: t.Optional[int] = 1
 673    job_retry_deadline_seconds: t.Optional[int] = None
 674    priority: t.Optional[BigQueryPriority] = None
 675    maximum_bytes_billed: t.Optional[int] = None
 676
 677    concurrent_tasks: int = 1
 678    register_comments: bool = True
 679
 680    type_: Literal["bigquery"] = Field(alias="type", default="bigquery")
 681
 682    @property
 683    def _connection_kwargs_keys(self) -> t.Set[str]:
 684        return set()
 685
 686    @property
 687    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 688        return engine_adapter.BigQueryEngineAdapter
 689
 690    @property
 691    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
 692        """The static connection kwargs for this connection"""
 693        import google.auth
 694        from google.api_core import client_info
 695        from google.oauth2 import credentials, service_account
 696
 697        if self.method == BigQueryConnectionMethod.OAUTH:
 698            creds, _ = google.auth.default(scopes=self.scopes)
 699        elif self.method == BigQueryConnectionMethod.SERVICE_ACCOUNT:
 700            creds = service_account.Credentials.from_service_account_file(
 701                self.keyfile, scopes=self.scopes
 702            )
 703        elif self.method == BigQueryConnectionMethod.SERVICE_ACCOUNT_JSON:
 704            creds = service_account.Credentials.from_service_account_info(
 705                self.keyfile_json, scopes=self.scopes
 706            )
 707        elif self.method == BigQueryConnectionMethod.OAUTH_SECRETS:
 708            creds = credentials.Credentials(
 709                token=self.token,
 710                refresh_token=self.refresh_token,
 711                client_id=self.client_id,
 712                client_secret=self.client_secret,
 713                token_uri=self.token_uri,
 714                scopes=self.scopes,
 715            )
 716        else:
 717            raise ConfigError("Invalid BigQuery Connection Method")
 718        client = google.cloud.bigquery.Client(
 719            project=self.execution_project or self.project,
 720            credentials=creds,
 721            location=self.location,
 722            client_info=client_info.ClientInfo(user_agent="sqlmesh"),
 723        )
 724
 725        return {
 726            "client": client,
 727        }
 728
 729    @property
 730    def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 731        return {
 732            k: v
 733            for k, v in self.dict().items()
 734            if k
 735            in {
 736                "job_creation_timeout_seconds",
 737                "job_execution_timeout_seconds",
 738                "job_retries",
 739                "job_retry_deadline_seconds",
 740                "priority",
 741                "maximum_bytes_billed",
 742            }
 743        }
 744
 745    @property
 746    def _connection_factory(self) -> t.Callable:
 747        from google.cloud.bigquery.dbapi import connect
 748
 749        return connect
 750
 751    def get_catalog(self) -> t.Optional[str]:
 752        return self.project
 753
 754
 755class GCPPostgresConnectionConfig(ConnectionConfig):
 756    """
 757    Postgres Connection Configuration for GCP.
 758
 759    Args:
 760        instance_connection_string: Connection name for the postgres instance.
 761        user: Postgres or IAM user's name
 762        password: The postgres user's password. Only needed when the user is a postgres user.
 763        enable_iam_auth: Set to True when user is an IAM user.
 764        db: Name of the db to connect to.
 765    """
 766
 767    instance_connection_string: str
 768    user: str
 769    password: t.Optional[str] = None
 770    enable_iam_auth: t.Optional[bool] = None
 771    db: str
 772    timeout: t.Optional[int] = None
 773
 774    driver: str = "pg8000"
 775    type_: Literal["gcp_postgres"] = Field(alias="type", default="gcp_postgres")
 776    concurrent_tasks: int = 4
 777    register_comments: bool = True
 778
 779    @model_validator(mode="before")
 780    @model_validator_v1_args
 781    def _validate_auth_method(
 782        cls, values: t.Dict[str, t.Optional[str]]
 783    ) -> t.Dict[str, t.Optional[str]]:
 784        password = values.get("password")
 785        enable_iam_auth = values.get("enable_iam_auth")
 786        if password and enable_iam_auth:
 787            raise ConfigError(
 788                "Invalid GCP Postgres connection configuration - both password and"
 789                " enable_iam_auth set. Use password when connecting to a postgres"
 790                " user and enable_iam_auth 'True' when connecting to an IAM user."
 791            )
 792        if not password and not enable_iam_auth:
 793            raise ConfigError(
 794                "GCP Postgres connection configuration requires either password set"
 795                " for a postgres user account or enable_iam_auth set to 'True'"
 796                " for an IAM user account."
 797            )
 798        return values
 799
 800    @property
 801    def _connection_kwargs_keys(self) -> t.Set[str]:
 802        return {
 803            "instance_connection_string",
 804            "driver",
 805            "user",
 806            "password",
 807            "db",
 808            "enable_iam_auth",
 809            "timeout",
 810        }
 811
 812    @property
 813    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 814        return engine_adapter.PostgresEngineAdapter
 815
 816    @property
 817    def _connection_factory(self) -> t.Callable:
 818        from google.cloud.sql.connector import Connector
 819
 820        return Connector().connect
 821
 822
 823class RedshiftConnectionConfig(ConnectionConfig):
 824    """
 825    Redshift Connection Configuration.
 826
 827    Arg Source: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/__init__.py#L146
 828    Note: A subset of properties were selected. Please open an issue/PR if you want to see more supported.
 829
 830    Args:
 831        user: The username to use for authentication with the Amazon Redshift cluster.
 832        password: The password to use for authentication with the Amazon Redshift cluster.
 833        database: The name of the database instance to connect to.
 834        host: The hostname of the Amazon Redshift cluster.
 835        port: The port number of the Amazon Redshift cluster. Default value is 5439.
 836        source_address: No description provided
 837        unix_sock: No description provided
 838        ssl: Is SSL enabled. Default value is ``True``. SSL must be enabled when authenticating using IAM.
 839        sslmode: The security of the connection to the Amazon Redshift cluster. 'verify-ca' and 'verify-full' are supported.
 840        timeout: The number of seconds before the connection to the server will timeout. By default there is no timeout.
 841        tcp_keepalive: Is `TCP keepalive <https://en.wikipedia.org/wiki/Keepalive#TCP_keepalive>`_ used. The default value is ``True``.
 842        application_name: Sets the application name. The default value is None.
 843        preferred_role: The IAM role preferred for the current connection.
 844        principal_arn: The ARN of the IAM entity (user or role) for which you are generating a policy.
 845        credentials_provider: The class name of the IdP that will be used for authenticating with the Amazon Redshift cluster.
 846        region: The AWS region where the Amazon Redshift cluster is located.
 847        cluster_identifier: The cluster identifier of the Amazon Redshift cluster.
 848        iam: If IAM authentication is enabled. Default value is False. IAM must be True when authenticating using an IdP.
 849        is_serverless: Redshift end-point is serverless or provisional. Default value false.
 850        serverless_acct_id: The account ID of the serverless. Default value None
 851        serverless_work_group: The name of work group for serverless end point. Default value None.
 852    """
 853
 854    user: t.Optional[str] = None
 855    password: t.Optional[str] = None
 856    database: t.Optional[str] = None
 857    host: t.Optional[str] = None
 858    port: t.Optional[int] = None
 859    source_address: t.Optional[str] = None
 860    unix_sock: t.Optional[str] = None
 861    ssl: t.Optional[bool] = None
 862    sslmode: t.Optional[str] = None
 863    timeout: t.Optional[int] = None
 864    tcp_keepalive: t.Optional[bool] = None
 865    application_name: t.Optional[str] = None
 866    preferred_role: t.Optional[str] = None
 867    principal_arn: t.Optional[str] = None
 868    credentials_provider: t.Optional[str] = None
 869    region: t.Optional[str] = None
 870    cluster_identifier: t.Optional[str] = None
 871    iam: t.Optional[bool] = None
 872    is_serverless: t.Optional[bool] = None
 873    serverless_acct_id: t.Optional[str] = None
 874    serverless_work_group: t.Optional[str] = None
 875
 876    concurrent_tasks: int = 4
 877    register_comments: bool = True
 878
 879    type_: Literal["redshift"] = Field(alias="type", default="redshift")
 880
 881    @property
 882    def _connection_kwargs_keys(self) -> t.Set[str]:
 883        return {
 884            "user",
 885            "password",
 886            "database",
 887            "host",
 888            "port",
 889            "source_address",
 890            "unix_sock",
 891            "ssl",
 892            "sslmode",
 893            "timeout",
 894            "tcp_keepalive",
 895            "application_name",
 896            "preferred_role",
 897            "principal_arn",
 898            "credentials_provider",
 899            "region",
 900            "cluster_identifier",
 901            "iam",
 902            "is_serverless",
 903            "serverless_acct_id",
 904            "serverless_work_group",
 905        }
 906
 907    @property
 908    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 909        return engine_adapter.RedshiftEngineAdapter
 910
 911    @property
 912    def _connection_factory(self) -> t.Callable:
 913        from redshift_connector import connect
 914
 915        return connect
 916
 917
 918class PostgresConnectionConfig(ConnectionConfig):
 919    host: str
 920    user: str
 921    password: str
 922    port: int
 923    database: str
 924    keepalives_idle: t.Optional[int] = None
 925    connect_timeout: int = 10
 926    role: t.Optional[str] = None
 927    sslmode: t.Optional[str] = None
 928
 929    concurrent_tasks: int = 4
 930    register_comments: bool = True
 931
 932    type_: Literal["postgres"] = Field(alias="type", default="postgres")
 933
 934    @property
 935    def _connection_kwargs_keys(self) -> t.Set[str]:
 936        return {
 937            "host",
 938            "user",
 939            "password",
 940            "port",
 941            "database",
 942            "keepalives_idle",
 943            "connect_timeout",
 944            "role",
 945            "sslmode",
 946        }
 947
 948    @property
 949    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 950        return engine_adapter.PostgresEngineAdapter
 951
 952    @property
 953    def _connection_factory(self) -> t.Callable:
 954        from psycopg2 import connect
 955
 956        return connect
 957
 958
 959class MySQLConnectionConfig(ConnectionConfig):
 960    host: str
 961    user: str
 962    password: str
 963    port: t.Optional[int] = None
 964    charset: t.Optional[str] = None
 965    ssl_disabled: t.Optional[bool] = None
 966
 967    concurrent_tasks: int = 4
 968    register_comments: bool = True
 969
 970    type_: Literal["mysql"] = Field(alias="type", default="mysql")
 971
 972    @property
 973    def _cursor_kwargs(self) -> t.Optional[t.Dict[str, t.Any]]:
 974        """Key-value arguments that will be passed during cursor construction."""
 975        return {"buffered": True}
 976
 977    @property
 978    def _connection_kwargs_keys(self) -> t.Set[str]:
 979        connection_keys = {
 980            "host",
 981            "user",
 982            "password",
 983            "port",
 984            "database",
 985        }
 986        if self.port is not None:
 987            connection_keys.add("port")
 988        if self.charset is not None:
 989            connection_keys.add("charset")
 990        if self.ssl_disabled is not None:
 991            connection_keys.add("ssl_disabled")
 992        return connection_keys
 993
 994    @property
 995    def _engine_adapter(self) -> t.Type[EngineAdapter]:
 996        return engine_adapter.MySQLEngineAdapter
 997
 998    @property
 999    def _connection_factory(self) -> t.Callable:
1000        from mysql.connector import connect
1001
1002        return connect
1003
1004
1005class MSSQLConnectionConfig(ConnectionConfig):
1006    host: str
1007    user: t.Optional[str] = None
1008    password: t.Optional[str] = None
1009    database: t.Optional[str] = ""
1010    timeout: t.Optional[int] = 0
1011    login_timeout: t.Optional[int] = 60
1012    charset: t.Optional[str] = "UTF-8"
1013    appname: t.Optional[str] = None
1014    port: t.Optional[int] = 1433
1015    conn_properties: t.Optional[t.Union[t.Iterable[str], str]] = None
1016    autocommit: t.Optional[bool] = False
1017    tds_version: t.Optional[str] = None
1018
1019    concurrent_tasks: int = 4
1020    register_comments: bool = True
1021
1022    type_: Literal["mssql"] = Field(alias="type", default="mssql")
1023
1024    @property
1025    def _connection_kwargs_keys(self) -> t.Set[str]:
1026        return {
1027            "host",
1028            "user",
1029            "password",
1030            "database",
1031            "timeout",
1032            "login_timeout",
1033            "charset",
1034            "appname",
1035            "port",
1036            "conn_properties",
1037            "autocommit",
1038            "tds_version",
1039        }
1040
1041    @property
1042    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1043        return engine_adapter.MSSQLEngineAdapter
1044
1045    @property
1046    def _connection_factory(self) -> t.Callable:
1047        import pymssql
1048
1049        return pymssql.connect
1050
1051
1052class SparkConnectionConfig(ConnectionConfig):
1053    """
1054    Vanilla Spark Connection Configuration. Use `DatabricksConnectionConfig` for Databricks.
1055    """
1056
1057    config_dir: t.Optional[str] = None
1058    catalog: t.Optional[str] = None
1059    config: t.Dict[str, t.Any] = {}
1060
1061    concurrent_tasks: int = 4
1062    register_comments: bool = True
1063
1064    type_: Literal["spark"] = Field(alias="type", default="spark")
1065
1066    @property
1067    def _connection_kwargs_keys(self) -> t.Set[str]:
1068        return {
1069            "catalog",
1070        }
1071
1072    @property
1073    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1074        return engine_adapter.SparkEngineAdapter
1075
1076    @property
1077    def _connection_factory(self) -> t.Callable:
1078        from sqlmesh.engines.spark.db_api.spark_session import connection
1079
1080        return connection
1081
1082    @property
1083    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
1084        from pyspark.conf import SparkConf
1085        from pyspark.sql import SparkSession
1086
1087        spark_config = SparkConf()
1088        if self.config:
1089            for k, v in self.config.items():
1090                spark_config.set(k, v)
1091
1092        if self.config_dir:
1093            os.environ["SPARK_CONF_DIR"] = self.config_dir
1094        return {
1095            "spark": SparkSession.builder.config(conf=spark_config)
1096            .enableHiveSupport()
1097            .getOrCreate(),
1098        }
1099
1100
1101class TrinoAuthenticationMethod(str, Enum):
1102    NO_AUTH = "no-auth"
1103    BASIC = "basic"
1104    LDAP = "ldap"
1105    KERBEROS = "kerberos"
1106    JWT = "jwt"
1107    CERTIFICATE = "certificate"
1108    OAUTH = "oauth"
1109
1110    @property
1111    def is_no_auth(self) -> bool:
1112        return self == self.NO_AUTH
1113
1114    @property
1115    def is_basic(self) -> bool:
1116        return self == self.BASIC
1117
1118    @property
1119    def is_ldap(self) -> bool:
1120        return self == self.LDAP
1121
1122    @property
1123    def is_kerberos(self) -> bool:
1124        return self == self.KERBEROS
1125
1126    @property
1127    def is_jwt(self) -> bool:
1128        return self == self.JWT
1129
1130    @property
1131    def is_certificate(self) -> bool:
1132        return self == self.CERTIFICATE
1133
1134    @property
1135    def is_oauth(self) -> bool:
1136        return self == self.OAUTH
1137
1138
1139class TrinoConnectionConfig(ConnectionConfig):
1140    method: TrinoAuthenticationMethod = TrinoAuthenticationMethod.NO_AUTH
1141    host: str
1142    user: str
1143    catalog: str
1144    port: t.Optional[int] = None
1145    http_scheme: Literal["http", "https"] = "https"
1146    # General Optional
1147    roles: t.Optional[t.Dict[str, str]] = None
1148    http_headers: t.Optional[t.Dict[str, str]] = None
1149    session_properties: t.Optional[t.Dict[str, str]] = None
1150    retries: int = 3
1151    timezone: t.Optional[str] = None
1152    # Basic/LDAP
1153    password: t.Optional[str] = None
1154    # LDAP
1155    impersonation_user: t.Optional[str] = None
1156    # Kerberos
1157    keytab: t.Optional[str] = None
1158    krb5_config: t.Optional[str] = None
1159    principal: t.Optional[str] = None
1160    service_name: str = "trino"
1161    hostname_override: t.Optional[str] = None
1162    mutual_authentication: bool = False
1163    force_preemptive: bool = False
1164    sanitize_mutual_error_response: bool = True
1165    delegate: bool = False
1166    # JWT
1167    jwt_token: t.Optional[str] = None
1168    # Certificate
1169    client_certificate: t.Optional[str] = None
1170    client_private_key: t.Optional[str] = None
1171    cert: t.Optional[str] = None
1172
1173    concurrent_tasks: int = 4
1174    register_comments: bool = True
1175
1176    type_: Literal["trino"] = Field(alias="type", default="trino")
1177
1178    @model_validator(mode="after")
1179    @model_validator_v1_args
1180    def _root_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
1181        port = values.get("port")
1182        if (
1183            values["http_scheme"] == "http"
1184            and not values["method"].is_no_auth
1185            and not values["method"].is_basic
1186        ):
1187            raise ConfigError("HTTP scheme can only be used with no-auth or basic method")
1188        if port is None:
1189            values["port"] = 80 if values["http_scheme"] == "http" else 443
1190        if (values["method"].is_ldap or values["method"].is_basic) and (
1191            not values["password"] or not values["user"]
1192        ):
1193            raise ConfigError(
1194                f"Username and Password must be provided if using {values['method'].value} authentication"
1195            )
1196        if values["method"].is_kerberos and (
1197            not values["principal"] or not values["keytab"] or not values["krb5_config"]
1198        ):
1199            raise ConfigError(
1200                "Kerberos requires the following fields: principal, keytab, and krb5_config"
1201            )
1202        if values["method"].is_jwt and not values["jwt_token"]:
1203            raise ConfigError("JWT requires `jwt_token` to be set")
1204        if values["method"].is_certificate and (
1205            not values["cert"]
1206            or not values["client_certificate"]
1207            or not values["client_private_key"]
1208        ):
1209            raise ConfigError(
1210                "Certificate requires the following fields: cert, client_certificate, and client_private_key"
1211            )
1212        return values
1213
1214    @property
1215    def _connection_kwargs_keys(self) -> t.Set[str]:
1216        kwargs = {
1217            "host",
1218            "port",
1219            "catalog",
1220            "roles",
1221            "http_scheme",
1222            "http_headers",
1223            "session_properties",
1224            "timezone",
1225        }
1226        return kwargs
1227
1228    @property
1229    def _engine_adapter(self) -> t.Type[EngineAdapter]:
1230        return engine_adapter.TrinoEngineAdapter
1231
1232    @property
1233    def _connection_factory(self) -> t.Callable:
1234        from trino.dbapi import connect
1235
1236        return connect
1237
1238    @property
1239    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
1240        from trino.auth import (
1241            BasicAuthentication,
1242            CertificateAuthentication,
1243            JWTAuthentication,
1244            KerberosAuthentication,
1245            OAuth2Authentication,
1246        )
1247
1248        if self.method.is_basic or self.method.is_ldap:
1249            auth = BasicAuthentication(self.user, self.password)
1250        elif self.method.is_kerberos:
1251            if self.keytab:
1252                os.environ["KRB5_CLIENT_KTNAME"] = self.keytab
1253            auth = KerberosAuthentication(
1254                config=self.krb5_config,
1255                service_name=self.service_name,
1256                principal=self.principal,
1257                mutual_authentication=self.mutual_authentication,
1258                ca_bundle=self.cert,
1259                force_preemptive=self.force_preemptive,
1260                hostname_override=self.hostname_override,
1261                sanitize_mutual_error_response=self.sanitize_mutual_error_response,
1262                delegate=self.delegate,
1263            )
1264        elif self.method.is_oauth:
1265            auth = OAuth2Authentication()
1266        elif self.method.is_jwt:
1267            auth = JWTAuthentication(self.jwt_token)
1268        elif self.method.is_certificate:
1269            auth = CertificateAuthentication(self.client_certificate, self.client_private_key)
1270        else:
1271            auth = None
1272
1273        return {
1274            "auth": auth,
1275            "user": self.impersonation_user or self.user,
1276            "max_attempts": self.retries,
1277            "verify": self.cert,
1278            "source": "sqlmesh",
1279        }
1280
1281
1282CONNECTION_CONFIG_TO_TYPE = {
1283    # Map all subclasses of ConnectionConfig to the value of their `type_` field.
1284    tpe.all_field_infos()["type_"].default: tpe
1285    for tpe in subclasses(
1286        __name__, ConnectionConfig, exclude=(ConnectionConfig, BaseDuckDBConnectionConfig)
1287    )
1288}
1289
1290
1291def parse_connection_config(v: t.Dict[str, t.Any]) -> ConnectionConfig:
1292    if "type" not in v:
1293        raise ConfigError("Missing connection type.")
1294
1295    connection_type = v["type"]
1296    if connection_type not in CONNECTION_CONFIG_TO_TYPE:
1297        raise ConfigError(f"Unknown connection type '{connection_type}'.")
1298
1299    return CONNECTION_CONFIG_TO_TYPE[connection_type](**v)
1300
1301
1302def _connection_config_validator(
1303    cls: t.Type, v: ConnectionConfig | t.Dict[str, t.Any] | None
1304) -> ConnectionConfig | None:
1305    if v is None or isinstance(v, ConnectionConfig):
1306        return v
1307    return parse_connection_config(v)
1308
1309
1310connection_config_validator = field_validator(
1311    "connection",
1312    "state_connection",
1313    "test_connection",
1314    "default_connection",
1315    "default_test_connection",
1316    mode="before",
1317    check_fields=False,
1318)(_connection_config_validator)
1319
1320
1321if t.TYPE_CHECKING:
1322    # TypeAlias hasn't been introduced until Python 3.10 which means that we can't use it
1323    # outside the TYPE_CHECKING guard.
1324    SerializableConnectionConfig: t.TypeAlias = ConnectionConfig  # type: ignore
1325elif PYDANTIC_MAJOR_VERSION >= 2:
1326    import pydantic
1327
1328    # Workaround for https://docs.pydantic.dev/latest/concepts/serialization/#serializing-with-duck-typing
1329    SerializableConnectionConfig = pydantic.SerializeAsAny[ConnectionConfig]  # type: ignore
1330else:
1331    SerializableConnectionConfig = ConnectionConfig  # type: ignore