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