EngineAdapter
Engine adapters are how SQLMesh connects and interacts with various data stores. They allow SQLMesh to generalize its functionality to different engines that have Python Database API 2.0-compliant connections. Rather than executing queries directly against your data stores, SQLMesh components such as the SnapshotEvaluator delegate them to engine adapters so these components can be engine-agnostic.
1""" 2# EngineAdapter 3 4Engine adapters are how SQLMesh connects and interacts with various data stores. They allow SQLMesh to 5generalize its functionality to different engines that have Python Database API 2.0-compliant 6connections. Rather than executing queries directly against your data stores, SQLMesh components such as 7the SnapshotEvaluator delegate them to engine adapters so these components can be engine-agnostic. 8""" 9 10from __future__ import annotations 11 12import contextlib 13import itertools 14import logging 15import sys 16import typing as t 17from functools import cached_property, partial 18 19from sqlglot import Dialect, exp 20from sqlglot.errors import ErrorLevel 21from sqlglot.helper import ensure_list, seq_get 22from sqlglot.optimizer.qualify_columns import quote_identifiers 23 24from sqlmesh.core.dialect import ( 25 add_table, 26 schema_, 27 select_from_values_for_batch_range, 28 to_schema, 29) 30from sqlmesh.core.engine_adapter.shared import ( 31 CatalogSupport, 32 CommentCreationTable, 33 CommentCreationView, 34 DataObject, 35 DataObjectType, 36 EngineRunMode, 37 InsertOverwriteStrategy, 38 SourceQuery, 39 set_catalog, 40) 41from sqlmesh.core.model.kind import TimeColumn 42from sqlmesh.core.schema_diff import SchemaDiffer, TableAlterOperation 43from sqlmesh.core.snapshot.execution_tracker import QueryExecutionTracker 44from sqlmesh.utils import ( 45 CorrelationId, 46 columns_to_types_all_known, 47 random_id, 48 get_source_columns_to_types, 49) 50from sqlmesh.utils.connection_pool import ConnectionPool, create_connection_pool 51from sqlmesh.utils.date import TimeLike, make_inclusive, to_time_column 52from sqlmesh.utils.errors import ( 53 MissingDefaultCatalogError, 54 SQLMeshError, 55 UnsupportedCatalogOperationError, 56) 57from sqlmesh.utils.pandas import columns_to_types_from_df 58 59if t.TYPE_CHECKING: 60 import pandas as pd 61 62 from sqlmesh.core._typing import SchemaName, SessionProperties, TableName 63 from sqlmesh.core.engine_adapter._typing import ( 64 DF, 65 BigframeSession, 66 GrantsConfig, 67 PySparkDataFrame, 68 PySparkSession, 69 Query, 70 QueryOrDF, 71 SnowparkSession, 72 ) 73 from sqlmesh.core.node import IntervalUnit 74 75logger = logging.getLogger(__name__) 76 77MERGE_TARGET_ALIAS = "__MERGE_TARGET__" 78MERGE_SOURCE_ALIAS = "__MERGE_SOURCE__" 79 80KEY_FOR_CREATABLE_TYPE = "CREATABLE_TYPE" 81 82 83@set_catalog() 84class EngineAdapter: 85 """Base class wrapping a Database API compliant connection. 86 87 The EngineAdapter is an easily-subclassable interface that interacts 88 with the underlying engine and data store. 89 90 Args: 91 connection_factory_or_pool: a callable which produces a new Database API-compliant 92 connection on every call. 93 dialect: The dialect with which this adapter is associated. 94 multithreaded: Indicates whether this adapter will be used by more than one thread. 95 """ 96 97 DIALECT = "" 98 DEFAULT_BATCH_SIZE = 10000 99 DATA_OBJECT_FILTER_BATCH_SIZE = 4000 100 SUPPORTS_TRANSACTIONS = True 101 SUPPORTS_INDEXES = False 102 COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_CTAS 103 COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_AND_COMMANDS 104 MAX_TABLE_COMMENT_LENGTH: t.Optional[int] = None 105 MAX_COLUMN_COMMENT_LENGTH: t.Optional[int] = None 106 INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT 107 SUPPORTS_MATERIALIZED_VIEWS = False 108 SUPPORTS_MATERIALIZED_VIEW_SCHEMA = False 109 SUPPORTS_VIEW_SCHEMA = True 110 SUPPORTS_CLONING = False 111 SUPPORTS_MANAGED_MODELS = False 112 SUPPORTS_CREATE_DROP_CATALOG = False 113 SUPPORTED_DROP_CASCADE_OBJECT_KINDS: t.List[str] = [] 114 SCHEMA_DIFFER_KWARGS: t.Dict[str, t.Any] = {} 115 SUPPORTS_TUPLE_IN = True 116 HAS_VIEW_BINDING = False 117 RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = True 118 SUPPORTS_REPLACE_TABLE = True 119 SUPPORTS_GRANTS = False 120 DEFAULT_CATALOG_TYPE = DIALECT 121 QUOTE_IDENTIFIERS_IN_VIEWS = True 122 MAX_IDENTIFIER_LENGTH: t.Optional[int] = None 123 ATTACH_CORRELATION_ID = True 124 SUPPORTS_QUERY_EXECUTION_TRACKING = False 125 SUPPORTS_METADATA_TABLE_LAST_MODIFIED_TS = False 126 RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES: t.FrozenSet[str] = frozenset() 127 """Physical property keys whose values may contain logical model references that 128 should be resolved to physical table names during property rendering. Engines that 129 need such resolution (e.g. StarRocks' excluded_trigger_tables) override this set.""" 130 131 def __init__( 132 self, 133 connection_factory_or_pool: t.Union[t.Callable[[], t.Any], ConnectionPool], 134 dialect: str = "", 135 sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None, 136 multithreaded: bool = False, 137 cursor_init: t.Optional[t.Callable[[t.Any], None]] = None, 138 default_catalog: t.Optional[str] = None, 139 execute_log_level: int = logging.DEBUG, 140 register_comments: bool = True, 141 pre_ping: bool = False, 142 pretty_sql: bool = False, 143 shared_connection: bool = False, 144 correlation_id: t.Optional[CorrelationId] = None, 145 schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None, 146 query_execution_tracker: t.Optional[QueryExecutionTracker] = None, 147 **kwargs: t.Any, 148 ): 149 self.dialect = dialect.lower() or self.DIALECT 150 self._connection_pool = ( 151 connection_factory_or_pool 152 if isinstance(connection_factory_or_pool, ConnectionPool) 153 else create_connection_pool( 154 connection_factory_or_pool, 155 multithreaded, 156 shared_connection=shared_connection, 157 cursor_init=cursor_init, 158 ) 159 ) 160 self._sql_gen_kwargs = sql_gen_kwargs or {} 161 self._default_catalog = default_catalog 162 self._execute_log_level = execute_log_level 163 self._extra_config = kwargs 164 self._register_comments = register_comments 165 self._pre_ping = pre_ping 166 self._pretty_sql = pretty_sql 167 self._multithreaded = multithreaded 168 self.correlation_id = correlation_id 169 self._schema_differ_overrides = schema_differ_overrides 170 self._query_execution_tracker = query_execution_tracker 171 self._data_object_cache: t.Dict[str, t.Optional[DataObject]] = {} 172 173 def with_settings(self, **kwargs: t.Any) -> EngineAdapter: 174 extra_kwargs = { 175 "null_connection": True, 176 "execute_log_level": kwargs.pop("execute_log_level", self._execute_log_level), 177 "correlation_id": kwargs.pop("correlation_id", self.correlation_id), 178 "query_execution_tracker": kwargs.pop( 179 "query_execution_tracker", self._query_execution_tracker 180 ), 181 **self._extra_config, 182 **kwargs, 183 } 184 185 adapter = self.__class__( 186 self._connection_pool, 187 dialect=self.dialect, 188 sql_gen_kwargs=self._sql_gen_kwargs, 189 default_catalog=self._default_catalog, 190 register_comments=self._register_comments, 191 multithreaded=self._multithreaded, 192 pretty_sql=self._pretty_sql, 193 **extra_kwargs, 194 ) 195 196 return adapter 197 198 @property 199 def cursor(self) -> t.Any: 200 return self._connection_pool.get_cursor() 201 202 @property 203 def connection(self) -> t.Any: 204 return self._connection_pool.get() 205 206 @property 207 def spark(self) -> t.Optional[PySparkSession]: 208 return None 209 210 @property 211 def snowpark(self) -> t.Optional[SnowparkSession]: 212 return None 213 214 @property 215 def bigframe(self) -> t.Optional[BigframeSession]: 216 return None 217 218 @property 219 def comments_enabled(self) -> bool: 220 return self._register_comments and self.COMMENT_CREATION_TABLE.is_supported 221 222 @property 223 def catalog_support(self) -> CatalogSupport: 224 return CatalogSupport.UNSUPPORTED 225 226 def supports_virtual_catalog(self) -> bool: 227 """Return True if this adapter can accept a virtual catalog for multi-gateway nesting alignment. 228 229 When a project mixes catalog-aware gateways (e.g. DuckDB) with catalog-unsupported gateways 230 (e.g. ClickHouse), all adapters need a uniform 3-level FQN so MappingSchema nesting stays 231 consistent. Adapters that return True here opt in to receiving an injected virtual catalog 232 via inject_virtual_catalog(), which causes the set_catalog decorator to strip the catalog 233 from DDL expressions rather than raising UnsupportedCatalogOperationError. 234 """ 235 return False 236 237 def inject_virtual_catalog(self, gateway: str) -> None: 238 """Inject a gateway name to configure the adapter's virtual catalog. 239 240 The adapter determines the final catalog name from the gateway name (e.g. ClickHouse 241 wraps it as __{gateway}__). Only call this on adapters that return True from 242 supports_virtual_catalog(). After injection, catalog_support should return 243 SINGLE_CATALOG_ONLY so the set_catalog decorator strips the virtual catalog from DDL 244 expressions instead of raising an error. 245 """ 246 raise NotImplementedError( 247 f"{self.dialect} does not support virtual catalog injection. " 248 "Override supports_virtual_catalog() to return True and implement inject_virtual_catalog()." 249 ) 250 251 @cached_property 252 def schema_differ(self) -> SchemaDiffer: 253 return SchemaDiffer( 254 **{ 255 **self.SCHEMA_DIFFER_KWARGS, 256 **(self._schema_differ_overrides or {}), 257 } 258 ) 259 260 @property 261 def _catalog_type_overrides(self) -> t.Dict[str, str]: 262 return self._extra_config.get("catalog_type_overrides") or {} 263 264 @classmethod 265 def _casted_columns( 266 cls, 267 target_columns_to_types: t.Dict[str, exp.DataType], 268 source_columns: t.Optional[t.List[str]] = None, 269 ) -> t.List[exp.Expr]: 270 source_columns_lookup = set(source_columns or target_columns_to_types) 271 return [ 272 exp.alias_( 273 exp.cast( 274 exp.column(column, quoted=True) 275 if column in source_columns_lookup 276 else exp.Null(), 277 to=kind, 278 ), 279 column, 280 copy=False, 281 quoted=True, 282 ) 283 for column, kind in target_columns_to_types.items() 284 ] 285 286 @property 287 def default_catalog(self) -> t.Optional[str]: 288 if self.catalog_support.is_unsupported: 289 return None 290 default_catalog = self._default_catalog or self.get_current_catalog() 291 if not default_catalog: 292 raise MissingDefaultCatalogError( 293 "Could not determine a default catalog despite it being supported." 294 ) 295 return default_catalog 296 297 @property 298 def engine_run_mode(self) -> EngineRunMode: 299 return EngineRunMode.SINGLE_MODE_ENGINE 300 301 def _get_source_queries( 302 self, 303 query_or_df: QueryOrDF, 304 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 305 target_table: TableName, 306 *, 307 batch_size: t.Optional[int] = None, 308 source_columns: t.Optional[t.List[str]] = None, 309 ) -> t.List[SourceQuery]: 310 import pandas as pd 311 312 batch_size = self.DEFAULT_BATCH_SIZE if batch_size is None else batch_size 313 if isinstance(query_or_df, exp.Query): 314 query_factory = lambda: query_or_df 315 if source_columns: 316 source_columns_lookup = set(source_columns) 317 if not target_columns_to_types: 318 raise SQLMeshError("columns_to_types must be set if source_columns is set") 319 if not set(target_columns_to_types).issubset(source_columns_lookup): 320 select_columns = [ 321 exp.column(c, quoted=True) 322 if c in source_columns_lookup 323 else exp.cast(exp.Null(), target_columns_to_types[c], copy=False).as_( 324 c, copy=False, quoted=True 325 ) 326 for c in target_columns_to_types 327 ] 328 query_factory = lambda: ( 329 exp.Select() 330 .select(*select_columns) 331 .from_(query_or_df.subquery("select_source_columns")) 332 ) 333 return [SourceQuery(query_factory=query_factory)] # type: ignore 334 335 if not target_columns_to_types: 336 raise SQLMeshError( 337 "It is expected that if a DataFrame is passed in then columns_to_types is set" 338 ) 339 340 if isinstance(query_or_df, pd.DataFrame) and query_or_df.empty: 341 raise SQLMeshError( 342 "Cannot construct source query from an empty DataFrame. This error is commonly " 343 "related to Python models that produce no data. For such models, consider yielding " 344 "from an empty generator if the resulting set is empty, i.e. use `yield from ()`." 345 ) 346 347 return self._df_to_source_queries( 348 query_or_df, 349 target_columns_to_types, 350 batch_size, 351 target_table=target_table, 352 source_columns=source_columns, 353 ) 354 355 def _df_to_source_queries( 356 self, 357 df: DF, 358 target_columns_to_types: t.Dict[str, exp.DataType], 359 batch_size: int, 360 target_table: TableName, 361 source_columns: t.Optional[t.List[str]] = None, 362 ) -> t.List[SourceQuery]: 363 import pandas as pd 364 365 assert isinstance(df, pd.DataFrame) 366 num_rows = len(df.index) 367 batch_size = sys.maxsize if batch_size == 0 else batch_size 368 369 # we need to ensure that the order of the columns in columns_to_types columns matches the order of the values 370 # they can differ if a user specifies columns() on a python model in a different order than what's in the DataFrame's emitted by that model 371 df = df[list(source_columns or target_columns_to_types)] 372 values = list(df.itertuples(index=False, name=None)) 373 374 return [ 375 SourceQuery( 376 query_factory=partial( 377 self._values_to_sql, 378 values=values, # type: ignore 379 target_columns_to_types=target_columns_to_types, 380 batch_start=i, 381 batch_end=min(i + batch_size, num_rows), 382 source_columns=source_columns, 383 ), 384 ) 385 for i in range(0, num_rows, batch_size) 386 ] 387 388 def _get_source_queries_and_columns_to_types( 389 self, 390 query_or_df: QueryOrDF, 391 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 392 target_table: TableName, 393 *, 394 batch_size: t.Optional[int] = None, 395 source_columns: t.Optional[t.List[str]] = None, 396 ) -> t.Tuple[t.List[SourceQuery], t.Optional[t.Dict[str, exp.DataType]]]: 397 target_columns_to_types, source_columns = self._columns_to_types( 398 query_or_df, target_columns_to_types, source_columns 399 ) 400 source_queries = self._get_source_queries( 401 query_or_df, 402 target_columns_to_types, 403 target_table=target_table, 404 batch_size=batch_size, 405 source_columns=source_columns, 406 ) 407 return source_queries, target_columns_to_types 408 409 @t.overload 410 def _columns_to_types( 411 self, 412 query_or_df: DF, 413 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 414 source_columns: t.Optional[t.List[str]] = None, 415 ) -> t.Tuple[t.Dict[str, exp.DataType], t.List[str]]: ... 416 417 @t.overload 418 def _columns_to_types( 419 self, 420 query_or_df: Query, 421 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 422 source_columns: t.Optional[t.List[str]] = None, 423 ) -> t.Tuple[t.Optional[t.Dict[str, exp.DataType]], t.Optional[t.List[str]]]: ... 424 425 def _columns_to_types( 426 self, 427 query_or_df: QueryOrDF, 428 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 429 source_columns: t.Optional[t.List[str]] = None, 430 ) -> t.Tuple[t.Optional[t.Dict[str, exp.DataType]], t.Optional[t.List[str]]]: 431 import pandas as pd 432 433 if not target_columns_to_types and isinstance(query_or_df, pd.DataFrame): 434 target_columns_to_types = columns_to_types_from_df(t.cast(pd.DataFrame, query_or_df)) 435 if not source_columns and target_columns_to_types: 436 source_columns = list(target_columns_to_types) 437 # source columns should only contain columns that are defined in the target. If there are extras then 438 # that means they are intended to be ignored and will be excluded 439 source_columns = ( 440 [x for x in source_columns if x in target_columns_to_types] 441 if source_columns and target_columns_to_types 442 else None 443 ) 444 return target_columns_to_types, source_columns 445 446 def recycle(self) -> None: 447 """Closes all open connections and releases all allocated resources associated with any thread 448 except the calling one.""" 449 self._connection_pool.close_all(exclude_calling_thread=True) 450 451 def close(self) -> t.Any: 452 """Closes all open connections and releases all allocated resources.""" 453 self._connection_pool.close_all() 454 455 def get_current_catalog(self) -> t.Optional[str]: 456 """Returns the catalog name of the current connection.""" 457 raise NotImplementedError() 458 459 def set_current_catalog(self, catalog: str) -> None: 460 """Sets the catalog name of the current connection.""" 461 raise NotImplementedError() 462 463 def get_catalog_type(self, catalog: t.Optional[str]) -> str: 464 """Intended to be overridden for data virtualization systems like Trino that, 465 depending on the target catalog, require slightly different properties to be set when creating / updating tables 466 """ 467 if self.catalog_support.is_unsupported: 468 raise UnsupportedCatalogOperationError( 469 f"{self.dialect} does not support catalogs and a catalog was provided: {catalog}" 470 ) 471 return ( 472 self._catalog_type_overrides.get(catalog, self.DEFAULT_CATALOG_TYPE) 473 if catalog 474 else self.DEFAULT_CATALOG_TYPE 475 ) 476 477 def get_catalog_type_from_table(self, table: TableName) -> str: 478 """Get the catalog type from a table name if it has a catalog specified, otherwise return the current catalog type""" 479 catalog = exp.to_table(table).catalog or self.get_current_catalog() 480 return self.get_catalog_type(catalog) 481 482 @property 483 def current_catalog_type(self) -> str: 484 # `get_catalog_type_from_table` should be used over this property. Reason is that the table that is the target 485 # of the operation is what matters and not the catalog type of the connection. 486 # This still remains for legacy reasons and should be refactored out. 487 return self.get_catalog_type(self.get_current_catalog()) 488 489 def replace_query( 490 self, 491 table_name: TableName, 492 query_or_df: QueryOrDF, 493 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 494 table_description: t.Optional[str] = None, 495 column_descriptions: t.Optional[t.Dict[str, str]] = None, 496 source_columns: t.Optional[t.List[str]] = None, 497 supports_replace_table_override: t.Optional[bool] = None, 498 **kwargs: t.Any, 499 ) -> None: 500 """Replaces an existing table with a query. 501 502 For partition based engines (hive, spark), insert override is used. For other systems, create or replace is used. 503 504 Args: 505 table_name: The name of the table (eg. prod.table) 506 query_or_df: The SQL query to run or a dataframe. 507 target_columns_to_types: Only used if a dataframe is provided. A mapping between the column name and its data type. 508 Expected to be ordered to match the order of values in the dataframe. 509 kwargs: Optional create table properties. 510 """ 511 target_table = exp.to_table(table_name) 512 513 target_data_object = self.get_data_object(target_table) 514 table_exists = target_data_object is not None 515 if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE): 516 table_exists = False 517 518 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 519 query_or_df, 520 target_columns_to_types, 521 target_table=target_table, 522 source_columns=source_columns, 523 ) 524 if not target_columns_to_types and table_exists: 525 target_columns_to_types = self.columns(target_table) 526 query = source_queries[0].query_factory() 527 self_referencing = any( 528 quote_identifiers(table) == quote_identifiers(target_table) 529 for table in query.find_all(exp.Table) 530 ) 531 # If a query references itself then it must have a table created regardless of approach used. 532 if self_referencing: 533 if not target_columns_to_types: 534 raise SQLMeshError( 535 f"Cannot create a self-referencing table {target_table.sql(dialect=self.dialect)} without knowing the column types. " 536 "Try casting the columns to an expected type or defining the columns in the model metadata. " 537 ) 538 self._create_table_from_columns( 539 target_table, 540 target_columns_to_types, 541 exists=True, 542 table_description=table_description, 543 column_descriptions=column_descriptions, 544 **kwargs, 545 ) 546 # All engines support `CREATE TABLE AS` so we use that if the table doesn't already exist and we 547 # use `CREATE OR REPLACE TABLE AS` if the engine supports it 548 supports_replace_table = ( 549 self.SUPPORTS_REPLACE_TABLE 550 if supports_replace_table_override is None 551 else supports_replace_table_override 552 ) 553 if supports_replace_table or not table_exists: 554 return self._create_table_from_source_queries( 555 target_table, 556 source_queries, 557 target_columns_to_types, 558 replace=supports_replace_table, 559 table_description=table_description, 560 column_descriptions=column_descriptions, 561 **kwargs, 562 ) 563 if self_referencing: 564 assert target_columns_to_types is not None 565 with self.temp_table( 566 self._select_columns(target_columns_to_types).from_(target_table), 567 name=target_table, 568 target_columns_to_types=target_columns_to_types, 569 **kwargs, 570 ) as temp_table: 571 for source_query in source_queries: 572 source_query.add_transform( 573 lambda node: ( # type: ignore 574 temp_table # type: ignore 575 if isinstance(node, exp.Table) 576 and quote_identifiers(node) == quote_identifiers(target_table) 577 else node 578 ) 579 ) 580 return self._insert_overwrite_by_condition( 581 target_table, 582 source_queries, 583 target_columns_to_types, 584 **kwargs, 585 ) 586 return self._insert_overwrite_by_condition( 587 target_table, 588 source_queries, 589 target_columns_to_types, 590 **kwargs, 591 ) 592 593 def create_index( 594 self, 595 table_name: TableName, 596 index_name: str, 597 columns: t.Tuple[str, ...], 598 exists: bool = True, 599 ) -> None: 600 """Creates a new index for the given table if supported 601 602 Args: 603 table_name: The name of the target table. 604 index_name: The name of the index. 605 columns: The list of columns that constitute the index. 606 exists: Indicates whether to include the IF NOT EXISTS check. 607 """ 608 if not self.SUPPORTS_INDEXES: 609 return 610 611 expression = exp.Create( 612 this=exp.Index( 613 this=exp.to_identifier(index_name), 614 table=exp.to_table(table_name), 615 params=exp.IndexParameters(columns=[exp.to_column(c) for c in columns]), 616 ), 617 kind="INDEX", 618 exists=exists, 619 ) 620 self.execute(expression) 621 622 def _pop_creatable_type_from_properties( 623 self, 624 properties: t.Dict[str, exp.Expr], 625 ) -> t.Optional[exp.Property]: 626 """Pop out the creatable_type from the properties dictionary (if exists (return it/remove it) else return none). 627 It also checks that none of the expressions are MATERIALIZE as that conflicts with the `materialize` parameter. 628 """ 629 for key in list(properties.keys()): 630 upper_key = key.upper() 631 if upper_key == KEY_FOR_CREATABLE_TYPE: 632 value = properties.pop(key).name 633 parsed_properties = exp.maybe_parse( 634 value, into=exp.Properties, dialect=self.dialect 635 ) 636 property, *others = parsed_properties.expressions 637 if others: 638 # Multiple properties are unsupported today, can look into it in the future if needed 639 raise SQLMeshError( 640 f"Invalid creatable_type value with multiple properties: {value}" 641 ) 642 if isinstance(property, exp.MaterializedProperty): 643 raise SQLMeshError( 644 f"Cannot use {value} as a creatable_type as it conflicts with the `materialize` parameter." 645 ) 646 return property 647 return None 648 649 def create_table( 650 self, 651 table_name: TableName, 652 target_columns_to_types: t.Dict[str, exp.DataType], 653 primary_key: t.Optional[t.Tuple[str, ...]] = None, 654 exists: bool = True, 655 table_description: t.Optional[str] = None, 656 column_descriptions: t.Optional[t.Dict[str, str]] = None, 657 **kwargs: t.Any, 658 ) -> None: 659 """Create a table using a DDL statement 660 661 Args: 662 table_name: The name of the table to create. Can be fully qualified or just table name. 663 target_columns_to_types: A mapping between the column name and its data type. 664 primary_key: Determines the table primary key. 665 exists: Indicates whether to include the IF NOT EXISTS check. 666 table_description: Optional table description from MODEL DDL. 667 column_descriptions: Optional column descriptions from model query. 668 kwargs: Optional create table properties. 669 """ 670 self._create_table_from_columns( 671 table_name, 672 target_columns_to_types, 673 primary_key, 674 exists, 675 table_description, 676 column_descriptions, 677 **kwargs, 678 ) 679 680 def create_managed_table( 681 self, 682 table_name: TableName, 683 query: Query, 684 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 685 partitioned_by: t.Optional[t.List[exp.Expr]] = None, 686 clustered_by: t.Optional[t.List[exp.Expr]] = None, 687 table_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 688 table_description: t.Optional[str] = None, 689 column_descriptions: t.Optional[t.Dict[str, str]] = None, 690 source_columns: t.Optional[t.List[str]] = None, 691 **kwargs: t.Any, 692 ) -> None: 693 """Create a managed table using a query. 694 695 "Managed" means that once the table is created, the data is kept up to date by the underlying database engine and not SQLMesh. 696 697 Args: 698 table_name: The name of the table to create. Can be fully qualified or just table name. 699 query: The SQL query for the engine to base the managed table on 700 target_columns_to_types: A mapping between the column name and its data type. 701 partitioned_by: The partition columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 702 clustered_by: The cluster columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 703 table_properties: Optional mapping of engine-specific properties to be set on the managed table 704 table_description: Optional table description from MODEL DDL. 705 column_descriptions: Optional column descriptions from model query. 706 kwargs: Optional create table properties. 707 """ 708 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}") 709 710 def ctas( 711 self, 712 table_name: TableName, 713 query_or_df: QueryOrDF, 714 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 715 exists: bool = True, 716 table_description: t.Optional[str] = None, 717 column_descriptions: t.Optional[t.Dict[str, str]] = None, 718 source_columns: t.Optional[t.List[str]] = None, 719 **kwargs: t.Any, 720 ) -> None: 721 """Create a table using a CTAS statement 722 723 Args: 724 table_name: The name of the table to create. Can be fully qualified or just table name. 725 query_or_df: The SQL query to run or a dataframe for the CTAS. 726 target_columns_to_types: A mapping between the column name and its data type. Required if using a DataFrame. 727 exists: Indicates whether to include the IF NOT EXISTS check. 728 table_description: Optional table description from MODEL DDL. 729 column_descriptions: Optional column descriptions from model query. 730 kwargs: Optional create table properties. 731 """ 732 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 733 query_or_df, 734 target_columns_to_types, 735 target_table=table_name, 736 source_columns=source_columns, 737 ) 738 return self._create_table_from_source_queries( 739 table_name, 740 source_queries, 741 target_columns_to_types, 742 exists, 743 table_description=table_description, 744 column_descriptions=column_descriptions, 745 **kwargs, 746 ) 747 748 def create_state_table( 749 self, 750 table_name: str, 751 target_columns_to_types: t.Dict[str, exp.DataType], 752 primary_key: t.Optional[t.Tuple[str, ...]] = None, 753 ) -> None: 754 """Create a table to store SQLMesh internal state. 755 756 Args: 757 table_name: The name of the table to create. Can be fully qualified or just table name. 758 target_columns_to_types: A mapping between the column name and its data type. 759 primary_key: Determines the table primary key. 760 """ 761 self.create_table( 762 table_name, 763 target_columns_to_types, 764 primary_key=primary_key, 765 ) 766 767 def _create_table_from_columns( 768 self, 769 table_name: TableName, 770 target_columns_to_types: t.Dict[str, exp.DataType], 771 primary_key: t.Optional[t.Tuple[str, ...]] = None, 772 exists: bool = True, 773 table_description: t.Optional[str] = None, 774 column_descriptions: t.Optional[t.Dict[str, str]] = None, 775 **kwargs: t.Any, 776 ) -> None: 777 """ 778 Create a table using a DDL statement. 779 780 Args: 781 table_name: The name of the table to create. Can be fully qualified or just table name. 782 target_columns_to_types: Mapping between the column name and its data type. 783 primary_key: Determines the table primary key. 784 exists: Indicates whether to include the IF NOT EXISTS check. 785 table_description: Optional table description from MODEL DDL. 786 column_descriptions: Optional column descriptions from model query. 787 kwargs: Optional create table properties. 788 """ 789 table = exp.to_table(table_name) 790 791 if not columns_to_types_all_known(target_columns_to_types): 792 # It is ok if the columns types are not known if the table already exists and IF NOT EXISTS is set 793 if exists and self.table_exists(table_name): 794 return 795 raise SQLMeshError( 796 "Cannot create a table without knowing the column types. " 797 "Try casting the columns to an expected type or defining the columns in the model metadata. " 798 f"Columns to types: {target_columns_to_types}" 799 ) 800 801 primary_key_expression = ( 802 [exp.PrimaryKey(expressions=[exp.to_column(k) for k in primary_key])] 803 if primary_key and self.SUPPORTS_INDEXES 804 else [] 805 ) 806 807 schema = self._build_schema_exp( 808 table, 809 target_columns_to_types, 810 column_descriptions, 811 primary_key_expression, 812 ) 813 814 self._create_table( 815 schema, 816 None, 817 exists=exists, 818 target_columns_to_types=target_columns_to_types, 819 table_description=table_description, 820 **kwargs, 821 ) 822 823 # Register comments with commands if the engine doesn't support comments in the schema or CREATE 824 if ( 825 table_description 826 and self.COMMENT_CREATION_TABLE.is_comment_command_only 827 and self.comments_enabled 828 ): 829 self._create_table_comment(table_name, table_description) 830 if ( 831 column_descriptions 832 and self.COMMENT_CREATION_TABLE.is_comment_command_only 833 and self.comments_enabled 834 ): 835 self._create_column_comments(table_name, column_descriptions) 836 837 def _build_schema_exp( 838 self, 839 table: exp.Table, 840 target_columns_to_types: t.Dict[str, exp.DataType], 841 column_descriptions: t.Optional[t.Dict[str, str]] = None, 842 expressions: t.Optional[t.List[exp.PrimaryKey]] = None, 843 is_view: bool = False, 844 materialized: bool = False, 845 ) -> exp.Schema: 846 """ 847 Build a schema expression for a table, columns, column comments, and additional schema properties. 848 """ 849 expressions = expressions or [] 850 851 return exp.Schema( 852 this=table, 853 expressions=self._build_column_defs( 854 target_columns_to_types=target_columns_to_types, 855 column_descriptions=column_descriptions, 856 is_view=is_view, 857 materialized=materialized, 858 ) 859 + expressions, 860 ) 861 862 def _build_column_defs( 863 self, 864 target_columns_to_types: t.Dict[str, exp.DataType], 865 column_descriptions: t.Optional[t.Dict[str, str]] = None, 866 is_view: bool = False, 867 materialized: bool = False, 868 ) -> t.List[exp.ColumnDef]: 869 engine_supports_schema_comments = ( 870 self.COMMENT_CREATION_VIEW.supports_schema_def 871 if is_view 872 else self.COMMENT_CREATION_TABLE.supports_schema_def 873 ) 874 return [ 875 self._build_column_def( 876 column, 877 column_descriptions=column_descriptions, 878 engine_supports_schema_comments=engine_supports_schema_comments, 879 col_type=None if is_view else kind, # don't include column data type for views 880 ) 881 for column, kind in target_columns_to_types.items() 882 ] 883 884 def _build_column_def( 885 self, 886 col_name: str, 887 column_descriptions: t.Optional[t.Dict[str, str]] = None, 888 engine_supports_schema_comments: bool = False, 889 col_type: t.Optional[exp.DATA_TYPE] = None, 890 nested_names: t.List[str] = [], 891 ) -> exp.ColumnDef: 892 return exp.ColumnDef( 893 this=exp.to_identifier(col_name), 894 kind=col_type, 895 constraints=( 896 self._build_col_comment_exp(col_name, column_descriptions) 897 if engine_supports_schema_comments and self.comments_enabled and column_descriptions 898 else None 899 ), 900 ) 901 902 def _build_col_comment_exp( 903 self, col_name: str, column_descriptions: t.Dict[str, str] 904 ) -> t.List[exp.ColumnConstraint]: 905 comment = column_descriptions.get(col_name, None) 906 if comment: 907 return [ 908 exp.ColumnConstraint( 909 kind=exp.CommentColumnConstraint( 910 this=exp.Literal.string(self._truncate_column_comment(comment)) 911 ) 912 ) 913 ] 914 return [] 915 916 def _create_table_from_source_queries( 917 self, 918 table_name: TableName, 919 source_queries: t.List[SourceQuery], 920 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 921 exists: bool = True, 922 replace: bool = False, 923 table_description: t.Optional[str] = None, 924 column_descriptions: t.Optional[t.Dict[str, str]] = None, 925 table_kind: t.Optional[str] = None, 926 track_rows_processed: bool = True, 927 **kwargs: t.Any, 928 ) -> None: 929 table = exp.to_table(table_name) 930 931 # CTAS calls do not usually include a schema expression. However, most engines 932 # permit them in CTAS expressions, and they allow us to register all column comments 933 # in a single call rather than in a separate comment command call for each column. 934 # 935 # This block conditionally builds a schema expression with column comments if the engine 936 # supports it and we have columns_to_types. column_to_types is required because the 937 # schema expression must include at least column name, data type, and the comment - 938 # for example, `(colname INTEGER COMMENT 'comment')`. 939 # 940 # column_to_types will be available when loading from a DataFrame (by converting from 941 # pandas to SQL types), when a model is "annotated" by explicitly specifying column 942 # types, and for evaluation methods like `LogicalReplaceQueryMixin.replace_query()` 943 # calls and SCD Type 2 model calls. 944 schema = None 945 target_columns_to_types_known = target_columns_to_types and columns_to_types_all_known( 946 target_columns_to_types 947 ) 948 if ( 949 column_descriptions 950 and target_columns_to_types_known 951 and self.COMMENT_CREATION_TABLE.is_in_schema_def_ctas 952 and self.comments_enabled 953 ): 954 schema = self._build_schema_exp(table, target_columns_to_types, column_descriptions) # type: ignore 955 956 with self.transaction(condition=len(source_queries) > 1): 957 for i, source_query in enumerate(source_queries): 958 with source_query as query: 959 if target_columns_to_types and target_columns_to_types_known: 960 query = self._order_projections_and_filter( 961 query, target_columns_to_types, coerce_types=True 962 ) 963 if i == 0: 964 self._create_table( 965 schema if schema else table, 966 query, 967 target_columns_to_types=target_columns_to_types, 968 exists=exists, 969 replace=replace, 970 table_description=table_description, 971 table_kind=table_kind, 972 track_rows_processed=track_rows_processed, 973 **kwargs, 974 ) 975 else: 976 self._insert_append_query( 977 table_name, 978 query, 979 target_columns_to_types or self.columns(table), 980 track_rows_processed=track_rows_processed, 981 ) 982 983 # Register comments with commands if the engine supports comments and we weren't able to 984 # register them with the CTAS call's schema expression. 985 if ( 986 table_description 987 and self.COMMENT_CREATION_TABLE.is_comment_command_only 988 and self.comments_enabled 989 ): 990 self._create_table_comment(table_name, table_description) 991 if column_descriptions and schema is None and self.comments_enabled: 992 self._create_column_comments(table_name, column_descriptions) 993 994 def _create_table( 995 self, 996 table_name_or_schema: t.Union[exp.Schema, TableName], 997 expression: t.Optional[exp.Expr], 998 exists: bool = True, 999 replace: bool = False, 1000 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1001 table_description: t.Optional[str] = None, 1002 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1003 table_kind: t.Optional[str] = None, 1004 track_rows_processed: bool = True, 1005 **kwargs: t.Any, 1006 ) -> None: 1007 self.execute( 1008 self._build_create_table_exp( 1009 table_name_or_schema, 1010 expression=expression, 1011 exists=exists, 1012 replace=replace, 1013 target_columns_to_types=target_columns_to_types, 1014 table_description=( 1015 table_description 1016 if self.COMMENT_CREATION_TABLE.supports_schema_def and self.comments_enabled 1017 else None 1018 ), 1019 table_kind=table_kind, 1020 **kwargs, 1021 ), 1022 track_rows_processed=track_rows_processed, 1023 ) 1024 # Extract table name to clear cache 1025 table_name = ( 1026 table_name_or_schema.this 1027 if isinstance(table_name_or_schema, exp.Schema) 1028 else table_name_or_schema 1029 ) 1030 self._clear_data_object_cache(table_name) 1031 1032 def _build_create_table_exp( 1033 self, 1034 table_name_or_schema: t.Union[exp.Schema, TableName], 1035 expression: t.Optional[exp.Expr], 1036 exists: bool = True, 1037 replace: bool = False, 1038 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1039 table_description: t.Optional[str] = None, 1040 table_kind: t.Optional[str] = None, 1041 **kwargs: t.Any, 1042 ) -> exp.Create: 1043 exists = False if replace else exists 1044 catalog_name = None 1045 if not isinstance(table_name_or_schema, exp.Schema): 1046 table_name_or_schema = exp.to_table(table_name_or_schema) 1047 catalog_name = table_name_or_schema.catalog 1048 else: 1049 if isinstance(table_name_or_schema.this, exp.Table): 1050 catalog_name = table_name_or_schema.this.catalog 1051 1052 properties = ( 1053 self._build_table_properties_exp( 1054 **kwargs, 1055 catalog_name=catalog_name, 1056 target_columns_to_types=target_columns_to_types, 1057 table_description=table_description, 1058 table_kind=table_kind, 1059 ) 1060 if kwargs or table_description 1061 else None 1062 ) 1063 return exp.Create( 1064 this=table_name_or_schema, 1065 kind=table_kind or "TABLE", 1066 replace=replace, 1067 exists=exists, 1068 expression=expression, 1069 properties=properties, 1070 ) 1071 1072 def create_table_like( 1073 self, 1074 target_table_name: TableName, 1075 source_table_name: TableName, 1076 exists: bool = True, 1077 **kwargs: t.Any, 1078 ) -> None: 1079 """Create a table to store SQLMesh internal state based on the definition of another table, including any 1080 column attributes and indexes defined in the original table. 1081 1082 Args: 1083 target_table_name: The name of the table to create. Can be fully qualified or just table name. 1084 source_table_name: The name of the table to base the new table on. 1085 """ 1086 self._create_table_like(target_table_name, source_table_name, exists=exists, **kwargs) 1087 self._clear_data_object_cache(target_table_name) 1088 1089 def clone_table( 1090 self, 1091 target_table_name: TableName, 1092 source_table_name: TableName, 1093 replace: bool = False, 1094 exists: bool = True, 1095 clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, 1096 **kwargs: t.Any, 1097 ) -> None: 1098 """Creates a table with the target name by cloning the source table. 1099 1100 Args: 1101 target_table_name: The name of the table that should be created. 1102 source_table_name: The name of the source table that should be cloned. 1103 replace: Whether or not to replace an existing table. 1104 exists: Indicates whether to include the IF NOT EXISTS check. 1105 """ 1106 if not self.SUPPORTS_CLONING: 1107 raise NotImplementedError(f"Engine does not support cloning: {type(self)}") 1108 1109 kwargs.pop("rendered_physical_properties", None) 1110 self.execute( 1111 exp.Create( 1112 this=exp.to_table(target_table_name), 1113 kind="TABLE", 1114 replace=replace, 1115 exists=exists, 1116 clone=exp.Clone( 1117 this=exp.to_table(source_table_name), 1118 **(clone_kwargs or {}), 1119 ), 1120 **kwargs, 1121 ) 1122 ) 1123 self._clear_data_object_cache(target_table_name) 1124 1125 def drop_data_object(self, data_object: DataObject, ignore_if_not_exists: bool = True) -> None: 1126 """Drops a data object of arbitrary type. 1127 1128 Args: 1129 data_object: The data object to drop. 1130 ignore_if_not_exists: If True, no error will be raised if the data object does not exist. 1131 """ 1132 if data_object.type.is_view: 1133 self.drop_view(data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists) 1134 elif data_object.type.is_materialized_view: 1135 self.drop_view( 1136 data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists, materialized=True 1137 ) 1138 elif data_object.type.is_table: 1139 self.drop_table(data_object.to_table(), exists=ignore_if_not_exists) 1140 elif data_object.type.is_managed_table: 1141 self.drop_managed_table(data_object.to_table(), exists=ignore_if_not_exists) 1142 else: 1143 raise SQLMeshError( 1144 f"Can't drop data object '{data_object.to_table().sql(dialect=self.dialect)}' of type '{data_object.type.value}'" 1145 ) 1146 1147 def drop_table(self, table_name: TableName, exists: bool = True, **kwargs: t.Any) -> None: 1148 """Drops a table. 1149 1150 Args: 1151 table_name: The name of the table to drop. 1152 exists: If exists, defaults to True. 1153 """ 1154 self._drop_object(name=table_name, exists=exists, **kwargs) 1155 1156 def drop_managed_table(self, table_name: TableName, exists: bool = True) -> None: 1157 """Drops a managed table. 1158 1159 Args: 1160 table_name: The name of the table to drop. 1161 exists: If exists, defaults to True. 1162 """ 1163 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}") 1164 1165 def _drop_object( 1166 self, 1167 name: TableName | SchemaName, 1168 exists: bool = True, 1169 kind: str = "TABLE", 1170 cascade: bool = False, 1171 **drop_args: t.Any, 1172 ) -> None: 1173 """Drops an object. 1174 1175 An object could be a DATABASE, SCHEMA, VIEW, TABLE, DYNAMIC TABLE, TEMPORARY TABLE etc depending on the :kind. 1176 1177 Args: 1178 name: The name of the table to drop. 1179 exists: If exists, defaults to True. 1180 kind: What kind of object to drop. Defaults to TABLE 1181 cascade: Whether or not to DROP ... CASCADE. 1182 Note that this is ignored for :kind's that are not present in self.SUPPORTED_DROP_CASCADE_OBJECT_KINDS 1183 **drop_args: Any extra arguments to set on the Drop expression 1184 """ 1185 if cascade and kind.upper() in self.SUPPORTED_DROP_CASCADE_OBJECT_KINDS: 1186 drop_args["cascade"] = cascade 1187 1188 self.execute(exp.Drop(this=exp.to_table(name), kind=kind, exists=exists, **drop_args)) 1189 self._clear_data_object_cache(name) 1190 1191 def get_alter_operations( 1192 self, 1193 current_table_name: TableName, 1194 target_table_name: TableName, 1195 *, 1196 ignore_destructive: bool = False, 1197 ignore_additive: bool = False, 1198 ) -> t.List[TableAlterOperation]: 1199 """ 1200 Determines the alter statements needed to change the current table into the structure of the target table. 1201 """ 1202 return t.cast( 1203 t.List[TableAlterOperation], 1204 self.schema_differ.compare_columns( 1205 current_table_name, 1206 self.columns(current_table_name), 1207 self.columns(target_table_name), 1208 ignore_destructive=ignore_destructive, 1209 ignore_additive=ignore_additive, 1210 ), 1211 ) 1212 1213 def alter_table( 1214 self, 1215 alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], 1216 ) -> None: 1217 """ 1218 Performs the alter statements to change the current table into the structure of the target table. 1219 """ 1220 with self.transaction(): 1221 for alter_expression in [ 1222 x.expression if isinstance(x, TableAlterOperation) else x for x in alter_expressions 1223 ]: 1224 self.execute(alter_expression) 1225 1226 def create_view( 1227 self, 1228 view_name: TableName, 1229 query_or_df: QueryOrDF, 1230 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1231 replace: bool = True, 1232 materialized: bool = False, 1233 materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, 1234 table_description: t.Optional[str] = None, 1235 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1236 view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 1237 source_columns: t.Optional[t.List[str]] = None, 1238 **create_kwargs: t.Any, 1239 ) -> None: 1240 """Create a view with a query or dataframe. 1241 1242 If a dataframe is passed in, it will be converted into a literal values statement. 1243 This should only be done if the dataframe is very small! 1244 1245 Args: 1246 view_name: The view name. 1247 query_or_df: A query or dataframe. 1248 target_columns_to_types: Columns to use in the view statement. 1249 replace: Whether or not to replace an existing view defaults to True. 1250 materialized: Whether to create a a materialized view. Only used for engines that support this feature. 1251 materialized_properties: Optional materialized view properties to add to the view. 1252 table_description: Optional table description from MODEL DDL. 1253 column_descriptions: Optional column descriptions from model query. 1254 view_properties: Optional view properties to add to the view. 1255 create_kwargs: Additional kwargs to pass into the Create expression 1256 """ 1257 import pandas as pd 1258 1259 if materialized_properties and not materialized: 1260 raise SQLMeshError("Materialized properties are only supported for materialized views") 1261 1262 query_or_df = self._native_df_to_pandas_df(query_or_df) 1263 1264 if isinstance(query_or_df, pd.DataFrame): 1265 values: t.List[t.Tuple[t.Any, ...]] = list( 1266 query_or_df.itertuples(index=False, name=None) 1267 ) 1268 target_columns_to_types, source_columns = self._columns_to_types( 1269 query_or_df, target_columns_to_types, source_columns 1270 ) 1271 if not target_columns_to_types: 1272 raise SQLMeshError("columns_to_types must be provided for dataframes") 1273 source_columns_to_types = get_source_columns_to_types( 1274 target_columns_to_types, source_columns 1275 ) 1276 query_or_df = self._values_to_sql( 1277 values, 1278 source_columns_to_types, 1279 batch_start=0, 1280 batch_end=len(values), 1281 ) 1282 1283 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1284 query_or_df, 1285 target_columns_to_types, 1286 batch_size=0, 1287 target_table=view_name, 1288 source_columns=source_columns, 1289 ) 1290 if len(source_queries) != 1: 1291 raise SQLMeshError("Only one source query is supported for creating views") 1292 1293 schema: t.Union[exp.Table, exp.Schema] = exp.to_table(view_name) 1294 if target_columns_to_types: 1295 schema = self._build_schema_exp( 1296 exp.to_table(view_name), 1297 target_columns_to_types, 1298 column_descriptions, 1299 is_view=True, 1300 materialized=materialized, 1301 ) 1302 1303 properties = create_kwargs.pop("properties", None) 1304 if not properties: 1305 properties = exp.Properties(expressions=[]) 1306 1307 if view_properties: 1308 table_type = self._pop_creatable_type_from_properties(view_properties) 1309 if table_type: 1310 properties.append("expressions", table_type) 1311 1312 if materialized and self.SUPPORTS_MATERIALIZED_VIEWS: 1313 properties.append("expressions", exp.MaterializedProperty()) 1314 1315 if not self.SUPPORTS_MATERIALIZED_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1316 schema = schema.this 1317 1318 if not self.SUPPORTS_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1319 schema = schema.this 1320 1321 if materialized_properties: 1322 partitioned_by = materialized_properties.pop("partitioned_by", None) 1323 clustered_by = materialized_properties.pop("clustered_by", None) 1324 if ( 1325 partitioned_by 1326 and ( 1327 partitioned_by_prop := self._build_partitioned_by_exp( 1328 partitioned_by, **materialized_properties 1329 ) 1330 ) 1331 is not None 1332 ): 1333 materialized_properties["catalog_name"] = exp.to_table(view_name).catalog 1334 properties.append("expressions", partitioned_by_prop) 1335 if ( 1336 clustered_by 1337 and ( 1338 clustered_by_prop := self._build_clustered_by_exp( 1339 clustered_by, **materialized_properties 1340 ) 1341 ) 1342 is not None 1343 ): 1344 properties.append("expressions", clustered_by_prop) 1345 1346 create_view_properties = self._build_view_properties_exp( 1347 view_properties, 1348 ( 1349 table_description 1350 if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled 1351 else None 1352 ), 1353 physical_cluster=create_kwargs.pop("physical_cluster", None), 1354 ) 1355 if create_view_properties: 1356 for view_property in create_view_properties.expressions: 1357 # Small hack to make sure SECURE goes at the beginning before materialized as required by Snowflake 1358 if isinstance(view_property, exp.SecureProperty): 1359 properties.set("expressions", view_property, index=0, overwrite=False) 1360 else: 1361 properties.append("expressions", view_property) 1362 1363 if properties.expressions: 1364 create_kwargs["properties"] = properties 1365 1366 if replace: 1367 self.drop_data_object_on_type_mismatch( 1368 self.get_data_object(view_name), 1369 DataObjectType.VIEW if not materialized else DataObjectType.MATERIALIZED_VIEW, 1370 ) 1371 1372 with source_queries[0] as query: 1373 self.execute( 1374 exp.Create( 1375 this=schema, 1376 kind="VIEW", 1377 replace=replace, 1378 expression=query, 1379 **create_kwargs, 1380 ), 1381 quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, 1382 ) 1383 1384 self._clear_data_object_cache(view_name) 1385 1386 # Register table comment with commands if the engine doesn't support doing it in CREATE 1387 if ( 1388 table_description 1389 and self.COMMENT_CREATION_VIEW.is_comment_command_only 1390 and self.comments_enabled 1391 ): 1392 self._create_table_comment(view_name, table_description, "VIEW") 1393 # Register column comments with commands if the engine doesn't support doing it in 1394 # CREATE or we couldn't do it in the CREATE schema definition because we don't have 1395 # columns_to_types 1396 if ( 1397 column_descriptions 1398 and ( 1399 self.COMMENT_CREATION_VIEW.is_comment_command_only 1400 or ( 1401 self.COMMENT_CREATION_VIEW.is_in_schema_def_and_commands 1402 and not target_columns_to_types 1403 ) 1404 ) 1405 and self.comments_enabled 1406 ): 1407 self._create_column_comments(view_name, column_descriptions, "VIEW", materialized) 1408 1409 @set_catalog() 1410 def create_schema( 1411 self, 1412 schema_name: SchemaName, 1413 ignore_if_exists: bool = True, 1414 warn_on_error: bool = True, 1415 properties: t.Optional[t.List[exp.Expr]] = None, 1416 ) -> None: 1417 properties = properties or [] 1418 return self._create_schema( 1419 schema_name=schema_name, 1420 ignore_if_exists=ignore_if_exists, 1421 warn_on_error=warn_on_error, 1422 properties=properties, 1423 kind="SCHEMA", 1424 ) 1425 1426 def _create_schema( 1427 self, 1428 schema_name: SchemaName, 1429 ignore_if_exists: bool, 1430 warn_on_error: bool, 1431 properties: t.List[exp.Expr], 1432 kind: str, 1433 ) -> None: 1434 """Create a schema from a name or qualified table name.""" 1435 try: 1436 self.execute( 1437 exp.Create( 1438 this=to_schema(schema_name), 1439 kind=kind, 1440 exists=ignore_if_exists, 1441 properties=exp.Properties( # this renders as '' (empty string) if expressions is empty 1442 expressions=properties 1443 ), 1444 ) 1445 ) 1446 except Exception as e: 1447 if not warn_on_error: 1448 raise 1449 logger.warning("Failed to create %s '%s': %s", kind.lower(), schema_name, e) 1450 1451 def drop_schema( 1452 self, 1453 schema_name: SchemaName, 1454 ignore_if_not_exists: bool = True, 1455 cascade: bool = False, 1456 **drop_args: t.Dict[str, exp.Expr], 1457 ) -> None: 1458 return self._drop_object( 1459 name=schema_name, 1460 exists=ignore_if_not_exists, 1461 kind="SCHEMA", 1462 cascade=cascade, 1463 **drop_args, 1464 ) 1465 1466 def drop_view( 1467 self, 1468 view_name: TableName, 1469 ignore_if_not_exists: bool = True, 1470 materialized: bool = False, 1471 **kwargs: t.Any, 1472 ) -> None: 1473 """Drop a view.""" 1474 self._drop_object( 1475 name=view_name, 1476 exists=ignore_if_not_exists, 1477 kind="VIEW", 1478 materialized=materialized and self.SUPPORTS_MATERIALIZED_VIEWS, 1479 **kwargs, 1480 ) 1481 1482 def create_catalog(self, catalog_name: str | exp.Identifier) -> None: 1483 return self._create_catalog(exp.parse_identifier(catalog_name, dialect=self.dialect)) 1484 1485 def _create_catalog(self, catalog_name: exp.Identifier) -> None: 1486 raise SQLMeshError( 1487 f"Unable to create catalog '{catalog_name.sql(dialect=self.dialect)}' as automatic catalog management is not implemented in the {self.dialect} engine." 1488 ) 1489 1490 def drop_catalog(self, catalog_name: str | exp.Identifier) -> None: 1491 return self._drop_catalog(exp.parse_identifier(catalog_name, dialect=self.dialect)) 1492 1493 def _drop_catalog(self, catalog_name: exp.Identifier) -> None: 1494 raise SQLMeshError( 1495 f"Unable to drop catalog '{catalog_name.sql(dialect=self.dialect)}' as automatic catalog management is not implemented in the {self.dialect} engine." 1496 ) 1497 1498 def columns( 1499 self, table_name: TableName, include_pseudo_columns: bool = False 1500 ) -> t.Dict[str, exp.DataType]: 1501 """Fetches column names and types for the target table.""" 1502 self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) 1503 describe_output = self.cursor.fetchall() 1504 return { 1505 # Note: MySQL returns the column type as bytes. 1506 column_name: exp.DataType.build(_decoded_str(column_type), dialect=self.dialect) 1507 for column_name, column_type, *_ in itertools.takewhile( 1508 lambda t: not t[0].startswith("#"), 1509 describe_output, 1510 ) 1511 if column_name and column_name.strip() and column_type and column_type.strip() 1512 } 1513 1514 def table_exists(self, table_name: TableName) -> bool: 1515 table = exp.to_table(table_name) 1516 data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name) 1517 if data_object_cache_key in self._data_object_cache: 1518 logger.debug("Table existence cache hit: %s", data_object_cache_key) 1519 return self._data_object_cache[data_object_cache_key] is not None 1520 1521 try: 1522 self.execute(exp.Describe(this=table, kind="TABLE")) 1523 return True 1524 except Exception: 1525 return False 1526 1527 def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> None: 1528 self.execute(exp.delete(table_name, where)) 1529 1530 def insert_append( 1531 self, 1532 table_name: TableName, 1533 query_or_df: QueryOrDF, 1534 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1535 track_rows_processed: bool = True, 1536 source_columns: t.Optional[t.List[str]] = None, 1537 ) -> None: 1538 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1539 query_or_df, 1540 target_columns_to_types, 1541 target_table=table_name, 1542 source_columns=source_columns, 1543 ) 1544 self._insert_append_source_queries( 1545 table_name, source_queries, target_columns_to_types, track_rows_processed 1546 ) 1547 1548 def _insert_append_source_queries( 1549 self, 1550 table_name: TableName, 1551 source_queries: t.List[SourceQuery], 1552 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1553 track_rows_processed: bool = True, 1554 ) -> None: 1555 with self.transaction(condition=len(source_queries) > 0): 1556 target_columns_to_types = target_columns_to_types or self.columns(table_name) 1557 for source_query in source_queries: 1558 with source_query as query: 1559 self._insert_append_query( 1560 table_name, 1561 query, 1562 target_columns_to_types, 1563 track_rows_processed=track_rows_processed, 1564 ) 1565 1566 def _insert_append_query( 1567 self, 1568 table_name: TableName, 1569 query: Query, 1570 target_columns_to_types: t.Dict[str, exp.DataType], 1571 order_projections: bool = True, 1572 track_rows_processed: bool = True, 1573 ) -> None: 1574 if order_projections: 1575 query = self._order_projections_and_filter(query, target_columns_to_types) 1576 self.execute( 1577 exp.insert(query, table_name, columns=list(target_columns_to_types)), 1578 track_rows_processed=track_rows_processed, 1579 ) 1580 1581 def insert_overwrite_by_partition( 1582 self, 1583 table_name: TableName, 1584 query_or_df: QueryOrDF, 1585 partitioned_by: t.List[exp.Expr], 1586 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1587 source_columns: t.Optional[t.List[str]] = None, 1588 ) -> None: 1589 if self.INSERT_OVERWRITE_STRATEGY.is_insert_overwrite: 1590 target_table = exp.to_table(table_name) 1591 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1592 query_or_df, 1593 target_columns_to_types, 1594 target_table=target_table, 1595 source_columns=source_columns, 1596 ) 1597 self._insert_overwrite_by_condition( 1598 table_name, source_queries, target_columns_to_types=target_columns_to_types 1599 ) 1600 else: 1601 self._replace_by_key( 1602 table_name, 1603 query_or_df, 1604 target_columns_to_types, 1605 partitioned_by, 1606 is_unique_key=False, 1607 source_columns=source_columns, 1608 ) 1609 1610 def insert_overwrite_by_time_partition( 1611 self, 1612 table_name: TableName, 1613 query_or_df: QueryOrDF, 1614 start: TimeLike, 1615 end: TimeLike, 1616 time_formatter: t.Callable[[TimeLike, t.Optional[t.Dict[str, exp.DataType]]], exp.Expr], 1617 time_column: TimeColumn | exp.Expr | str, 1618 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1619 source_columns: t.Optional[t.List[str]] = None, 1620 **kwargs: t.Any, 1621 ) -> None: 1622 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1623 query_or_df, 1624 target_columns_to_types, 1625 target_table=table_name, 1626 source_columns=source_columns, 1627 ) 1628 if not target_columns_to_types or not columns_to_types_all_known(target_columns_to_types): 1629 target_columns_to_types = self.columns(table_name) 1630 low, high = [ 1631 time_formatter(dt, target_columns_to_types) 1632 for dt in make_inclusive(start, end, self.dialect) 1633 ] 1634 if isinstance(time_column, TimeColumn): 1635 time_column = time_column.column 1636 where = exp.Between( 1637 this=exp.to_column(time_column) if isinstance(time_column, str) else time_column, 1638 low=low, 1639 high=high, 1640 ) 1641 return self._insert_overwrite_by_time_partition( 1642 table_name, source_queries, target_columns_to_types, where, **kwargs 1643 ) 1644 1645 def _insert_overwrite_by_time_partition( 1646 self, 1647 table_name: TableName, 1648 source_queries: t.List[SourceQuery], 1649 target_columns_to_types: t.Dict[str, exp.DataType], 1650 where: exp.Condition, 1651 **kwargs: t.Any, 1652 ) -> None: 1653 return self._insert_overwrite_by_condition( 1654 table_name, source_queries, target_columns_to_types, where, **kwargs 1655 ) 1656 1657 def _values_to_sql( 1658 self, 1659 values: t.List[t.Tuple[t.Any, ...]], 1660 target_columns_to_types: t.Dict[str, exp.DataType], 1661 batch_start: int, 1662 batch_end: int, 1663 alias: str = "t", 1664 source_columns: t.Optional[t.List[str]] = None, 1665 ) -> Query: 1666 return select_from_values_for_batch_range( 1667 values=values, 1668 target_columns_to_types=target_columns_to_types, 1669 batch_start=batch_start, 1670 batch_end=batch_end, 1671 alias=alias, 1672 source_columns=source_columns, 1673 ) 1674 1675 def _insert_overwrite_by_condition( 1676 self, 1677 table_name: TableName, 1678 source_queries: t.List[SourceQuery], 1679 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1680 where: t.Optional[exp.Condition] = None, 1681 insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None, 1682 **kwargs: t.Any, 1683 ) -> None: 1684 table = exp.to_table(table_name) 1685 insert_overwrite_strategy = ( 1686 insert_overwrite_strategy_override or self.INSERT_OVERWRITE_STRATEGY 1687 ) 1688 with self.transaction( 1689 condition=len(source_queries) > 0 or insert_overwrite_strategy.is_delete_insert 1690 ): 1691 target_columns_to_types = target_columns_to_types or self.columns(table_name) 1692 for i, source_query in enumerate(source_queries): 1693 with source_query as query: 1694 query = self._order_projections_and_filter( 1695 query, target_columns_to_types, where=where 1696 ) 1697 if i > 0 or insert_overwrite_strategy.is_delete_insert: 1698 if i == 0: 1699 self.delete_from(table_name, where=where or exp.true()) 1700 self._insert_append_query( 1701 table_name, 1702 query, 1703 target_columns_to_types=target_columns_to_types, 1704 order_projections=False, 1705 ) 1706 elif insert_overwrite_strategy.is_merge: 1707 columns = [exp.column(col) for col in target_columns_to_types] 1708 when_not_matched_by_source = exp.When( 1709 matched=False, 1710 source=True, 1711 condition=where, 1712 then=exp.Delete(), 1713 ) 1714 when_not_matched_by_target = exp.When( 1715 matched=False, 1716 source=False, 1717 then=exp.Insert( 1718 this=exp.Tuple(expressions=columns), 1719 expression=exp.Tuple(expressions=columns), 1720 ), 1721 ) 1722 self._merge( 1723 target_table=table_name, 1724 query=query, 1725 on=exp.false(), 1726 whens=exp.Whens( 1727 expressions=[when_not_matched_by_source, when_not_matched_by_target] 1728 ), 1729 ) 1730 else: 1731 insert_exp = exp.insert( 1732 query, 1733 table, 1734 columns=( 1735 list(target_columns_to_types) 1736 if not insert_overwrite_strategy.is_replace_where 1737 else None 1738 ), 1739 overwrite=insert_overwrite_strategy.is_insert_overwrite, 1740 ) 1741 if insert_overwrite_strategy.is_replace_where: 1742 insert_exp.set("where", where or exp.true()) 1743 self.execute(insert_exp, track_rows_processed=True) 1744 1745 def update_table( 1746 self, 1747 table_name: TableName, 1748 properties: t.Dict[str, t.Any], 1749 where: t.Optional[str | exp.Condition] = None, 1750 ) -> None: 1751 self.execute(exp.update(table_name, properties, where=where)) 1752 1753 def _merge( 1754 self, 1755 target_table: TableName, 1756 query: Query, 1757 on: exp.Expr, 1758 whens: exp.Whens, 1759 ) -> None: 1760 this = exp.alias_(exp.to_table(target_table), alias=MERGE_TARGET_ALIAS, table=True) 1761 using = exp.alias_( 1762 exp.Subquery(this=query), alias=MERGE_SOURCE_ALIAS, copy=False, table=True 1763 ) 1764 self.execute( 1765 exp.Merge(this=this, using=using, on=on, whens=whens), track_rows_processed=True 1766 ) 1767 1768 def scd_type_2_by_time( 1769 self, 1770 target_table: TableName, 1771 source_table: QueryOrDF, 1772 unique_key: t.Sequence[exp.Expr], 1773 valid_from_col: exp.Column, 1774 valid_to_col: exp.Column, 1775 execution_time: t.Union[TimeLike, exp.Column], 1776 updated_at_col: exp.Column, 1777 invalidate_hard_deletes: bool = True, 1778 updated_at_as_valid_from: bool = False, 1779 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1780 table_description: t.Optional[str] = None, 1781 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1782 truncate: bool = False, 1783 source_columns: t.Optional[t.List[str]] = None, 1784 **kwargs: t.Any, 1785 ) -> None: 1786 self._scd_type_2( 1787 target_table=target_table, 1788 source_table=source_table, 1789 unique_key=unique_key, 1790 valid_from_col=valid_from_col, 1791 valid_to_col=valid_to_col, 1792 execution_time=execution_time, 1793 updated_at_col=updated_at_col, 1794 invalidate_hard_deletes=invalidate_hard_deletes, 1795 updated_at_as_valid_from=updated_at_as_valid_from, 1796 target_columns_to_types=target_columns_to_types, 1797 table_description=table_description, 1798 column_descriptions=column_descriptions, 1799 truncate=truncate, 1800 source_columns=source_columns, 1801 **kwargs, 1802 ) 1803 1804 def scd_type_2_by_column( 1805 self, 1806 target_table: TableName, 1807 source_table: QueryOrDF, 1808 unique_key: t.Sequence[exp.Expr], 1809 valid_from_col: exp.Column, 1810 valid_to_col: exp.Column, 1811 execution_time: t.Union[TimeLike, exp.Column], 1812 check_columns: t.Union[exp.Star, t.Sequence[exp.Expr]], 1813 invalidate_hard_deletes: bool = True, 1814 execution_time_as_valid_from: bool = False, 1815 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1816 table_description: t.Optional[str] = None, 1817 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1818 truncate: bool = False, 1819 source_columns: t.Optional[t.List[str]] = None, 1820 **kwargs: t.Any, 1821 ) -> None: 1822 self._scd_type_2( 1823 target_table=target_table, 1824 source_table=source_table, 1825 unique_key=unique_key, 1826 valid_from_col=valid_from_col, 1827 valid_to_col=valid_to_col, 1828 execution_time=execution_time, 1829 check_columns=check_columns, 1830 target_columns_to_types=target_columns_to_types, 1831 invalidate_hard_deletes=invalidate_hard_deletes, 1832 execution_time_as_valid_from=execution_time_as_valid_from, 1833 table_description=table_description, 1834 column_descriptions=column_descriptions, 1835 truncate=truncate, 1836 source_columns=source_columns, 1837 **kwargs, 1838 ) 1839 1840 def _scd_type_2( 1841 self, 1842 target_table: TableName, 1843 source_table: QueryOrDF, 1844 unique_key: t.Sequence[exp.Expr], 1845 valid_from_col: exp.Column, 1846 valid_to_col: exp.Column, 1847 execution_time: t.Union[TimeLike, exp.Column], 1848 invalidate_hard_deletes: bool = True, 1849 updated_at_col: t.Optional[exp.Column] = None, 1850 check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Expr]]] = None, 1851 updated_at_as_valid_from: bool = False, 1852 execution_time_as_valid_from: bool = False, 1853 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1854 table_description: t.Optional[str] = None, 1855 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1856 truncate: bool = False, 1857 source_columns: t.Optional[t.List[str]] = None, 1858 **kwargs: t.Any, 1859 ) -> None: 1860 def remove_managed_columns( 1861 cols_to_types: t.Dict[str, exp.DataType], 1862 ) -> t.Dict[str, exp.DataType]: 1863 return { 1864 k: v for k, v in cols_to_types.items() if k not in {valid_from_name, valid_to_name} 1865 } 1866 1867 valid_from_name = valid_from_col.name 1868 valid_to_name = valid_to_col.name 1869 target_columns_to_types = target_columns_to_types or self.columns(target_table) 1870 if ( 1871 valid_from_name not in target_columns_to_types 1872 or valid_to_name not in target_columns_to_types 1873 or not columns_to_types_all_known(target_columns_to_types) 1874 ): 1875 target_columns_to_types = self.columns(target_table) 1876 unmanaged_columns_to_types = ( 1877 remove_managed_columns(target_columns_to_types) if target_columns_to_types else None 1878 ) 1879 source_queries, unmanaged_columns_to_types = self._get_source_queries_and_columns_to_types( 1880 source_table, 1881 unmanaged_columns_to_types, 1882 target_table=target_table, 1883 batch_size=0, 1884 source_columns=source_columns, 1885 ) 1886 updated_at_name = updated_at_col.name if updated_at_col else None 1887 if not target_columns_to_types: 1888 raise SQLMeshError(f"Could not get columns_to_types. Does {target_table} exist?") 1889 unmanaged_columns_to_types = unmanaged_columns_to_types or remove_managed_columns( 1890 target_columns_to_types 1891 ) 1892 if not unique_key: 1893 raise SQLMeshError("unique_key must be provided for SCD Type 2") 1894 if check_columns and updated_at_col: 1895 raise SQLMeshError( 1896 "Cannot use both `check_columns` and `updated_at_name` for SCD Type 2" 1897 ) 1898 if check_columns and updated_at_as_valid_from: 1899 raise SQLMeshError( 1900 "Cannot use both `check_columns` and `updated_at_as_valid_from` for SCD Type 2" 1901 ) 1902 if execution_time_as_valid_from and not check_columns: 1903 raise SQLMeshError( 1904 "Cannot use `execution_time_as_valid_from` without `check_columns` for SCD Type 2" 1905 ) 1906 if updated_at_name and updated_at_name not in target_columns_to_types: 1907 raise SQLMeshError( 1908 f"Column {updated_at_name} not found in {target_table}. Table must contain an `updated_at` timestamp for SCD Type 2" 1909 ) 1910 time_data_type = target_columns_to_types[valid_from_name] 1911 select_source_columns: t.List[t.Union[str, exp.Alias]] = [ 1912 col for col in unmanaged_columns_to_types if col != updated_at_name 1913 ] 1914 table_columns = [exp.column(c, quoted=True) for c in target_columns_to_types] 1915 if updated_at_name: 1916 select_source_columns.append( 1917 exp.cast(updated_at_col, time_data_type).as_(updated_at_col.this) # type: ignore 1918 ) 1919 1920 # If a star is provided, we include all unmanaged columns in the check. 1921 # This unnecessarily includes unique key columns but since they are used in the join, and therefore we know 1922 # they are equal or not, the extra check is not a problem and we gain simplified logic here. 1923 # If we want to change this, then we just need to check the expressions in unique_key and pull out the 1924 # column names and then remove them from the unmanaged_columns 1925 if check_columns: 1926 # Handle both Star directly and [Star()] (which can happen during serialization/deserialization) 1927 if isinstance(seq_get(ensure_list(check_columns), 0), exp.Star): 1928 check_columns = [exp.column(col) for col in unmanaged_columns_to_types] 1929 execution_ts = ( 1930 exp.cast(execution_time, time_data_type, dialect=self.dialect) 1931 if isinstance(execution_time, exp.Column) 1932 else to_time_column(execution_time, time_data_type, self.dialect, nullable=True) 1933 ) 1934 if updated_at_as_valid_from: 1935 if not updated_at_col: 1936 raise SQLMeshError( 1937 "Cannot use `updated_at_as_valid_from` without `updated_at_name` for SCD Type 2" 1938 ) 1939 update_valid_from_start: t.Union[str, exp.Expr] = updated_at_col 1940 # If using check_columns and the user doesn't always want execution_time for valid from 1941 # then we only use epoch 0 if we are truncating the table and loading rows for the first time. 1942 # All future new rows should have execution time. 1943 elif check_columns and (execution_time_as_valid_from or not truncate): 1944 update_valid_from_start = execution_ts 1945 else: 1946 update_valid_from_start = to_time_column( 1947 "1970-01-01 00:00:00+00:00", time_data_type, self.dialect, nullable=True 1948 ) 1949 insert_valid_from_start = execution_ts if check_columns else updated_at_col # type: ignore 1950 # joined._exists IS NULL is saying "if the row is deleted" 1951 delete_check = ( 1952 exp.column("_exists", "joined").is_(exp.Null()) if invalidate_hard_deletes else None 1953 ) 1954 prefixed_valid_to_col = valid_to_col.copy() 1955 prefixed_valid_to_col.this.set("this", f"t_{prefixed_valid_to_col.name}") 1956 prefixed_valid_from_col = valid_from_col.copy() 1957 prefixed_valid_from_col.this.set("this", f"t_{valid_from_col.name}") 1958 if check_columns: 1959 row_check_conditions = [] 1960 for col in check_columns: 1961 col_qualified = col.copy() 1962 col_qualified.set("table", exp.to_identifier("joined")) 1963 1964 t_col = col_qualified.copy() 1965 for column in t_col.find_all(exp.Column): 1966 column.this.set("this", f"t_{column.name}") 1967 1968 row_check_conditions.extend( 1969 [ 1970 col_qualified.neq(t_col), 1971 exp.and_(t_col.is_(exp.Null()), col_qualified.is_(exp.Null()).not_()), 1972 exp.and_(t_col.is_(exp.Null()).not_(), col_qualified.is_(exp.Null())), 1973 ] 1974 ) 1975 row_value_check = exp.or_(*row_check_conditions) 1976 unique_key_conditions = [] 1977 for key in unique_key: 1978 key_qualified = key.copy() 1979 key_qualified.set("table", exp.to_identifier("joined")) 1980 t_key = key_qualified.copy() 1981 for col in t_key.find_all(exp.Column): 1982 col.this.set("this", f"t_{col.name}") 1983 unique_key_conditions.extend( 1984 [t_key.is_(exp.Null()).not_(), key_qualified.is_(exp.Null()).not_()] 1985 ) 1986 unique_key_check = exp.and_(*unique_key_conditions) 1987 # unique_key_check is saying "if the row is updated" 1988 # row_value_check is saying "if the row has changed" 1989 updated_row_filter = exp.and_(unique_key_check, row_value_check) 1990 valid_to_case_stmt = ( 1991 exp.Case() 1992 .when( 1993 exp.and_( 1994 exp.or_( 1995 delete_check, 1996 updated_row_filter, 1997 ) 1998 ), 1999 execution_ts, 2000 ) 2001 .else_(prefixed_valid_to_col) 2002 .as_(valid_to_col.this) 2003 ) 2004 valid_from_case_stmt = exp.func( 2005 "COALESCE", 2006 prefixed_valid_from_col, 2007 update_valid_from_start, 2008 ).as_(valid_from_col.this) 2009 else: 2010 assert updated_at_col is not None 2011 updated_at_col_qualified = updated_at_col.copy() 2012 updated_at_col_qualified.set("table", exp.to_identifier("joined")) 2013 prefixed_updated_at_col = updated_at_col_qualified.copy() 2014 prefixed_updated_at_col.this.set("this", f"t_{updated_at_col_qualified.name}") 2015 updated_row_filter = updated_at_col_qualified > prefixed_updated_at_col 2016 2017 valid_to_case_stmt_builder = exp.Case().when( 2018 updated_row_filter, updated_at_col_qualified 2019 ) 2020 if delete_check: 2021 valid_to_case_stmt_builder = valid_to_case_stmt_builder.when( 2022 delete_check, execution_ts 2023 ) 2024 valid_to_case_stmt = valid_to_case_stmt_builder.else_(prefixed_valid_to_col).as_( 2025 valid_to_col.this 2026 ) 2027 2028 valid_from_case_stmt = ( 2029 exp.Case() 2030 .when( 2031 exp.and_( 2032 prefixed_valid_from_col.is_(exp.Null()), 2033 exp.column("_exists", "latest_deleted").is_(exp.Null()).not_(), 2034 ), 2035 exp.Case() 2036 .when( 2037 exp.column(valid_to_col.this, "latest_deleted") > updated_at_col, 2038 exp.column(valid_to_col.this, "latest_deleted"), 2039 ) 2040 .else_(updated_at_col), 2041 ) 2042 .when(prefixed_valid_from_col.is_(exp.Null()), update_valid_from_start) 2043 .else_(prefixed_valid_from_col) 2044 ).as_(valid_from_col.this) 2045 2046 existing_rows_query = exp.select(*table_columns, exp.true().as_("_exists")).from_( 2047 target_table 2048 ) 2049 if truncate: 2050 existing_rows_query = existing_rows_query.limit(0) 2051 2052 with source_queries[0] as source_query: 2053 prefixed_columns_to_types = [] 2054 for column in target_columns_to_types: 2055 prefixed_col = exp.column(column).copy() 2056 prefixed_col.this.set("this", f"t_{prefixed_col.name}") 2057 prefixed_columns_to_types.append(prefixed_col) 2058 prefixed_unmanaged_columns = [] 2059 for column in unmanaged_columns_to_types: 2060 prefixed_col = exp.column(column).copy() 2061 prefixed_col.this.set("this", f"t_{prefixed_col.name}") 2062 prefixed_unmanaged_columns.append(prefixed_col) 2063 query = ( 2064 exp.Select() # type: ignore 2065 .select(*table_columns) 2066 .from_("static") 2067 .union( 2068 exp.select(*table_columns).from_("updated_rows"), 2069 distinct=False, 2070 ) 2071 .union( 2072 exp.select(*table_columns).from_("inserted_rows"), 2073 distinct=False, 2074 ) 2075 .with_( 2076 "source", 2077 exp.select(exp.true().as_("_exists"), *select_source_columns) 2078 .distinct(*unique_key) 2079 .from_( 2080 self.use_server_nulls_for_unmatched_after_join(source_query).subquery( # type: ignore 2081 "raw_source" 2082 ) 2083 ), 2084 ) 2085 # Historical Records that Do Not Change 2086 .with_( 2087 "static", 2088 existing_rows_query.where(valid_to_col.is_(exp.Null()).not_()), 2089 ) 2090 # Latest Records that can be updated 2091 .with_( 2092 "latest", 2093 existing_rows_query.where(valid_to_col.is_(exp.Null())), 2094 ) 2095 # Deleted records which can be used to determine `valid_from` for undeleted source records 2096 .with_( 2097 "deleted", 2098 exp.select(*[exp.column(col, "static") for col in target_columns_to_types]) 2099 .from_("static") 2100 .join( 2101 "latest", 2102 on=exp.and_( 2103 *[ 2104 add_table(key, "static").eq(add_table(key, "latest")) 2105 for key in unique_key 2106 ] 2107 ), 2108 join_type="left", 2109 ) 2110 .where(exp.column(valid_to_col.this, "latest").is_(exp.Null())), 2111 ) 2112 # Get the latest `valid_to` deleted record for each unique key 2113 .with_( 2114 "latest_deleted", 2115 exp.select( 2116 exp.true().as_("_exists"), 2117 *(part.as_(f"_key{i}") for i, part in enumerate(unique_key)), 2118 exp.Max(this=valid_to_col).as_(valid_to_col.this), 2119 ) 2120 .from_("deleted") 2121 .group_by(*unique_key), 2122 ) 2123 # Do a full join between latest records and source table in order to combine them together 2124 # MySQL doesn't support full join so going to do a left then right join and remove dups with union 2125 # We do a left/right and filter right on only matching to remove the need to do union distinct 2126 # which allows scd type 2 to be compatible with unhashable data types 2127 .with_( 2128 "joined", 2129 exp.select( 2130 exp.column("_exists", table="source").as_("_exists"), 2131 *( 2132 exp.column(col, table="latest").as_(prefixed_columns_to_types[i].this) 2133 for i, col in enumerate(target_columns_to_types) 2134 ), 2135 *( 2136 exp.column(col, table="source").as_(col) 2137 for col in unmanaged_columns_to_types 2138 ), 2139 ) 2140 .from_("latest") 2141 .join( 2142 "source", 2143 on=exp.and_( 2144 *[ 2145 add_table(key, "latest").eq(add_table(key, "source")) 2146 for key in unique_key 2147 ] 2148 ), 2149 join_type="left", 2150 ) 2151 .union( 2152 exp.select( 2153 exp.column("_exists", table="source").as_("_exists"), 2154 *( 2155 exp.column(col, table="latest").as_( 2156 prefixed_columns_to_types[i].this 2157 ) 2158 for i, col in enumerate(target_columns_to_types) 2159 ), 2160 *( 2161 exp.column(col, table="source").as_(col) 2162 for col in unmanaged_columns_to_types 2163 ), 2164 ) 2165 .from_("latest") 2166 .join( 2167 "source", 2168 on=exp.and_( 2169 *[ 2170 add_table(key, "latest").eq(add_table(key, "source")) 2171 for key in unique_key 2172 ] 2173 ), 2174 join_type="right", 2175 ) 2176 .where(exp.column("_exists", table="latest").is_(exp.Null())), 2177 distinct=False, 2178 ), 2179 ) 2180 # Get deleted, new, no longer current, or unchanged records 2181 .with_( 2182 "updated_rows", 2183 exp.select( 2184 *( 2185 exp.func( 2186 "COALESCE", 2187 exp.column(prefixed_unmanaged_columns[i].this, table="joined"), 2188 exp.column(col, table="joined"), 2189 ).as_(col) 2190 for i, col in enumerate(unmanaged_columns_to_types) 2191 ), 2192 valid_from_case_stmt, 2193 valid_to_case_stmt, 2194 ) 2195 .from_("joined") 2196 .join( 2197 "latest_deleted", 2198 on=exp.and_( 2199 *[ 2200 add_table(part, "joined").eq( 2201 exp.column(f"_key{i}", "latest_deleted") 2202 ) 2203 for i, part in enumerate(unique_key) 2204 ] 2205 ), 2206 join_type="left", 2207 ), 2208 ) 2209 # Get records that have been "updated" which means inserting a new record with previous `valid_from` 2210 .with_( 2211 "inserted_rows", 2212 exp.select( 2213 *unmanaged_columns_to_types, 2214 insert_valid_from_start.as_(valid_from_col.this), # type: ignore 2215 to_time_column(exp.null(), time_data_type, self.dialect, nullable=True).as_( 2216 valid_to_col.this 2217 ), 2218 ) 2219 .from_("joined") 2220 .where(updated_row_filter), 2221 ) 2222 ) 2223 2224 self.replace_query( 2225 target_table, 2226 self.ensure_nulls_for_unmatched_after_join(query), 2227 target_columns_to_types=target_columns_to_types, 2228 table_description=table_description, 2229 column_descriptions=column_descriptions, 2230 **kwargs, 2231 ) 2232 2233 def merge( 2234 self, 2235 target_table: TableName, 2236 source_table: QueryOrDF, 2237 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 2238 unique_key: t.Sequence[exp.Expr], 2239 when_matched: t.Optional[exp.Whens] = None, 2240 merge_filter: t.Optional[exp.Expr] = None, 2241 source_columns: t.Optional[t.List[str]] = None, 2242 **kwargs: t.Any, 2243 ) -> None: 2244 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2245 source_table, 2246 target_columns_to_types, 2247 target_table=target_table, 2248 source_columns=source_columns, 2249 ) 2250 target_columns_to_types = target_columns_to_types or self.columns(target_table) 2251 on = exp.and_( 2252 *( 2253 add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS)) 2254 for part in unique_key 2255 ) 2256 ) 2257 if merge_filter: 2258 on = exp.and_(merge_filter, on) 2259 2260 if not when_matched: 2261 match_expressions = [ 2262 exp.When( 2263 matched=True, 2264 source=False, 2265 then=exp.Update( 2266 expressions=[ 2267 exp.column(col, MERGE_TARGET_ALIAS).eq( 2268 exp.column(col, MERGE_SOURCE_ALIAS) 2269 ) 2270 for col in target_columns_to_types 2271 ], 2272 ), 2273 ) 2274 ] 2275 else: 2276 match_expressions = when_matched.copy().expressions 2277 2278 match_expressions.append( 2279 exp.When( 2280 matched=False, 2281 source=False, 2282 then=exp.Insert( 2283 this=exp.Tuple( 2284 expressions=[exp.column(col) for col in target_columns_to_types] 2285 ), 2286 expression=exp.Tuple( 2287 expressions=[ 2288 exp.column(col, MERGE_SOURCE_ALIAS) for col in target_columns_to_types 2289 ] 2290 ), 2291 ), 2292 ) 2293 ) 2294 for source_query in source_queries: 2295 with source_query as query: 2296 self._merge( 2297 target_table=target_table, 2298 query=query, 2299 on=on, 2300 whens=exp.Whens(expressions=match_expressions), 2301 ) 2302 2303 def rename_table( 2304 self, 2305 old_table_name: TableName, 2306 new_table_name: TableName, 2307 ) -> None: 2308 new_table = exp.to_table(new_table_name) 2309 if new_table.catalog: 2310 old_table = exp.to_table(old_table_name) 2311 catalog = old_table.catalog or self.get_current_catalog() 2312 if catalog != new_table.catalog: 2313 raise UnsupportedCatalogOperationError( 2314 "Tried to rename table across catalogs which is not supported" 2315 ) 2316 self._rename_table(old_table_name, new_table_name) 2317 self._clear_data_object_cache(old_table_name) 2318 self._clear_data_object_cache(new_table_name) 2319 2320 def get_data_object( 2321 self, target_name: TableName, safe_to_cache: bool = False 2322 ) -> t.Optional[DataObject]: 2323 target_table = exp.to_table(target_name) 2324 existing_data_objects = self.get_data_objects( 2325 schema_(target_table.db, target_table.catalog), 2326 {target_table.name}, 2327 safe_to_cache=safe_to_cache, 2328 ) 2329 if existing_data_objects: 2330 return existing_data_objects[0] 2331 return None 2332 2333 def get_data_objects( 2334 self, 2335 schema_name: SchemaName, 2336 object_names: t.Optional[t.Set[str]] = None, 2337 safe_to_cache: bool = False, 2338 ) -> t.List[DataObject]: 2339 """Lists all data objects in the target schema. 2340 2341 Args: 2342 schema_name: The name of the schema to list data objects from. 2343 object_names: If provided, only return data objects with these names. 2344 safe_to_cache: Whether it is safe to cache the results of this call. 2345 2346 Returns: 2347 A list of data objects in the target schema. 2348 """ 2349 if object_names is not None: 2350 if not object_names: 2351 return [] 2352 2353 # Check cache for each object name 2354 target_schema = to_schema(schema_name) 2355 cached_objects = [] 2356 missing_names = set() 2357 2358 for name in object_names: 2359 cache_key = _get_data_object_cache_key( 2360 target_schema.catalog, target_schema.db, name 2361 ) 2362 if cache_key in self._data_object_cache: 2363 logger.debug("Data object cache hit: %s", cache_key) 2364 data_object = self._data_object_cache[cache_key] 2365 # If the object is none, then the table was previously looked for but not found 2366 if data_object: 2367 cached_objects.append(data_object) 2368 else: 2369 logger.debug("Data object cache miss: %s", cache_key) 2370 missing_names.add(name) 2371 2372 # Fetch missing objects from database 2373 if missing_names: 2374 object_names_list = list(missing_names) 2375 batches = [ 2376 object_names_list[i : i + self.DATA_OBJECT_FILTER_BATCH_SIZE] 2377 for i in range(0, len(object_names_list), self.DATA_OBJECT_FILTER_BATCH_SIZE) 2378 ] 2379 2380 fetched_objects = [] 2381 fetched_object_names = set() 2382 for batch in batches: 2383 objects = self._get_data_objects(schema_name, set(batch)) 2384 for obj in objects: 2385 if safe_to_cache: 2386 cache_key = _get_data_object_cache_key( 2387 obj.catalog, obj.schema_name, obj.name 2388 ) 2389 self._data_object_cache[cache_key] = obj 2390 fetched_objects.append(obj) 2391 fetched_object_names.add(obj.name) 2392 2393 if safe_to_cache: 2394 for missing_name in missing_names - fetched_object_names: 2395 cache_key = _get_data_object_cache_key( 2396 target_schema.catalog, target_schema.db, missing_name 2397 ) 2398 self._data_object_cache[cache_key] = None 2399 2400 return cached_objects + fetched_objects 2401 2402 return cached_objects 2403 2404 fetched_objects = self._get_data_objects(schema_name) 2405 if safe_to_cache: 2406 for obj in fetched_objects: 2407 cache_key = _get_data_object_cache_key(obj.catalog, obj.schema_name, obj.name) 2408 self._data_object_cache[cache_key] = obj 2409 return fetched_objects 2410 2411 def fetchone( 2412 self, 2413 query: t.Union[exp.Expr, str], 2414 ignore_unsupported_errors: bool = False, 2415 quote_identifiers: bool = False, 2416 ) -> t.Optional[t.Tuple]: 2417 with self.transaction(): 2418 self.execute( 2419 query, 2420 ignore_unsupported_errors=ignore_unsupported_errors, 2421 quote_identifiers=quote_identifiers, 2422 ) 2423 return self.cursor.fetchone() 2424 2425 def fetchall( 2426 self, 2427 query: t.Union[exp.Expr, str], 2428 ignore_unsupported_errors: bool = False, 2429 quote_identifiers: bool = False, 2430 ) -> t.List[t.Tuple]: 2431 with self.transaction(): 2432 self.execute( 2433 query, 2434 ignore_unsupported_errors=ignore_unsupported_errors, 2435 quote_identifiers=quote_identifiers, 2436 ) 2437 return self.cursor.fetchall() 2438 2439 def _fetch_native_df( 2440 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2441 ) -> DF: 2442 """Fetches a DataFrame that can be either Pandas or PySpark from the cursor""" 2443 with self.transaction(): 2444 self.execute(query, quote_identifiers=quote_identifiers) 2445 return self.cursor.fetchdf() 2446 2447 def _native_df_to_pandas_df( 2448 self, 2449 query_or_df: QueryOrDF, 2450 ) -> t.Union[Query, pd.DataFrame]: 2451 """ 2452 Take a "native" DataFrame (eg Pyspark, Bigframe, Snowpark etc) and convert it to Pandas 2453 """ 2454 import pandas as pd 2455 2456 if isinstance(query_or_df, (exp.Query, pd.DataFrame)): 2457 return query_or_df 2458 2459 # EngineAdapter subclasses that have native DataFrame types should override this 2460 raise NotImplementedError(f"Unable to convert {type(query_or_df)} to Pandas") 2461 2462 def fetchdf( 2463 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2464 ) -> pd.DataFrame: 2465 """Fetches a Pandas DataFrame from the cursor""" 2466 import pandas as pd 2467 2468 df = self._fetch_native_df(query, quote_identifiers=quote_identifiers) 2469 if not isinstance(df, pd.DataFrame): 2470 raise NotImplementedError( 2471 "The cursor's `fetch_native_df` method is not returning a pandas DataFrame. Need to update `fetchdf` so a Pandas DataFrame is returned" 2472 ) 2473 return df 2474 2475 def fetch_pyspark_df( 2476 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2477 ) -> PySparkDataFrame: 2478 """Fetches a PySpark DataFrame from the cursor""" 2479 raise NotImplementedError(f"Engine does not support PySpark DataFrames: {type(self)}") 2480 2481 @property 2482 def wap_enabled(self) -> bool: 2483 """Returns whether WAP is enabled for this engine.""" 2484 return self._extra_config.get("wap_enabled", False) 2485 2486 def wap_supported(self, table_name: TableName) -> bool: 2487 """Returns whether WAP for the target table is supported.""" 2488 return False 2489 2490 def wap_table_name(self, table_name: TableName, wap_id: str) -> str: 2491 """Returns the updated table name for the given WAP ID. 2492 2493 Args: 2494 table_name: The name of the target table. 2495 wap_id: The WAP ID to prepare. 2496 2497 Returns: 2498 The updated table name that should be used for writing. 2499 """ 2500 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2501 2502 def wap_prepare(self, table_name: TableName, wap_id: str) -> str: 2503 """Prepares the target table for WAP and returns the updated table name. 2504 2505 Args: 2506 table_name: The name of the target table. 2507 wap_id: The WAP ID to prepare. 2508 2509 Returns: 2510 The updated table name that should be used for writing. 2511 """ 2512 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2513 2514 def wap_publish(self, table_name: TableName, wap_id: str) -> None: 2515 """Publishes changes with the given WAP ID to the target table. 2516 2517 Args: 2518 table_name: The name of the target table. 2519 wap_id: The WAP ID to publish. 2520 """ 2521 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2522 2523 def sync_grants_config( 2524 self, 2525 table: exp.Table, 2526 grants_config: GrantsConfig, 2527 table_type: DataObjectType = DataObjectType.TABLE, 2528 ) -> None: 2529 """Applies the grants_config to a table authoritatively. 2530 It first compares the specified grants against the current grants, and then 2531 applies the diffs to the table by revoking and granting privileges as needed. 2532 2533 Args: 2534 table: The table/view to apply grants to. 2535 grants_config: Dictionary mapping privileges to lists of grantees. 2536 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 2537 """ 2538 if not self.SUPPORTS_GRANTS: 2539 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 2540 2541 current_grants = self._get_current_grants_config(table) 2542 new_grants, revoked_grants = self._diff_grants_configs(grants_config, current_grants) 2543 revoke_exprs = self._revoke_grants_config_expr(table, revoked_grants, table_type) 2544 grant_exprs = self._apply_grants_config_expr(table, new_grants, table_type) 2545 dcl_exprs = revoke_exprs + grant_exprs 2546 2547 if dcl_exprs: 2548 self.execute(dcl_exprs) 2549 2550 @contextlib.contextmanager 2551 def transaction( 2552 self, 2553 condition: t.Optional[bool] = None, 2554 ) -> t.Iterator[None]: 2555 """A transaction context manager.""" 2556 if ( 2557 self._connection_pool.is_transaction_active 2558 or not self.SUPPORTS_TRANSACTIONS 2559 or (condition is not None and not condition) 2560 ): 2561 yield 2562 return 2563 2564 if self._pre_ping: 2565 try: 2566 logger.debug("Pinging the database to check the connection") 2567 self.ping() 2568 except Exception: 2569 logger.info("Connection to the database was lost. Reconnecting...") 2570 self._connection_pool.close() 2571 2572 self._connection_pool.begin() 2573 try: 2574 yield 2575 except Exception as e: 2576 self._connection_pool.rollback() 2577 raise e 2578 else: 2579 self._connection_pool.commit() 2580 2581 @contextlib.contextmanager 2582 def session(self, properties: SessionProperties) -> t.Iterator[None]: 2583 """A session context manager.""" 2584 if self._is_session_active(): 2585 yield 2586 return 2587 2588 self._begin_session(properties) 2589 try: 2590 yield 2591 finally: 2592 self._end_session() 2593 2594 def _begin_session(self, properties: SessionProperties) -> t.Any: 2595 """Begin a new session.""" 2596 2597 def _end_session(self) -> None: 2598 """End the existing session.""" 2599 2600 def _is_session_active(self) -> bool: 2601 """Indicates whether or not a session is active.""" 2602 return False 2603 2604 def execute( 2605 self, 2606 expressions: t.Union[str, exp.Expr, t.Sequence[exp.Expr]], 2607 ignore_unsupported_errors: bool = False, 2608 quote_identifiers: bool = True, 2609 track_rows_processed: bool = False, 2610 **kwargs: t.Any, 2611 ) -> None: 2612 """Execute a sql query.""" 2613 to_sql_kwargs = ( 2614 {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {} 2615 ) 2616 with self.transaction(): 2617 for e in ensure_list(expressions): 2618 if isinstance(e, exp.Expr): 2619 self._check_identifier_length(e) 2620 sql = self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs) 2621 else: 2622 sql = t.cast(str, e) 2623 2624 sql = self._attach_correlation_id(sql) 2625 2626 self._log_sql( 2627 sql, 2628 expression=e if isinstance(e, exp.Expr) else None, 2629 quote_identifiers=quote_identifiers, 2630 ) 2631 self._execute(sql, track_rows_processed, **kwargs) 2632 2633 def _attach_correlation_id(self, sql: str) -> str: 2634 if self.ATTACH_CORRELATION_ID and self.correlation_id: 2635 return f"/* {self.correlation_id} */ {sql}" 2636 return sql 2637 2638 def _log_sql( 2639 self, 2640 sql: str, 2641 expression: t.Optional[exp.Expr] = None, 2642 quote_identifiers: bool = True, 2643 ) -> None: 2644 if not logger.isEnabledFor(self._execute_log_level): 2645 return 2646 2647 sql_to_log = sql 2648 if expression is not None and not isinstance(expression, exp.Query): 2649 values = expression.find(exp.Values) 2650 if values: 2651 values.set("expressions", [exp.to_identifier("<REDACTED VALUES>")]) 2652 sql_to_log = self._to_sql(expression, quote=quote_identifiers) 2653 2654 logger.log(self._execute_log_level, "Executing SQL: %s", sql_to_log) 2655 2656 def _record_execution_stats( 2657 self, sql: str, rowcount: t.Optional[int] = None, bytes_processed: t.Optional[int] = None 2658 ) -> None: 2659 if self._query_execution_tracker: 2660 self._query_execution_tracker.record_execution(sql, rowcount, bytes_processed) 2661 2662 def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: 2663 self.cursor.execute(sql, **kwargs) 2664 2665 if ( 2666 self.SUPPORTS_QUERY_EXECUTION_TRACKING 2667 and track_rows_processed 2668 and self._query_execution_tracker 2669 and self._query_execution_tracker.is_tracking() 2670 ): 2671 if ( 2672 rowcount := getattr(self.cursor, "rowcount", None) 2673 ) is not None and rowcount is not None: 2674 try: 2675 self._record_execution_stats(sql, int(rowcount)) 2676 except (TypeError, ValueError): 2677 return 2678 2679 @contextlib.contextmanager 2680 def temp_table( 2681 self, 2682 query_or_df: QueryOrDF, 2683 name: TableName = "diff", 2684 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2685 source_columns: t.Optional[t.List[str]] = None, 2686 **kwargs: t.Any, 2687 ) -> t.Iterator[exp.Table]: 2688 """A context manager for working a temp table. 2689 2690 The table will be created with a random guid and cleaned up after the block. 2691 2692 Args: 2693 query_or_df: The query or df to create a temp table for. 2694 name: The base name of the temp table. 2695 target_columns_to_types: A mapping between the column name and its data type. 2696 2697 Yields: 2698 The table expression 2699 """ 2700 name = exp.to_table(name) 2701 # ensure that we use default catalog if none is not specified 2702 if isinstance(name, exp.Table) and not name.catalog and name.db and self.default_catalog: 2703 name.set("catalog", exp.parse_identifier(self.default_catalog)) 2704 2705 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2706 query_or_df, 2707 target_columns_to_types=target_columns_to_types, 2708 target_table=name, 2709 source_columns=source_columns, 2710 ) 2711 2712 with self.transaction(): 2713 table = self._get_temp_table(name) 2714 if table.db: 2715 self.create_schema(schema_(table.args["db"], table.args.get("catalog"))) 2716 self._create_table_from_source_queries( 2717 table, 2718 source_queries, 2719 target_columns_to_types, 2720 exists=True, 2721 table_description=None, 2722 column_descriptions=None, 2723 track_rows_processed=False, 2724 **kwargs, 2725 ) 2726 2727 try: 2728 yield table 2729 finally: 2730 self.drop_table(table) 2731 2732 def _table_or_view_properties_to_expressions( 2733 self, table_or_view_properties: t.Optional[t.Dict[str, exp.Expr]] = None 2734 ) -> t.List[exp.Property]: 2735 """Converts model properties (either physical or virtual) to a list of property expressions.""" 2736 if not table_or_view_properties: 2737 return [] 2738 return [ 2739 exp.Property(this=key, value=value.copy()) 2740 for key, value in table_or_view_properties.items() 2741 ] 2742 2743 def _build_partitioned_by_exp( 2744 self, 2745 partitioned_by: t.List[exp.Expr], 2746 *, 2747 partition_interval_unit: t.Optional[IntervalUnit] = None, 2748 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2749 catalog_name: t.Optional[str] = None, 2750 **kwargs: t.Any, 2751 ) -> t.Optional[t.Union[exp.PartitionedByProperty, exp.Property]]: 2752 return None 2753 2754 def _build_clustered_by_exp( 2755 self, 2756 clustered_by: t.List[exp.Expr], 2757 **kwargs: t.Any, 2758 ) -> t.Optional[exp.Cluster]: 2759 return None 2760 2761 def adjust_physical_properties_for_incremental( 2762 self, 2763 physical_properties: t.Dict[str, t.Any], 2764 *, 2765 requires_delete_capable_table: bool, 2766 unique_key: t.Optional[t.List[exp.Expr]], 2767 model_name: str, 2768 ) -> t.Dict[str, t.Any]: 2769 """Adjusts physical properties for an incremental model before the table is created. 2770 2771 Some engines require a specific physical table layout before they can run the DELETE/MERGE 2772 statements that incremental model kinds rely on (e.g. StarRocks only supports those on 2773 PRIMARY KEY tables). This hook lets each engine derive or validate the required properties 2774 while keeping the generic evaluator free of engine-specific branching. 2775 2776 Args: 2777 physical_properties: The model's physical properties. 2778 requires_delete_capable_table: Whether the model kind issues DELETE/MERGE statements 2779 (as opposed to append-only INSERTs), as determined by the generic evaluator. 2780 unique_key: The model's unique key, populated only when the kind allows promoting it to 2781 an engine-specific key (i.e. INCREMENTAL_BY_UNIQUE_KEY); otherwise None. 2782 model_name: The model name, for use in diagnostics. 2783 2784 Returns: 2785 The (possibly adjusted) physical properties. Implementations own the given mapping and 2786 may mutate it in place; the base implementation returns it unchanged. 2787 """ 2788 return physical_properties 2789 2790 def _build_table_properties_exp( 2791 self, 2792 catalog_name: t.Optional[str] = None, 2793 table_format: t.Optional[str] = None, 2794 storage_format: t.Optional[str] = None, 2795 partitioned_by: t.Optional[t.List[exp.Expr]] = None, 2796 partition_interval_unit: t.Optional[IntervalUnit] = None, 2797 clustered_by: t.Optional[t.List[exp.Expr]] = None, 2798 table_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 2799 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2800 table_description: t.Optional[str] = None, 2801 table_kind: t.Optional[str] = None, 2802 **kwargs: t.Any, 2803 ) -> t.Optional[exp.Properties]: 2804 """Creates a SQLGlot table properties expression for ddl.""" 2805 properties: t.List[exp.Expr] = [] 2806 2807 if table_description: 2808 properties.append( 2809 exp.SchemaCommentProperty( 2810 this=exp.Literal.string(self._truncate_table_comment(table_description)) 2811 ) 2812 ) 2813 2814 if table_properties: 2815 table_type = self._pop_creatable_type_from_properties(table_properties) 2816 properties.extend(ensure_list(table_type)) 2817 2818 if properties: 2819 return exp.Properties(expressions=properties) 2820 return None 2821 2822 def _build_view_properties_exp( 2823 self, 2824 view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 2825 table_description: t.Optional[str] = None, 2826 **kwargs: t.Any, 2827 ) -> t.Optional[exp.Properties]: 2828 """Creates a SQLGlot table properties expression for view""" 2829 properties: t.List[exp.Expr] = [] 2830 2831 if table_description: 2832 properties.append( 2833 exp.SchemaCommentProperty( 2834 this=exp.Literal.string(self._truncate_table_comment(table_description)) 2835 ) 2836 ) 2837 2838 if properties: 2839 return exp.Properties(expressions=properties) 2840 return None 2841 2842 def _truncate_comment(self, comment: str, length: t.Optional[int]) -> str: 2843 return comment[:length] if length else comment 2844 2845 def _truncate_table_comment(self, comment: str) -> str: 2846 return self._truncate_comment(comment, self.MAX_TABLE_COMMENT_LENGTH) 2847 2848 def _truncate_column_comment(self, comment: str) -> str: 2849 return self._truncate_comment(comment, self.MAX_COLUMN_COMMENT_LENGTH) 2850 2851 def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: 2852 """ 2853 Converts an expression to a SQL string. Has a set of default kwargs to apply, and then default 2854 kwargs defined for the given dialect, and then kwargs provided by the user when defining the engine 2855 adapter, and then finally kwargs provided by the user when calling this method. 2856 """ 2857 sql_gen_kwargs = { 2858 "dialect": self.dialect, 2859 "pretty": self._pretty_sql, 2860 "comments": False, 2861 **self._sql_gen_kwargs, 2862 **kwargs, 2863 } 2864 2865 expression = expression.copy() 2866 2867 if quote: 2868 quote_identifiers(expression) 2869 2870 return expression.sql(**sql_gen_kwargs, copy=False) # type: ignore 2871 2872 def _clear_data_object_cache(self, table_name: t.Optional[TableName] = None) -> None: 2873 """Clears the cache entry for the given table name, or clears the entire cache if table_name is None.""" 2874 if table_name is None: 2875 logger.debug("Clearing entire data object cache") 2876 self._data_object_cache.clear() 2877 else: 2878 table = exp.to_table(table_name) 2879 cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name) 2880 logger.debug("Clearing data object cache key: %s", cache_key) 2881 self._data_object_cache.pop(cache_key, None) 2882 2883 def _get_data_objects( 2884 self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None 2885 ) -> t.List[DataObject]: 2886 """ 2887 Returns all the data objects that exist in the given schema and optionally catalog. 2888 """ 2889 raise NotImplementedError() 2890 2891 def _get_temp_table( 2892 self, table: TableName, table_only: bool = False, quoted: bool = True 2893 ) -> exp.Table: 2894 """ 2895 Returns the name of the temp table that should be used for the given table name. 2896 """ 2897 table = t.cast(exp.Table, exp.to_table(table).copy()) 2898 table.set( 2899 "this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}", quoted=quoted) 2900 ) 2901 2902 if table_only: 2903 table.set("db", None) 2904 table.set("catalog", None) 2905 2906 return table 2907 2908 def _order_projections_and_filter( 2909 self, 2910 query: Query, 2911 target_columns_to_types: t.Dict[str, exp.DataType], 2912 where: t.Optional[exp.Expr] = None, 2913 coerce_types: bool = False, 2914 ) -> Query: 2915 if not isinstance(query, exp.Query) or ( 2916 not where and not coerce_types and query.named_selects == list(target_columns_to_types) 2917 ): 2918 return query 2919 2920 query = t.cast(exp.Query, query.copy()) 2921 with_ = query.args.pop("with_", None) 2922 2923 select_exprs: t.List[exp.Expr] = [ 2924 exp.column(c, quoted=True) for c in target_columns_to_types 2925 ] 2926 if coerce_types and columns_to_types_all_known(target_columns_to_types): 2927 select_exprs = [ 2928 exp.cast(select_exprs[i], col_tpe).as_(col, quoted=True) 2929 for i, (col, col_tpe) in enumerate(target_columns_to_types.items()) 2930 ] 2931 2932 query = exp.select(*select_exprs).from_(query.subquery("_subquery", copy=False), copy=False) 2933 if where: 2934 query = query.where(where, copy=False) 2935 2936 if with_: 2937 query.set("with_", with_) 2938 2939 return query 2940 2941 def _truncate_table(self, table_name: TableName) -> None: 2942 table = exp.to_table(table_name) 2943 self.execute(f"TRUNCATE TABLE {table.sql(dialect=self.dialect, identify=True)}") 2944 2945 def drop_data_object_on_type_mismatch( 2946 self, data_object: t.Optional[DataObject], expected_type: DataObjectType 2947 ) -> bool: 2948 """Drops a data object if it exists and is not of the expected type. 2949 2950 Args: 2951 data_object: The data object to check. 2952 expected_type: The expected type of the data object. 2953 2954 Returns: 2955 True if the data object was dropped, False otherwise. 2956 """ 2957 if data_object is None or data_object.type == expected_type: 2958 return False 2959 2960 logger.warning( 2961 "Target data object '%s' is a %s and not a %s, dropping it", 2962 data_object.to_table().sql(dialect=self.dialect), 2963 data_object.type.value, 2964 expected_type.value, 2965 ) 2966 self.drop_data_object(data_object) 2967 return True 2968 2969 def _replace_by_key( 2970 self, 2971 target_table: TableName, 2972 source_table: QueryOrDF, 2973 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 2974 key: t.Sequence[exp.Expr], 2975 is_unique_key: bool, 2976 source_columns: t.Optional[t.List[str]] = None, 2977 ) -> None: 2978 if target_columns_to_types is None: 2979 target_columns_to_types = self.columns(target_table) 2980 2981 temp_table = self._get_temp_table(target_table) 2982 key_exp = ( 2983 exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key, dialect=self.dialect) 2984 if len(key) > 1 2985 else key[0] 2986 ) 2987 column_names = list(target_columns_to_types or []) 2988 2989 with self.transaction(): 2990 self.ctas( 2991 temp_table, 2992 source_table, 2993 target_columns_to_types=target_columns_to_types, 2994 exists=False, 2995 source_columns=source_columns, 2996 ) 2997 2998 try: 2999 delete_query = exp.select(key_exp).from_(temp_table) 3000 insert_query = self._select_columns(target_columns_to_types).from_(temp_table) 3001 if not is_unique_key: 3002 delete_query = delete_query.distinct() 3003 else: 3004 insert_query = insert_query.distinct(*key) 3005 3006 insert_statement = exp.insert( 3007 insert_query, 3008 target_table, 3009 columns=column_names, 3010 ) 3011 delete_filter = key_exp.isin(query=delete_query) 3012 3013 if not self.INSERT_OVERWRITE_STRATEGY.is_replace_where: 3014 self.delete_from(target_table, delete_filter) 3015 else: 3016 insert_statement.set("where", delete_filter) 3017 insert_statement.set("this", exp.to_table(target_table)) 3018 3019 self.execute(insert_statement, track_rows_processed=True) 3020 finally: 3021 self.drop_table(temp_table) 3022 3023 def _build_create_comment_table_exp( 3024 self, table: exp.Table, table_comment: str, table_kind: str 3025 ) -> exp.Comment | str: 3026 return exp.Comment( 3027 this=table, 3028 kind=table_kind, 3029 expression=exp.Literal.string(self._truncate_table_comment(table_comment)), 3030 ) 3031 3032 def _create_table_comment( 3033 self, table_name: TableName, table_comment: str, table_kind: str = "TABLE" 3034 ) -> None: 3035 table = exp.to_table(table_name) 3036 3037 try: 3038 self.execute(self._build_create_comment_table_exp(table, table_comment, table_kind)) 3039 except Exception: 3040 logger.warning( 3041 f"Table comment for '{table.alias_or_name}' not registered - this may be due to limited permissions", 3042 exc_info=True, 3043 ) 3044 3045 def _build_create_comment_column_exp( 3046 self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE" 3047 ) -> exp.Comment | str: 3048 return exp.Comment( 3049 this=exp.column(column_name, *reversed(table.parts)), # type: ignore 3050 kind="COLUMN", 3051 expression=exp.Literal.string(self._truncate_column_comment(column_comment)), 3052 ) 3053 3054 def _create_column_comments( 3055 self, 3056 table_name: TableName, 3057 column_comments: t.Dict[str, str], 3058 table_kind: str = "TABLE", 3059 materialized_view: bool = False, 3060 ) -> None: 3061 table = exp.to_table(table_name) 3062 3063 for col, comment in column_comments.items(): 3064 try: 3065 self.execute(self._build_create_comment_column_exp(table, col, comment, table_kind)) 3066 except Exception: 3067 logger.warning( 3068 f"Column comments for column '{col}' in table '{table.alias_or_name}' not registered - this may be due to limited permissions", 3069 exc_info=True, 3070 ) 3071 3072 def _create_table_like( 3073 self, 3074 target_table_name: TableName, 3075 source_table_name: TableName, 3076 exists: bool, 3077 **kwargs: t.Any, 3078 ) -> None: 3079 self.create_table(target_table_name, self.columns(source_table_name), exists=exists) 3080 3081 def _rename_table( 3082 self, 3083 old_table_name: TableName, 3084 new_table_name: TableName, 3085 ) -> None: 3086 self.execute(exp.rename_table(old_table_name, new_table_name)) 3087 3088 def ensure_nulls_for_unmatched_after_join( 3089 self, 3090 query: Query, 3091 ) -> Query: 3092 return query 3093 3094 def use_server_nulls_for_unmatched_after_join( 3095 self, 3096 query: Query, 3097 ) -> Query: 3098 return query 3099 3100 def ping(self) -> None: 3101 try: 3102 self._execute(exp.select("1").sql(dialect=self.dialect)) 3103 finally: 3104 self._connection_pool.close_cursor() 3105 3106 @classmethod 3107 def _select_columns( 3108 cls, columns: t.Iterable[str], source_columns: t.Optional[t.List[str]] = None 3109 ) -> exp.Select: 3110 return exp.select( 3111 *( 3112 exp.column(c, quoted=True) 3113 if c in (source_columns or columns) 3114 else exp.alias_(exp.Null(), c, quoted=True) 3115 for c in columns 3116 ) 3117 ) 3118 3119 def _check_identifier_length(self, expression: exp.Expr) -> None: 3120 if self.MAX_IDENTIFIER_LENGTH is None or not isinstance(expression, exp.DDL): 3121 return 3122 3123 for identifier in expression.find_all(exp.Identifier): 3124 name = identifier.name 3125 name_length = len(name) 3126 if name_length > self.MAX_IDENTIFIER_LENGTH: 3127 raise SQLMeshError( 3128 f"Identifier name '{name}' (length {name_length}) exceeds {self.dialect.capitalize()}'s max identifier limit of {self.MAX_IDENTIFIER_LENGTH} characters" 3129 ) 3130 3131 def get_table_last_modified_ts(self, table_names: t.List[TableName]) -> t.List[int]: 3132 raise NotImplementedError() 3133 3134 @classmethod 3135 def _diff_grants_configs( 3136 cls, new_config: GrantsConfig, old_config: GrantsConfig 3137 ) -> t.Tuple[GrantsConfig, GrantsConfig]: 3138 """Compute additions and removals between two grants configurations. 3139 3140 This method compares new (desired) and old (current) GrantsConfigs case-insensitively 3141 for both privilege keys and grantees, while preserving original casing 3142 in the output GrantsConfigs. 3143 3144 Args: 3145 new_config: Desired grants configuration (specified by the user). 3146 old_config: Current grants configuration (returned by the database). 3147 3148 Returns: 3149 A tuple of (additions, removals) GrantsConfig where: 3150 - additions contains privileges/grantees present in new_config but not in old_config 3151 - additions uses keys and grantee strings from new_config (user-specified casing) 3152 - removals contains privileges/grantees present in old_config but not in new_config 3153 - removals uses keys and grantee strings from old_config (database-returned casing) 3154 3155 Notes: 3156 - Comparison is case-insensitive using casefold(); original casing is preserved in results. 3157 - Overlapping grantees (case-insensitive) are excluded from the results. 3158 """ 3159 3160 def _diffs(config1: GrantsConfig, config2: GrantsConfig) -> GrantsConfig: 3161 diffs: GrantsConfig = {} 3162 cf_config2 = {k.casefold(): {g.casefold() for g in v} for k, v in config2.items()} 3163 for key, grantees in config1.items(): 3164 cf_key = key.casefold() 3165 3166 # Missing key (add all grantees) 3167 if cf_key not in cf_config2: 3168 diffs[key] = grantees.copy() 3169 continue 3170 3171 # Include only grantees not in config2 3172 cf_grantees2 = cf_config2[cf_key] 3173 diff_grantees = [] 3174 for grantee in grantees: 3175 if grantee.casefold() not in cf_grantees2: 3176 diff_grantees.append(grantee) 3177 if diff_grantees: 3178 diffs[key] = diff_grantees 3179 return diffs 3180 3181 return _diffs(new_config, old_config), _diffs(old_config, new_config) 3182 3183 def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig: 3184 """Returns current grants for a table as a dictionary. 3185 3186 This method queries the database and returns the current grants/permissions 3187 for the given table, parsed into a dictionary format. The it handles 3188 case-insensitive comparison between these current grants and the desired 3189 grants from model configuration. 3190 3191 Args: 3192 table: The table/view to query grants for. 3193 3194 Returns: 3195 Dictionary mapping permissions to lists of grantees. Permission names 3196 should be returned as the database provides them (typically uppercase 3197 for standard SQL permissions, but engine-specific roles may vary). 3198 3199 Raises: 3200 NotImplementedError: If the engine does not support grants. 3201 """ 3202 if not self.SUPPORTS_GRANTS: 3203 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3204 raise NotImplementedError("Subclass must implement get_current_grants") 3205 3206 def _apply_grants_config_expr( 3207 self, 3208 table: exp.Table, 3209 grants_config: GrantsConfig, 3210 table_type: DataObjectType = DataObjectType.TABLE, 3211 ) -> t.List[exp.Expr]: 3212 """Returns SQLGlot Grant expressions to apply grants to a table. 3213 3214 Args: 3215 table: The table/view to grant permissions on. 3216 grants_config: Dictionary mapping permissions to lists of grantees. 3217 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 3218 3219 Returns: 3220 List of SQLGlot expressions for grant operations. 3221 3222 Raises: 3223 NotImplementedError: If the engine does not support grants. 3224 """ 3225 if not self.SUPPORTS_GRANTS: 3226 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3227 raise NotImplementedError("Subclass must implement _apply_grants_config_expr") 3228 3229 def _revoke_grants_config_expr( 3230 self, 3231 table: exp.Table, 3232 grants_config: GrantsConfig, 3233 table_type: DataObjectType = DataObjectType.TABLE, 3234 ) -> t.List[exp.Expr]: 3235 """Returns SQLGlot expressions to revoke grants from a table. 3236 3237 Args: 3238 table: The table/view to revoke permissions from. 3239 grants_config: Dictionary mapping permissions to lists of grantees. 3240 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 3241 3242 Returns: 3243 List of SQLGlot expressions for revoke operations. 3244 3245 Raises: 3246 NotImplementedError: If the engine does not support grants. 3247 """ 3248 if not self.SUPPORTS_GRANTS: 3249 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3250 raise NotImplementedError("Subclass must implement _revoke_grants_config_expr") 3251 3252 3253class EngineAdapterWithIndexSupport(EngineAdapter): 3254 SUPPORTS_INDEXES = True 3255 3256 3257def _decoded_str(value: t.Union[str, bytes]) -> str: 3258 if isinstance(value, bytes): 3259 return value.decode("utf-8") 3260 return value 3261 3262 3263def _get_data_object_cache_key(catalog: t.Optional[str], schema_name: str, object_name: str) -> str: 3264 """Returns a cache key for a data object based on its fully qualified name.""" 3265 catalog = f"{catalog}." if catalog else "" 3266 return f"{catalog}{schema_name}.{object_name}"
logger =
<Logger sqlmesh.core.engine_adapter.base (WARNING)>
MERGE_TARGET_ALIAS =
'__MERGE_TARGET__'
MERGE_SOURCE_ALIAS =
'__MERGE_SOURCE__'
KEY_FOR_CREATABLE_TYPE =
'CREATABLE_TYPE'
@set_catalog()
class
EngineAdapter:
84@set_catalog() 85class EngineAdapter: 86 """Base class wrapping a Database API compliant connection. 87 88 The EngineAdapter is an easily-subclassable interface that interacts 89 with the underlying engine and data store. 90 91 Args: 92 connection_factory_or_pool: a callable which produces a new Database API-compliant 93 connection on every call. 94 dialect: The dialect with which this adapter is associated. 95 multithreaded: Indicates whether this adapter will be used by more than one thread. 96 """ 97 98 DIALECT = "" 99 DEFAULT_BATCH_SIZE = 10000 100 DATA_OBJECT_FILTER_BATCH_SIZE = 4000 101 SUPPORTS_TRANSACTIONS = True 102 SUPPORTS_INDEXES = False 103 COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_CTAS 104 COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_AND_COMMANDS 105 MAX_TABLE_COMMENT_LENGTH: t.Optional[int] = None 106 MAX_COLUMN_COMMENT_LENGTH: t.Optional[int] = None 107 INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT 108 SUPPORTS_MATERIALIZED_VIEWS = False 109 SUPPORTS_MATERIALIZED_VIEW_SCHEMA = False 110 SUPPORTS_VIEW_SCHEMA = True 111 SUPPORTS_CLONING = False 112 SUPPORTS_MANAGED_MODELS = False 113 SUPPORTS_CREATE_DROP_CATALOG = False 114 SUPPORTED_DROP_CASCADE_OBJECT_KINDS: t.List[str] = [] 115 SCHEMA_DIFFER_KWARGS: t.Dict[str, t.Any] = {} 116 SUPPORTS_TUPLE_IN = True 117 HAS_VIEW_BINDING = False 118 RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = True 119 SUPPORTS_REPLACE_TABLE = True 120 SUPPORTS_GRANTS = False 121 DEFAULT_CATALOG_TYPE = DIALECT 122 QUOTE_IDENTIFIERS_IN_VIEWS = True 123 MAX_IDENTIFIER_LENGTH: t.Optional[int] = None 124 ATTACH_CORRELATION_ID = True 125 SUPPORTS_QUERY_EXECUTION_TRACKING = False 126 SUPPORTS_METADATA_TABLE_LAST_MODIFIED_TS = False 127 RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES: t.FrozenSet[str] = frozenset() 128 """Physical property keys whose values may contain logical model references that 129 should be resolved to physical table names during property rendering. Engines that 130 need such resolution (e.g. StarRocks' excluded_trigger_tables) override this set.""" 131 132 def __init__( 133 self, 134 connection_factory_or_pool: t.Union[t.Callable[[], t.Any], ConnectionPool], 135 dialect: str = "", 136 sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None, 137 multithreaded: bool = False, 138 cursor_init: t.Optional[t.Callable[[t.Any], None]] = None, 139 default_catalog: t.Optional[str] = None, 140 execute_log_level: int = logging.DEBUG, 141 register_comments: bool = True, 142 pre_ping: bool = False, 143 pretty_sql: bool = False, 144 shared_connection: bool = False, 145 correlation_id: t.Optional[CorrelationId] = None, 146 schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None, 147 query_execution_tracker: t.Optional[QueryExecutionTracker] = None, 148 **kwargs: t.Any, 149 ): 150 self.dialect = dialect.lower() or self.DIALECT 151 self._connection_pool = ( 152 connection_factory_or_pool 153 if isinstance(connection_factory_or_pool, ConnectionPool) 154 else create_connection_pool( 155 connection_factory_or_pool, 156 multithreaded, 157 shared_connection=shared_connection, 158 cursor_init=cursor_init, 159 ) 160 ) 161 self._sql_gen_kwargs = sql_gen_kwargs or {} 162 self._default_catalog = default_catalog 163 self._execute_log_level = execute_log_level 164 self._extra_config = kwargs 165 self._register_comments = register_comments 166 self._pre_ping = pre_ping 167 self._pretty_sql = pretty_sql 168 self._multithreaded = multithreaded 169 self.correlation_id = correlation_id 170 self._schema_differ_overrides = schema_differ_overrides 171 self._query_execution_tracker = query_execution_tracker 172 self._data_object_cache: t.Dict[str, t.Optional[DataObject]] = {} 173 174 def with_settings(self, **kwargs: t.Any) -> EngineAdapter: 175 extra_kwargs = { 176 "null_connection": True, 177 "execute_log_level": kwargs.pop("execute_log_level", self._execute_log_level), 178 "correlation_id": kwargs.pop("correlation_id", self.correlation_id), 179 "query_execution_tracker": kwargs.pop( 180 "query_execution_tracker", self._query_execution_tracker 181 ), 182 **self._extra_config, 183 **kwargs, 184 } 185 186 adapter = self.__class__( 187 self._connection_pool, 188 dialect=self.dialect, 189 sql_gen_kwargs=self._sql_gen_kwargs, 190 default_catalog=self._default_catalog, 191 register_comments=self._register_comments, 192 multithreaded=self._multithreaded, 193 pretty_sql=self._pretty_sql, 194 **extra_kwargs, 195 ) 196 197 return adapter 198 199 @property 200 def cursor(self) -> t.Any: 201 return self._connection_pool.get_cursor() 202 203 @property 204 def connection(self) -> t.Any: 205 return self._connection_pool.get() 206 207 @property 208 def spark(self) -> t.Optional[PySparkSession]: 209 return None 210 211 @property 212 def snowpark(self) -> t.Optional[SnowparkSession]: 213 return None 214 215 @property 216 def bigframe(self) -> t.Optional[BigframeSession]: 217 return None 218 219 @property 220 def comments_enabled(self) -> bool: 221 return self._register_comments and self.COMMENT_CREATION_TABLE.is_supported 222 223 @property 224 def catalog_support(self) -> CatalogSupport: 225 return CatalogSupport.UNSUPPORTED 226 227 def supports_virtual_catalog(self) -> bool: 228 """Return True if this adapter can accept a virtual catalog for multi-gateway nesting alignment. 229 230 When a project mixes catalog-aware gateways (e.g. DuckDB) with catalog-unsupported gateways 231 (e.g. ClickHouse), all adapters need a uniform 3-level FQN so MappingSchema nesting stays 232 consistent. Adapters that return True here opt in to receiving an injected virtual catalog 233 via inject_virtual_catalog(), which causes the set_catalog decorator to strip the catalog 234 from DDL expressions rather than raising UnsupportedCatalogOperationError. 235 """ 236 return False 237 238 def inject_virtual_catalog(self, gateway: str) -> None: 239 """Inject a gateway name to configure the adapter's virtual catalog. 240 241 The adapter determines the final catalog name from the gateway name (e.g. ClickHouse 242 wraps it as __{gateway}__). Only call this on adapters that return True from 243 supports_virtual_catalog(). After injection, catalog_support should return 244 SINGLE_CATALOG_ONLY so the set_catalog decorator strips the virtual catalog from DDL 245 expressions instead of raising an error. 246 """ 247 raise NotImplementedError( 248 f"{self.dialect} does not support virtual catalog injection. " 249 "Override supports_virtual_catalog() to return True and implement inject_virtual_catalog()." 250 ) 251 252 @cached_property 253 def schema_differ(self) -> SchemaDiffer: 254 return SchemaDiffer( 255 **{ 256 **self.SCHEMA_DIFFER_KWARGS, 257 **(self._schema_differ_overrides or {}), 258 } 259 ) 260 261 @property 262 def _catalog_type_overrides(self) -> t.Dict[str, str]: 263 return self._extra_config.get("catalog_type_overrides") or {} 264 265 @classmethod 266 def _casted_columns( 267 cls, 268 target_columns_to_types: t.Dict[str, exp.DataType], 269 source_columns: t.Optional[t.List[str]] = None, 270 ) -> t.List[exp.Expr]: 271 source_columns_lookup = set(source_columns or target_columns_to_types) 272 return [ 273 exp.alias_( 274 exp.cast( 275 exp.column(column, quoted=True) 276 if column in source_columns_lookup 277 else exp.Null(), 278 to=kind, 279 ), 280 column, 281 copy=False, 282 quoted=True, 283 ) 284 for column, kind in target_columns_to_types.items() 285 ] 286 287 @property 288 def default_catalog(self) -> t.Optional[str]: 289 if self.catalog_support.is_unsupported: 290 return None 291 default_catalog = self._default_catalog or self.get_current_catalog() 292 if not default_catalog: 293 raise MissingDefaultCatalogError( 294 "Could not determine a default catalog despite it being supported." 295 ) 296 return default_catalog 297 298 @property 299 def engine_run_mode(self) -> EngineRunMode: 300 return EngineRunMode.SINGLE_MODE_ENGINE 301 302 def _get_source_queries( 303 self, 304 query_or_df: QueryOrDF, 305 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 306 target_table: TableName, 307 *, 308 batch_size: t.Optional[int] = None, 309 source_columns: t.Optional[t.List[str]] = None, 310 ) -> t.List[SourceQuery]: 311 import pandas as pd 312 313 batch_size = self.DEFAULT_BATCH_SIZE if batch_size is None else batch_size 314 if isinstance(query_or_df, exp.Query): 315 query_factory = lambda: query_or_df 316 if source_columns: 317 source_columns_lookup = set(source_columns) 318 if not target_columns_to_types: 319 raise SQLMeshError("columns_to_types must be set if source_columns is set") 320 if not set(target_columns_to_types).issubset(source_columns_lookup): 321 select_columns = [ 322 exp.column(c, quoted=True) 323 if c in source_columns_lookup 324 else exp.cast(exp.Null(), target_columns_to_types[c], copy=False).as_( 325 c, copy=False, quoted=True 326 ) 327 for c in target_columns_to_types 328 ] 329 query_factory = lambda: ( 330 exp.Select() 331 .select(*select_columns) 332 .from_(query_or_df.subquery("select_source_columns")) 333 ) 334 return [SourceQuery(query_factory=query_factory)] # type: ignore 335 336 if not target_columns_to_types: 337 raise SQLMeshError( 338 "It is expected that if a DataFrame is passed in then columns_to_types is set" 339 ) 340 341 if isinstance(query_or_df, pd.DataFrame) and query_or_df.empty: 342 raise SQLMeshError( 343 "Cannot construct source query from an empty DataFrame. This error is commonly " 344 "related to Python models that produce no data. For such models, consider yielding " 345 "from an empty generator if the resulting set is empty, i.e. use `yield from ()`." 346 ) 347 348 return self._df_to_source_queries( 349 query_or_df, 350 target_columns_to_types, 351 batch_size, 352 target_table=target_table, 353 source_columns=source_columns, 354 ) 355 356 def _df_to_source_queries( 357 self, 358 df: DF, 359 target_columns_to_types: t.Dict[str, exp.DataType], 360 batch_size: int, 361 target_table: TableName, 362 source_columns: t.Optional[t.List[str]] = None, 363 ) -> t.List[SourceQuery]: 364 import pandas as pd 365 366 assert isinstance(df, pd.DataFrame) 367 num_rows = len(df.index) 368 batch_size = sys.maxsize if batch_size == 0 else batch_size 369 370 # we need to ensure that the order of the columns in columns_to_types columns matches the order of the values 371 # they can differ if a user specifies columns() on a python model in a different order than what's in the DataFrame's emitted by that model 372 df = df[list(source_columns or target_columns_to_types)] 373 values = list(df.itertuples(index=False, name=None)) 374 375 return [ 376 SourceQuery( 377 query_factory=partial( 378 self._values_to_sql, 379 values=values, # type: ignore 380 target_columns_to_types=target_columns_to_types, 381 batch_start=i, 382 batch_end=min(i + batch_size, num_rows), 383 source_columns=source_columns, 384 ), 385 ) 386 for i in range(0, num_rows, batch_size) 387 ] 388 389 def _get_source_queries_and_columns_to_types( 390 self, 391 query_or_df: QueryOrDF, 392 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 393 target_table: TableName, 394 *, 395 batch_size: t.Optional[int] = None, 396 source_columns: t.Optional[t.List[str]] = None, 397 ) -> t.Tuple[t.List[SourceQuery], t.Optional[t.Dict[str, exp.DataType]]]: 398 target_columns_to_types, source_columns = self._columns_to_types( 399 query_or_df, target_columns_to_types, source_columns 400 ) 401 source_queries = self._get_source_queries( 402 query_or_df, 403 target_columns_to_types, 404 target_table=target_table, 405 batch_size=batch_size, 406 source_columns=source_columns, 407 ) 408 return source_queries, target_columns_to_types 409 410 @t.overload 411 def _columns_to_types( 412 self, 413 query_or_df: DF, 414 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 415 source_columns: t.Optional[t.List[str]] = None, 416 ) -> t.Tuple[t.Dict[str, exp.DataType], t.List[str]]: ... 417 418 @t.overload 419 def _columns_to_types( 420 self, 421 query_or_df: Query, 422 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 423 source_columns: t.Optional[t.List[str]] = None, 424 ) -> t.Tuple[t.Optional[t.Dict[str, exp.DataType]], t.Optional[t.List[str]]]: ... 425 426 def _columns_to_types( 427 self, 428 query_or_df: QueryOrDF, 429 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 430 source_columns: t.Optional[t.List[str]] = None, 431 ) -> t.Tuple[t.Optional[t.Dict[str, exp.DataType]], t.Optional[t.List[str]]]: 432 import pandas as pd 433 434 if not target_columns_to_types and isinstance(query_or_df, pd.DataFrame): 435 target_columns_to_types = columns_to_types_from_df(t.cast(pd.DataFrame, query_or_df)) 436 if not source_columns and target_columns_to_types: 437 source_columns = list(target_columns_to_types) 438 # source columns should only contain columns that are defined in the target. If there are extras then 439 # that means they are intended to be ignored and will be excluded 440 source_columns = ( 441 [x for x in source_columns if x in target_columns_to_types] 442 if source_columns and target_columns_to_types 443 else None 444 ) 445 return target_columns_to_types, source_columns 446 447 def recycle(self) -> None: 448 """Closes all open connections and releases all allocated resources associated with any thread 449 except the calling one.""" 450 self._connection_pool.close_all(exclude_calling_thread=True) 451 452 def close(self) -> t.Any: 453 """Closes all open connections and releases all allocated resources.""" 454 self._connection_pool.close_all() 455 456 def get_current_catalog(self) -> t.Optional[str]: 457 """Returns the catalog name of the current connection.""" 458 raise NotImplementedError() 459 460 def set_current_catalog(self, catalog: str) -> None: 461 """Sets the catalog name of the current connection.""" 462 raise NotImplementedError() 463 464 def get_catalog_type(self, catalog: t.Optional[str]) -> str: 465 """Intended to be overridden for data virtualization systems like Trino that, 466 depending on the target catalog, require slightly different properties to be set when creating / updating tables 467 """ 468 if self.catalog_support.is_unsupported: 469 raise UnsupportedCatalogOperationError( 470 f"{self.dialect} does not support catalogs and a catalog was provided: {catalog}" 471 ) 472 return ( 473 self._catalog_type_overrides.get(catalog, self.DEFAULT_CATALOG_TYPE) 474 if catalog 475 else self.DEFAULT_CATALOG_TYPE 476 ) 477 478 def get_catalog_type_from_table(self, table: TableName) -> str: 479 """Get the catalog type from a table name if it has a catalog specified, otherwise return the current catalog type""" 480 catalog = exp.to_table(table).catalog or self.get_current_catalog() 481 return self.get_catalog_type(catalog) 482 483 @property 484 def current_catalog_type(self) -> str: 485 # `get_catalog_type_from_table` should be used over this property. Reason is that the table that is the target 486 # of the operation is what matters and not the catalog type of the connection. 487 # This still remains for legacy reasons and should be refactored out. 488 return self.get_catalog_type(self.get_current_catalog()) 489 490 def replace_query( 491 self, 492 table_name: TableName, 493 query_or_df: QueryOrDF, 494 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 495 table_description: t.Optional[str] = None, 496 column_descriptions: t.Optional[t.Dict[str, str]] = None, 497 source_columns: t.Optional[t.List[str]] = None, 498 supports_replace_table_override: t.Optional[bool] = None, 499 **kwargs: t.Any, 500 ) -> None: 501 """Replaces an existing table with a query. 502 503 For partition based engines (hive, spark), insert override is used. For other systems, create or replace is used. 504 505 Args: 506 table_name: The name of the table (eg. prod.table) 507 query_or_df: The SQL query to run or a dataframe. 508 target_columns_to_types: Only used if a dataframe is provided. A mapping between the column name and its data type. 509 Expected to be ordered to match the order of values in the dataframe. 510 kwargs: Optional create table properties. 511 """ 512 target_table = exp.to_table(table_name) 513 514 target_data_object = self.get_data_object(target_table) 515 table_exists = target_data_object is not None 516 if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE): 517 table_exists = False 518 519 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 520 query_or_df, 521 target_columns_to_types, 522 target_table=target_table, 523 source_columns=source_columns, 524 ) 525 if not target_columns_to_types and table_exists: 526 target_columns_to_types = self.columns(target_table) 527 query = source_queries[0].query_factory() 528 self_referencing = any( 529 quote_identifiers(table) == quote_identifiers(target_table) 530 for table in query.find_all(exp.Table) 531 ) 532 # If a query references itself then it must have a table created regardless of approach used. 533 if self_referencing: 534 if not target_columns_to_types: 535 raise SQLMeshError( 536 f"Cannot create a self-referencing table {target_table.sql(dialect=self.dialect)} without knowing the column types. " 537 "Try casting the columns to an expected type or defining the columns in the model metadata. " 538 ) 539 self._create_table_from_columns( 540 target_table, 541 target_columns_to_types, 542 exists=True, 543 table_description=table_description, 544 column_descriptions=column_descriptions, 545 **kwargs, 546 ) 547 # All engines support `CREATE TABLE AS` so we use that if the table doesn't already exist and we 548 # use `CREATE OR REPLACE TABLE AS` if the engine supports it 549 supports_replace_table = ( 550 self.SUPPORTS_REPLACE_TABLE 551 if supports_replace_table_override is None 552 else supports_replace_table_override 553 ) 554 if supports_replace_table or not table_exists: 555 return self._create_table_from_source_queries( 556 target_table, 557 source_queries, 558 target_columns_to_types, 559 replace=supports_replace_table, 560 table_description=table_description, 561 column_descriptions=column_descriptions, 562 **kwargs, 563 ) 564 if self_referencing: 565 assert target_columns_to_types is not None 566 with self.temp_table( 567 self._select_columns(target_columns_to_types).from_(target_table), 568 name=target_table, 569 target_columns_to_types=target_columns_to_types, 570 **kwargs, 571 ) as temp_table: 572 for source_query in source_queries: 573 source_query.add_transform( 574 lambda node: ( # type: ignore 575 temp_table # type: ignore 576 if isinstance(node, exp.Table) 577 and quote_identifiers(node) == quote_identifiers(target_table) 578 else node 579 ) 580 ) 581 return self._insert_overwrite_by_condition( 582 target_table, 583 source_queries, 584 target_columns_to_types, 585 **kwargs, 586 ) 587 return self._insert_overwrite_by_condition( 588 target_table, 589 source_queries, 590 target_columns_to_types, 591 **kwargs, 592 ) 593 594 def create_index( 595 self, 596 table_name: TableName, 597 index_name: str, 598 columns: t.Tuple[str, ...], 599 exists: bool = True, 600 ) -> None: 601 """Creates a new index for the given table if supported 602 603 Args: 604 table_name: The name of the target table. 605 index_name: The name of the index. 606 columns: The list of columns that constitute the index. 607 exists: Indicates whether to include the IF NOT EXISTS check. 608 """ 609 if not self.SUPPORTS_INDEXES: 610 return 611 612 expression = exp.Create( 613 this=exp.Index( 614 this=exp.to_identifier(index_name), 615 table=exp.to_table(table_name), 616 params=exp.IndexParameters(columns=[exp.to_column(c) for c in columns]), 617 ), 618 kind="INDEX", 619 exists=exists, 620 ) 621 self.execute(expression) 622 623 def _pop_creatable_type_from_properties( 624 self, 625 properties: t.Dict[str, exp.Expr], 626 ) -> t.Optional[exp.Property]: 627 """Pop out the creatable_type from the properties dictionary (if exists (return it/remove it) else return none). 628 It also checks that none of the expressions are MATERIALIZE as that conflicts with the `materialize` parameter. 629 """ 630 for key in list(properties.keys()): 631 upper_key = key.upper() 632 if upper_key == KEY_FOR_CREATABLE_TYPE: 633 value = properties.pop(key).name 634 parsed_properties = exp.maybe_parse( 635 value, into=exp.Properties, dialect=self.dialect 636 ) 637 property, *others = parsed_properties.expressions 638 if others: 639 # Multiple properties are unsupported today, can look into it in the future if needed 640 raise SQLMeshError( 641 f"Invalid creatable_type value with multiple properties: {value}" 642 ) 643 if isinstance(property, exp.MaterializedProperty): 644 raise SQLMeshError( 645 f"Cannot use {value} as a creatable_type as it conflicts with the `materialize` parameter." 646 ) 647 return property 648 return None 649 650 def create_table( 651 self, 652 table_name: TableName, 653 target_columns_to_types: t.Dict[str, exp.DataType], 654 primary_key: t.Optional[t.Tuple[str, ...]] = None, 655 exists: bool = True, 656 table_description: t.Optional[str] = None, 657 column_descriptions: t.Optional[t.Dict[str, str]] = None, 658 **kwargs: t.Any, 659 ) -> None: 660 """Create a table using a DDL statement 661 662 Args: 663 table_name: The name of the table to create. Can be fully qualified or just table name. 664 target_columns_to_types: A mapping between the column name and its data type. 665 primary_key: Determines the table primary key. 666 exists: Indicates whether to include the IF NOT EXISTS check. 667 table_description: Optional table description from MODEL DDL. 668 column_descriptions: Optional column descriptions from model query. 669 kwargs: Optional create table properties. 670 """ 671 self._create_table_from_columns( 672 table_name, 673 target_columns_to_types, 674 primary_key, 675 exists, 676 table_description, 677 column_descriptions, 678 **kwargs, 679 ) 680 681 def create_managed_table( 682 self, 683 table_name: TableName, 684 query: Query, 685 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 686 partitioned_by: t.Optional[t.List[exp.Expr]] = None, 687 clustered_by: t.Optional[t.List[exp.Expr]] = None, 688 table_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 689 table_description: t.Optional[str] = None, 690 column_descriptions: t.Optional[t.Dict[str, str]] = None, 691 source_columns: t.Optional[t.List[str]] = None, 692 **kwargs: t.Any, 693 ) -> None: 694 """Create a managed table using a query. 695 696 "Managed" means that once the table is created, the data is kept up to date by the underlying database engine and not SQLMesh. 697 698 Args: 699 table_name: The name of the table to create. Can be fully qualified or just table name. 700 query: The SQL query for the engine to base the managed table on 701 target_columns_to_types: A mapping between the column name and its data type. 702 partitioned_by: The partition columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 703 clustered_by: The cluster columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 704 table_properties: Optional mapping of engine-specific properties to be set on the managed table 705 table_description: Optional table description from MODEL DDL. 706 column_descriptions: Optional column descriptions from model query. 707 kwargs: Optional create table properties. 708 """ 709 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}") 710 711 def ctas( 712 self, 713 table_name: TableName, 714 query_or_df: QueryOrDF, 715 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 716 exists: bool = True, 717 table_description: t.Optional[str] = None, 718 column_descriptions: t.Optional[t.Dict[str, str]] = None, 719 source_columns: t.Optional[t.List[str]] = None, 720 **kwargs: t.Any, 721 ) -> None: 722 """Create a table using a CTAS statement 723 724 Args: 725 table_name: The name of the table to create. Can be fully qualified or just table name. 726 query_or_df: The SQL query to run or a dataframe for the CTAS. 727 target_columns_to_types: A mapping between the column name and its data type. Required if using a DataFrame. 728 exists: Indicates whether to include the IF NOT EXISTS check. 729 table_description: Optional table description from MODEL DDL. 730 column_descriptions: Optional column descriptions from model query. 731 kwargs: Optional create table properties. 732 """ 733 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 734 query_or_df, 735 target_columns_to_types, 736 target_table=table_name, 737 source_columns=source_columns, 738 ) 739 return self._create_table_from_source_queries( 740 table_name, 741 source_queries, 742 target_columns_to_types, 743 exists, 744 table_description=table_description, 745 column_descriptions=column_descriptions, 746 **kwargs, 747 ) 748 749 def create_state_table( 750 self, 751 table_name: str, 752 target_columns_to_types: t.Dict[str, exp.DataType], 753 primary_key: t.Optional[t.Tuple[str, ...]] = None, 754 ) -> None: 755 """Create a table to store SQLMesh internal state. 756 757 Args: 758 table_name: The name of the table to create. Can be fully qualified or just table name. 759 target_columns_to_types: A mapping between the column name and its data type. 760 primary_key: Determines the table primary key. 761 """ 762 self.create_table( 763 table_name, 764 target_columns_to_types, 765 primary_key=primary_key, 766 ) 767 768 def _create_table_from_columns( 769 self, 770 table_name: TableName, 771 target_columns_to_types: t.Dict[str, exp.DataType], 772 primary_key: t.Optional[t.Tuple[str, ...]] = None, 773 exists: bool = True, 774 table_description: t.Optional[str] = None, 775 column_descriptions: t.Optional[t.Dict[str, str]] = None, 776 **kwargs: t.Any, 777 ) -> None: 778 """ 779 Create a table using a DDL statement. 780 781 Args: 782 table_name: The name of the table to create. Can be fully qualified or just table name. 783 target_columns_to_types: Mapping between the column name and its data type. 784 primary_key: Determines the table primary key. 785 exists: Indicates whether to include the IF NOT EXISTS check. 786 table_description: Optional table description from MODEL DDL. 787 column_descriptions: Optional column descriptions from model query. 788 kwargs: Optional create table properties. 789 """ 790 table = exp.to_table(table_name) 791 792 if not columns_to_types_all_known(target_columns_to_types): 793 # It is ok if the columns types are not known if the table already exists and IF NOT EXISTS is set 794 if exists and self.table_exists(table_name): 795 return 796 raise SQLMeshError( 797 "Cannot create a table without knowing the column types. " 798 "Try casting the columns to an expected type or defining the columns in the model metadata. " 799 f"Columns to types: {target_columns_to_types}" 800 ) 801 802 primary_key_expression = ( 803 [exp.PrimaryKey(expressions=[exp.to_column(k) for k in primary_key])] 804 if primary_key and self.SUPPORTS_INDEXES 805 else [] 806 ) 807 808 schema = self._build_schema_exp( 809 table, 810 target_columns_to_types, 811 column_descriptions, 812 primary_key_expression, 813 ) 814 815 self._create_table( 816 schema, 817 None, 818 exists=exists, 819 target_columns_to_types=target_columns_to_types, 820 table_description=table_description, 821 **kwargs, 822 ) 823 824 # Register comments with commands if the engine doesn't support comments in the schema or CREATE 825 if ( 826 table_description 827 and self.COMMENT_CREATION_TABLE.is_comment_command_only 828 and self.comments_enabled 829 ): 830 self._create_table_comment(table_name, table_description) 831 if ( 832 column_descriptions 833 and self.COMMENT_CREATION_TABLE.is_comment_command_only 834 and self.comments_enabled 835 ): 836 self._create_column_comments(table_name, column_descriptions) 837 838 def _build_schema_exp( 839 self, 840 table: exp.Table, 841 target_columns_to_types: t.Dict[str, exp.DataType], 842 column_descriptions: t.Optional[t.Dict[str, str]] = None, 843 expressions: t.Optional[t.List[exp.PrimaryKey]] = None, 844 is_view: bool = False, 845 materialized: bool = False, 846 ) -> exp.Schema: 847 """ 848 Build a schema expression for a table, columns, column comments, and additional schema properties. 849 """ 850 expressions = expressions or [] 851 852 return exp.Schema( 853 this=table, 854 expressions=self._build_column_defs( 855 target_columns_to_types=target_columns_to_types, 856 column_descriptions=column_descriptions, 857 is_view=is_view, 858 materialized=materialized, 859 ) 860 + expressions, 861 ) 862 863 def _build_column_defs( 864 self, 865 target_columns_to_types: t.Dict[str, exp.DataType], 866 column_descriptions: t.Optional[t.Dict[str, str]] = None, 867 is_view: bool = False, 868 materialized: bool = False, 869 ) -> t.List[exp.ColumnDef]: 870 engine_supports_schema_comments = ( 871 self.COMMENT_CREATION_VIEW.supports_schema_def 872 if is_view 873 else self.COMMENT_CREATION_TABLE.supports_schema_def 874 ) 875 return [ 876 self._build_column_def( 877 column, 878 column_descriptions=column_descriptions, 879 engine_supports_schema_comments=engine_supports_schema_comments, 880 col_type=None if is_view else kind, # don't include column data type for views 881 ) 882 for column, kind in target_columns_to_types.items() 883 ] 884 885 def _build_column_def( 886 self, 887 col_name: str, 888 column_descriptions: t.Optional[t.Dict[str, str]] = None, 889 engine_supports_schema_comments: bool = False, 890 col_type: t.Optional[exp.DATA_TYPE] = None, 891 nested_names: t.List[str] = [], 892 ) -> exp.ColumnDef: 893 return exp.ColumnDef( 894 this=exp.to_identifier(col_name), 895 kind=col_type, 896 constraints=( 897 self._build_col_comment_exp(col_name, column_descriptions) 898 if engine_supports_schema_comments and self.comments_enabled and column_descriptions 899 else None 900 ), 901 ) 902 903 def _build_col_comment_exp( 904 self, col_name: str, column_descriptions: t.Dict[str, str] 905 ) -> t.List[exp.ColumnConstraint]: 906 comment = column_descriptions.get(col_name, None) 907 if comment: 908 return [ 909 exp.ColumnConstraint( 910 kind=exp.CommentColumnConstraint( 911 this=exp.Literal.string(self._truncate_column_comment(comment)) 912 ) 913 ) 914 ] 915 return [] 916 917 def _create_table_from_source_queries( 918 self, 919 table_name: TableName, 920 source_queries: t.List[SourceQuery], 921 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 922 exists: bool = True, 923 replace: bool = False, 924 table_description: t.Optional[str] = None, 925 column_descriptions: t.Optional[t.Dict[str, str]] = None, 926 table_kind: t.Optional[str] = None, 927 track_rows_processed: bool = True, 928 **kwargs: t.Any, 929 ) -> None: 930 table = exp.to_table(table_name) 931 932 # CTAS calls do not usually include a schema expression. However, most engines 933 # permit them in CTAS expressions, and they allow us to register all column comments 934 # in a single call rather than in a separate comment command call for each column. 935 # 936 # This block conditionally builds a schema expression with column comments if the engine 937 # supports it and we have columns_to_types. column_to_types is required because the 938 # schema expression must include at least column name, data type, and the comment - 939 # for example, `(colname INTEGER COMMENT 'comment')`. 940 # 941 # column_to_types will be available when loading from a DataFrame (by converting from 942 # pandas to SQL types), when a model is "annotated" by explicitly specifying column 943 # types, and for evaluation methods like `LogicalReplaceQueryMixin.replace_query()` 944 # calls and SCD Type 2 model calls. 945 schema = None 946 target_columns_to_types_known = target_columns_to_types and columns_to_types_all_known( 947 target_columns_to_types 948 ) 949 if ( 950 column_descriptions 951 and target_columns_to_types_known 952 and self.COMMENT_CREATION_TABLE.is_in_schema_def_ctas 953 and self.comments_enabled 954 ): 955 schema = self._build_schema_exp(table, target_columns_to_types, column_descriptions) # type: ignore 956 957 with self.transaction(condition=len(source_queries) > 1): 958 for i, source_query in enumerate(source_queries): 959 with source_query as query: 960 if target_columns_to_types and target_columns_to_types_known: 961 query = self._order_projections_and_filter( 962 query, target_columns_to_types, coerce_types=True 963 ) 964 if i == 0: 965 self._create_table( 966 schema if schema else table, 967 query, 968 target_columns_to_types=target_columns_to_types, 969 exists=exists, 970 replace=replace, 971 table_description=table_description, 972 table_kind=table_kind, 973 track_rows_processed=track_rows_processed, 974 **kwargs, 975 ) 976 else: 977 self._insert_append_query( 978 table_name, 979 query, 980 target_columns_to_types or self.columns(table), 981 track_rows_processed=track_rows_processed, 982 ) 983 984 # Register comments with commands if the engine supports comments and we weren't able to 985 # register them with the CTAS call's schema expression. 986 if ( 987 table_description 988 and self.COMMENT_CREATION_TABLE.is_comment_command_only 989 and self.comments_enabled 990 ): 991 self._create_table_comment(table_name, table_description) 992 if column_descriptions and schema is None and self.comments_enabled: 993 self._create_column_comments(table_name, column_descriptions) 994 995 def _create_table( 996 self, 997 table_name_or_schema: t.Union[exp.Schema, TableName], 998 expression: t.Optional[exp.Expr], 999 exists: bool = True, 1000 replace: bool = False, 1001 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1002 table_description: t.Optional[str] = None, 1003 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1004 table_kind: t.Optional[str] = None, 1005 track_rows_processed: bool = True, 1006 **kwargs: t.Any, 1007 ) -> None: 1008 self.execute( 1009 self._build_create_table_exp( 1010 table_name_or_schema, 1011 expression=expression, 1012 exists=exists, 1013 replace=replace, 1014 target_columns_to_types=target_columns_to_types, 1015 table_description=( 1016 table_description 1017 if self.COMMENT_CREATION_TABLE.supports_schema_def and self.comments_enabled 1018 else None 1019 ), 1020 table_kind=table_kind, 1021 **kwargs, 1022 ), 1023 track_rows_processed=track_rows_processed, 1024 ) 1025 # Extract table name to clear cache 1026 table_name = ( 1027 table_name_or_schema.this 1028 if isinstance(table_name_or_schema, exp.Schema) 1029 else table_name_or_schema 1030 ) 1031 self._clear_data_object_cache(table_name) 1032 1033 def _build_create_table_exp( 1034 self, 1035 table_name_or_schema: t.Union[exp.Schema, TableName], 1036 expression: t.Optional[exp.Expr], 1037 exists: bool = True, 1038 replace: bool = False, 1039 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1040 table_description: t.Optional[str] = None, 1041 table_kind: t.Optional[str] = None, 1042 **kwargs: t.Any, 1043 ) -> exp.Create: 1044 exists = False if replace else exists 1045 catalog_name = None 1046 if not isinstance(table_name_or_schema, exp.Schema): 1047 table_name_or_schema = exp.to_table(table_name_or_schema) 1048 catalog_name = table_name_or_schema.catalog 1049 else: 1050 if isinstance(table_name_or_schema.this, exp.Table): 1051 catalog_name = table_name_or_schema.this.catalog 1052 1053 properties = ( 1054 self._build_table_properties_exp( 1055 **kwargs, 1056 catalog_name=catalog_name, 1057 target_columns_to_types=target_columns_to_types, 1058 table_description=table_description, 1059 table_kind=table_kind, 1060 ) 1061 if kwargs or table_description 1062 else None 1063 ) 1064 return exp.Create( 1065 this=table_name_or_schema, 1066 kind=table_kind or "TABLE", 1067 replace=replace, 1068 exists=exists, 1069 expression=expression, 1070 properties=properties, 1071 ) 1072 1073 def create_table_like( 1074 self, 1075 target_table_name: TableName, 1076 source_table_name: TableName, 1077 exists: bool = True, 1078 **kwargs: t.Any, 1079 ) -> None: 1080 """Create a table to store SQLMesh internal state based on the definition of another table, including any 1081 column attributes and indexes defined in the original table. 1082 1083 Args: 1084 target_table_name: The name of the table to create. Can be fully qualified or just table name. 1085 source_table_name: The name of the table to base the new table on. 1086 """ 1087 self._create_table_like(target_table_name, source_table_name, exists=exists, **kwargs) 1088 self._clear_data_object_cache(target_table_name) 1089 1090 def clone_table( 1091 self, 1092 target_table_name: TableName, 1093 source_table_name: TableName, 1094 replace: bool = False, 1095 exists: bool = True, 1096 clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, 1097 **kwargs: t.Any, 1098 ) -> None: 1099 """Creates a table with the target name by cloning the source table. 1100 1101 Args: 1102 target_table_name: The name of the table that should be created. 1103 source_table_name: The name of the source table that should be cloned. 1104 replace: Whether or not to replace an existing table. 1105 exists: Indicates whether to include the IF NOT EXISTS check. 1106 """ 1107 if not self.SUPPORTS_CLONING: 1108 raise NotImplementedError(f"Engine does not support cloning: {type(self)}") 1109 1110 kwargs.pop("rendered_physical_properties", None) 1111 self.execute( 1112 exp.Create( 1113 this=exp.to_table(target_table_name), 1114 kind="TABLE", 1115 replace=replace, 1116 exists=exists, 1117 clone=exp.Clone( 1118 this=exp.to_table(source_table_name), 1119 **(clone_kwargs or {}), 1120 ), 1121 **kwargs, 1122 ) 1123 ) 1124 self._clear_data_object_cache(target_table_name) 1125 1126 def drop_data_object(self, data_object: DataObject, ignore_if_not_exists: bool = True) -> None: 1127 """Drops a data object of arbitrary type. 1128 1129 Args: 1130 data_object: The data object to drop. 1131 ignore_if_not_exists: If True, no error will be raised if the data object does not exist. 1132 """ 1133 if data_object.type.is_view: 1134 self.drop_view(data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists) 1135 elif data_object.type.is_materialized_view: 1136 self.drop_view( 1137 data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists, materialized=True 1138 ) 1139 elif data_object.type.is_table: 1140 self.drop_table(data_object.to_table(), exists=ignore_if_not_exists) 1141 elif data_object.type.is_managed_table: 1142 self.drop_managed_table(data_object.to_table(), exists=ignore_if_not_exists) 1143 else: 1144 raise SQLMeshError( 1145 f"Can't drop data object '{data_object.to_table().sql(dialect=self.dialect)}' of type '{data_object.type.value}'" 1146 ) 1147 1148 def drop_table(self, table_name: TableName, exists: bool = True, **kwargs: t.Any) -> None: 1149 """Drops a table. 1150 1151 Args: 1152 table_name: The name of the table to drop. 1153 exists: If exists, defaults to True. 1154 """ 1155 self._drop_object(name=table_name, exists=exists, **kwargs) 1156 1157 def drop_managed_table(self, table_name: TableName, exists: bool = True) -> None: 1158 """Drops a managed table. 1159 1160 Args: 1161 table_name: The name of the table to drop. 1162 exists: If exists, defaults to True. 1163 """ 1164 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}") 1165 1166 def _drop_object( 1167 self, 1168 name: TableName | SchemaName, 1169 exists: bool = True, 1170 kind: str = "TABLE", 1171 cascade: bool = False, 1172 **drop_args: t.Any, 1173 ) -> None: 1174 """Drops an object. 1175 1176 An object could be a DATABASE, SCHEMA, VIEW, TABLE, DYNAMIC TABLE, TEMPORARY TABLE etc depending on the :kind. 1177 1178 Args: 1179 name: The name of the table to drop. 1180 exists: If exists, defaults to True. 1181 kind: What kind of object to drop. Defaults to TABLE 1182 cascade: Whether or not to DROP ... CASCADE. 1183 Note that this is ignored for :kind's that are not present in self.SUPPORTED_DROP_CASCADE_OBJECT_KINDS 1184 **drop_args: Any extra arguments to set on the Drop expression 1185 """ 1186 if cascade and kind.upper() in self.SUPPORTED_DROP_CASCADE_OBJECT_KINDS: 1187 drop_args["cascade"] = cascade 1188 1189 self.execute(exp.Drop(this=exp.to_table(name), kind=kind, exists=exists, **drop_args)) 1190 self._clear_data_object_cache(name) 1191 1192 def get_alter_operations( 1193 self, 1194 current_table_name: TableName, 1195 target_table_name: TableName, 1196 *, 1197 ignore_destructive: bool = False, 1198 ignore_additive: bool = False, 1199 ) -> t.List[TableAlterOperation]: 1200 """ 1201 Determines the alter statements needed to change the current table into the structure of the target table. 1202 """ 1203 return t.cast( 1204 t.List[TableAlterOperation], 1205 self.schema_differ.compare_columns( 1206 current_table_name, 1207 self.columns(current_table_name), 1208 self.columns(target_table_name), 1209 ignore_destructive=ignore_destructive, 1210 ignore_additive=ignore_additive, 1211 ), 1212 ) 1213 1214 def alter_table( 1215 self, 1216 alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], 1217 ) -> None: 1218 """ 1219 Performs the alter statements to change the current table into the structure of the target table. 1220 """ 1221 with self.transaction(): 1222 for alter_expression in [ 1223 x.expression if isinstance(x, TableAlterOperation) else x for x in alter_expressions 1224 ]: 1225 self.execute(alter_expression) 1226 1227 def create_view( 1228 self, 1229 view_name: TableName, 1230 query_or_df: QueryOrDF, 1231 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1232 replace: bool = True, 1233 materialized: bool = False, 1234 materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, 1235 table_description: t.Optional[str] = None, 1236 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1237 view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 1238 source_columns: t.Optional[t.List[str]] = None, 1239 **create_kwargs: t.Any, 1240 ) -> None: 1241 """Create a view with a query or dataframe. 1242 1243 If a dataframe is passed in, it will be converted into a literal values statement. 1244 This should only be done if the dataframe is very small! 1245 1246 Args: 1247 view_name: The view name. 1248 query_or_df: A query or dataframe. 1249 target_columns_to_types: Columns to use in the view statement. 1250 replace: Whether or not to replace an existing view defaults to True. 1251 materialized: Whether to create a a materialized view. Only used for engines that support this feature. 1252 materialized_properties: Optional materialized view properties to add to the view. 1253 table_description: Optional table description from MODEL DDL. 1254 column_descriptions: Optional column descriptions from model query. 1255 view_properties: Optional view properties to add to the view. 1256 create_kwargs: Additional kwargs to pass into the Create expression 1257 """ 1258 import pandas as pd 1259 1260 if materialized_properties and not materialized: 1261 raise SQLMeshError("Materialized properties are only supported for materialized views") 1262 1263 query_or_df = self._native_df_to_pandas_df(query_or_df) 1264 1265 if isinstance(query_or_df, pd.DataFrame): 1266 values: t.List[t.Tuple[t.Any, ...]] = list( 1267 query_or_df.itertuples(index=False, name=None) 1268 ) 1269 target_columns_to_types, source_columns = self._columns_to_types( 1270 query_or_df, target_columns_to_types, source_columns 1271 ) 1272 if not target_columns_to_types: 1273 raise SQLMeshError("columns_to_types must be provided for dataframes") 1274 source_columns_to_types = get_source_columns_to_types( 1275 target_columns_to_types, source_columns 1276 ) 1277 query_or_df = self._values_to_sql( 1278 values, 1279 source_columns_to_types, 1280 batch_start=0, 1281 batch_end=len(values), 1282 ) 1283 1284 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1285 query_or_df, 1286 target_columns_to_types, 1287 batch_size=0, 1288 target_table=view_name, 1289 source_columns=source_columns, 1290 ) 1291 if len(source_queries) != 1: 1292 raise SQLMeshError("Only one source query is supported for creating views") 1293 1294 schema: t.Union[exp.Table, exp.Schema] = exp.to_table(view_name) 1295 if target_columns_to_types: 1296 schema = self._build_schema_exp( 1297 exp.to_table(view_name), 1298 target_columns_to_types, 1299 column_descriptions, 1300 is_view=True, 1301 materialized=materialized, 1302 ) 1303 1304 properties = create_kwargs.pop("properties", None) 1305 if not properties: 1306 properties = exp.Properties(expressions=[]) 1307 1308 if view_properties: 1309 table_type = self._pop_creatable_type_from_properties(view_properties) 1310 if table_type: 1311 properties.append("expressions", table_type) 1312 1313 if materialized and self.SUPPORTS_MATERIALIZED_VIEWS: 1314 properties.append("expressions", exp.MaterializedProperty()) 1315 1316 if not self.SUPPORTS_MATERIALIZED_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1317 schema = schema.this 1318 1319 if not self.SUPPORTS_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1320 schema = schema.this 1321 1322 if materialized_properties: 1323 partitioned_by = materialized_properties.pop("partitioned_by", None) 1324 clustered_by = materialized_properties.pop("clustered_by", None) 1325 if ( 1326 partitioned_by 1327 and ( 1328 partitioned_by_prop := self._build_partitioned_by_exp( 1329 partitioned_by, **materialized_properties 1330 ) 1331 ) 1332 is not None 1333 ): 1334 materialized_properties["catalog_name"] = exp.to_table(view_name).catalog 1335 properties.append("expressions", partitioned_by_prop) 1336 if ( 1337 clustered_by 1338 and ( 1339 clustered_by_prop := self._build_clustered_by_exp( 1340 clustered_by, **materialized_properties 1341 ) 1342 ) 1343 is not None 1344 ): 1345 properties.append("expressions", clustered_by_prop) 1346 1347 create_view_properties = self._build_view_properties_exp( 1348 view_properties, 1349 ( 1350 table_description 1351 if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled 1352 else None 1353 ), 1354 physical_cluster=create_kwargs.pop("physical_cluster", None), 1355 ) 1356 if create_view_properties: 1357 for view_property in create_view_properties.expressions: 1358 # Small hack to make sure SECURE goes at the beginning before materialized as required by Snowflake 1359 if isinstance(view_property, exp.SecureProperty): 1360 properties.set("expressions", view_property, index=0, overwrite=False) 1361 else: 1362 properties.append("expressions", view_property) 1363 1364 if properties.expressions: 1365 create_kwargs["properties"] = properties 1366 1367 if replace: 1368 self.drop_data_object_on_type_mismatch( 1369 self.get_data_object(view_name), 1370 DataObjectType.VIEW if not materialized else DataObjectType.MATERIALIZED_VIEW, 1371 ) 1372 1373 with source_queries[0] as query: 1374 self.execute( 1375 exp.Create( 1376 this=schema, 1377 kind="VIEW", 1378 replace=replace, 1379 expression=query, 1380 **create_kwargs, 1381 ), 1382 quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, 1383 ) 1384 1385 self._clear_data_object_cache(view_name) 1386 1387 # Register table comment with commands if the engine doesn't support doing it in CREATE 1388 if ( 1389 table_description 1390 and self.COMMENT_CREATION_VIEW.is_comment_command_only 1391 and self.comments_enabled 1392 ): 1393 self._create_table_comment(view_name, table_description, "VIEW") 1394 # Register column comments with commands if the engine doesn't support doing it in 1395 # CREATE or we couldn't do it in the CREATE schema definition because we don't have 1396 # columns_to_types 1397 if ( 1398 column_descriptions 1399 and ( 1400 self.COMMENT_CREATION_VIEW.is_comment_command_only 1401 or ( 1402 self.COMMENT_CREATION_VIEW.is_in_schema_def_and_commands 1403 and not target_columns_to_types 1404 ) 1405 ) 1406 and self.comments_enabled 1407 ): 1408 self._create_column_comments(view_name, column_descriptions, "VIEW", materialized) 1409 1410 @set_catalog() 1411 def create_schema( 1412 self, 1413 schema_name: SchemaName, 1414 ignore_if_exists: bool = True, 1415 warn_on_error: bool = True, 1416 properties: t.Optional[t.List[exp.Expr]] = None, 1417 ) -> None: 1418 properties = properties or [] 1419 return self._create_schema( 1420 schema_name=schema_name, 1421 ignore_if_exists=ignore_if_exists, 1422 warn_on_error=warn_on_error, 1423 properties=properties, 1424 kind="SCHEMA", 1425 ) 1426 1427 def _create_schema( 1428 self, 1429 schema_name: SchemaName, 1430 ignore_if_exists: bool, 1431 warn_on_error: bool, 1432 properties: t.List[exp.Expr], 1433 kind: str, 1434 ) -> None: 1435 """Create a schema from a name or qualified table name.""" 1436 try: 1437 self.execute( 1438 exp.Create( 1439 this=to_schema(schema_name), 1440 kind=kind, 1441 exists=ignore_if_exists, 1442 properties=exp.Properties( # this renders as '' (empty string) if expressions is empty 1443 expressions=properties 1444 ), 1445 ) 1446 ) 1447 except Exception as e: 1448 if not warn_on_error: 1449 raise 1450 logger.warning("Failed to create %s '%s': %s", kind.lower(), schema_name, e) 1451 1452 def drop_schema( 1453 self, 1454 schema_name: SchemaName, 1455 ignore_if_not_exists: bool = True, 1456 cascade: bool = False, 1457 **drop_args: t.Dict[str, exp.Expr], 1458 ) -> None: 1459 return self._drop_object( 1460 name=schema_name, 1461 exists=ignore_if_not_exists, 1462 kind="SCHEMA", 1463 cascade=cascade, 1464 **drop_args, 1465 ) 1466 1467 def drop_view( 1468 self, 1469 view_name: TableName, 1470 ignore_if_not_exists: bool = True, 1471 materialized: bool = False, 1472 **kwargs: t.Any, 1473 ) -> None: 1474 """Drop a view.""" 1475 self._drop_object( 1476 name=view_name, 1477 exists=ignore_if_not_exists, 1478 kind="VIEW", 1479 materialized=materialized and self.SUPPORTS_MATERIALIZED_VIEWS, 1480 **kwargs, 1481 ) 1482 1483 def create_catalog(self, catalog_name: str | exp.Identifier) -> None: 1484 return self._create_catalog(exp.parse_identifier(catalog_name, dialect=self.dialect)) 1485 1486 def _create_catalog(self, catalog_name: exp.Identifier) -> None: 1487 raise SQLMeshError( 1488 f"Unable to create catalog '{catalog_name.sql(dialect=self.dialect)}' as automatic catalog management is not implemented in the {self.dialect} engine." 1489 ) 1490 1491 def drop_catalog(self, catalog_name: str | exp.Identifier) -> None: 1492 return self._drop_catalog(exp.parse_identifier(catalog_name, dialect=self.dialect)) 1493 1494 def _drop_catalog(self, catalog_name: exp.Identifier) -> None: 1495 raise SQLMeshError( 1496 f"Unable to drop catalog '{catalog_name.sql(dialect=self.dialect)}' as automatic catalog management is not implemented in the {self.dialect} engine." 1497 ) 1498 1499 def columns( 1500 self, table_name: TableName, include_pseudo_columns: bool = False 1501 ) -> t.Dict[str, exp.DataType]: 1502 """Fetches column names and types for the target table.""" 1503 self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) 1504 describe_output = self.cursor.fetchall() 1505 return { 1506 # Note: MySQL returns the column type as bytes. 1507 column_name: exp.DataType.build(_decoded_str(column_type), dialect=self.dialect) 1508 for column_name, column_type, *_ in itertools.takewhile( 1509 lambda t: not t[0].startswith("#"), 1510 describe_output, 1511 ) 1512 if column_name and column_name.strip() and column_type and column_type.strip() 1513 } 1514 1515 def table_exists(self, table_name: TableName) -> bool: 1516 table = exp.to_table(table_name) 1517 data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name) 1518 if data_object_cache_key in self._data_object_cache: 1519 logger.debug("Table existence cache hit: %s", data_object_cache_key) 1520 return self._data_object_cache[data_object_cache_key] is not None 1521 1522 try: 1523 self.execute(exp.Describe(this=table, kind="TABLE")) 1524 return True 1525 except Exception: 1526 return False 1527 1528 def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> None: 1529 self.execute(exp.delete(table_name, where)) 1530 1531 def insert_append( 1532 self, 1533 table_name: TableName, 1534 query_or_df: QueryOrDF, 1535 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1536 track_rows_processed: bool = True, 1537 source_columns: t.Optional[t.List[str]] = None, 1538 ) -> None: 1539 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1540 query_or_df, 1541 target_columns_to_types, 1542 target_table=table_name, 1543 source_columns=source_columns, 1544 ) 1545 self._insert_append_source_queries( 1546 table_name, source_queries, target_columns_to_types, track_rows_processed 1547 ) 1548 1549 def _insert_append_source_queries( 1550 self, 1551 table_name: TableName, 1552 source_queries: t.List[SourceQuery], 1553 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1554 track_rows_processed: bool = True, 1555 ) -> None: 1556 with self.transaction(condition=len(source_queries) > 0): 1557 target_columns_to_types = target_columns_to_types or self.columns(table_name) 1558 for source_query in source_queries: 1559 with source_query as query: 1560 self._insert_append_query( 1561 table_name, 1562 query, 1563 target_columns_to_types, 1564 track_rows_processed=track_rows_processed, 1565 ) 1566 1567 def _insert_append_query( 1568 self, 1569 table_name: TableName, 1570 query: Query, 1571 target_columns_to_types: t.Dict[str, exp.DataType], 1572 order_projections: bool = True, 1573 track_rows_processed: bool = True, 1574 ) -> None: 1575 if order_projections: 1576 query = self._order_projections_and_filter(query, target_columns_to_types) 1577 self.execute( 1578 exp.insert(query, table_name, columns=list(target_columns_to_types)), 1579 track_rows_processed=track_rows_processed, 1580 ) 1581 1582 def insert_overwrite_by_partition( 1583 self, 1584 table_name: TableName, 1585 query_or_df: QueryOrDF, 1586 partitioned_by: t.List[exp.Expr], 1587 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1588 source_columns: t.Optional[t.List[str]] = None, 1589 ) -> None: 1590 if self.INSERT_OVERWRITE_STRATEGY.is_insert_overwrite: 1591 target_table = exp.to_table(table_name) 1592 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1593 query_or_df, 1594 target_columns_to_types, 1595 target_table=target_table, 1596 source_columns=source_columns, 1597 ) 1598 self._insert_overwrite_by_condition( 1599 table_name, source_queries, target_columns_to_types=target_columns_to_types 1600 ) 1601 else: 1602 self._replace_by_key( 1603 table_name, 1604 query_or_df, 1605 target_columns_to_types, 1606 partitioned_by, 1607 is_unique_key=False, 1608 source_columns=source_columns, 1609 ) 1610 1611 def insert_overwrite_by_time_partition( 1612 self, 1613 table_name: TableName, 1614 query_or_df: QueryOrDF, 1615 start: TimeLike, 1616 end: TimeLike, 1617 time_formatter: t.Callable[[TimeLike, t.Optional[t.Dict[str, exp.DataType]]], exp.Expr], 1618 time_column: TimeColumn | exp.Expr | str, 1619 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1620 source_columns: t.Optional[t.List[str]] = None, 1621 **kwargs: t.Any, 1622 ) -> None: 1623 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1624 query_or_df, 1625 target_columns_to_types, 1626 target_table=table_name, 1627 source_columns=source_columns, 1628 ) 1629 if not target_columns_to_types or not columns_to_types_all_known(target_columns_to_types): 1630 target_columns_to_types = self.columns(table_name) 1631 low, high = [ 1632 time_formatter(dt, target_columns_to_types) 1633 for dt in make_inclusive(start, end, self.dialect) 1634 ] 1635 if isinstance(time_column, TimeColumn): 1636 time_column = time_column.column 1637 where = exp.Between( 1638 this=exp.to_column(time_column) if isinstance(time_column, str) else time_column, 1639 low=low, 1640 high=high, 1641 ) 1642 return self._insert_overwrite_by_time_partition( 1643 table_name, source_queries, target_columns_to_types, where, **kwargs 1644 ) 1645 1646 def _insert_overwrite_by_time_partition( 1647 self, 1648 table_name: TableName, 1649 source_queries: t.List[SourceQuery], 1650 target_columns_to_types: t.Dict[str, exp.DataType], 1651 where: exp.Condition, 1652 **kwargs: t.Any, 1653 ) -> None: 1654 return self._insert_overwrite_by_condition( 1655 table_name, source_queries, target_columns_to_types, where, **kwargs 1656 ) 1657 1658 def _values_to_sql( 1659 self, 1660 values: t.List[t.Tuple[t.Any, ...]], 1661 target_columns_to_types: t.Dict[str, exp.DataType], 1662 batch_start: int, 1663 batch_end: int, 1664 alias: str = "t", 1665 source_columns: t.Optional[t.List[str]] = None, 1666 ) -> Query: 1667 return select_from_values_for_batch_range( 1668 values=values, 1669 target_columns_to_types=target_columns_to_types, 1670 batch_start=batch_start, 1671 batch_end=batch_end, 1672 alias=alias, 1673 source_columns=source_columns, 1674 ) 1675 1676 def _insert_overwrite_by_condition( 1677 self, 1678 table_name: TableName, 1679 source_queries: t.List[SourceQuery], 1680 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1681 where: t.Optional[exp.Condition] = None, 1682 insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None, 1683 **kwargs: t.Any, 1684 ) -> None: 1685 table = exp.to_table(table_name) 1686 insert_overwrite_strategy = ( 1687 insert_overwrite_strategy_override or self.INSERT_OVERWRITE_STRATEGY 1688 ) 1689 with self.transaction( 1690 condition=len(source_queries) > 0 or insert_overwrite_strategy.is_delete_insert 1691 ): 1692 target_columns_to_types = target_columns_to_types or self.columns(table_name) 1693 for i, source_query in enumerate(source_queries): 1694 with source_query as query: 1695 query = self._order_projections_and_filter( 1696 query, target_columns_to_types, where=where 1697 ) 1698 if i > 0 or insert_overwrite_strategy.is_delete_insert: 1699 if i == 0: 1700 self.delete_from(table_name, where=where or exp.true()) 1701 self._insert_append_query( 1702 table_name, 1703 query, 1704 target_columns_to_types=target_columns_to_types, 1705 order_projections=False, 1706 ) 1707 elif insert_overwrite_strategy.is_merge: 1708 columns = [exp.column(col) for col in target_columns_to_types] 1709 when_not_matched_by_source = exp.When( 1710 matched=False, 1711 source=True, 1712 condition=where, 1713 then=exp.Delete(), 1714 ) 1715 when_not_matched_by_target = exp.When( 1716 matched=False, 1717 source=False, 1718 then=exp.Insert( 1719 this=exp.Tuple(expressions=columns), 1720 expression=exp.Tuple(expressions=columns), 1721 ), 1722 ) 1723 self._merge( 1724 target_table=table_name, 1725 query=query, 1726 on=exp.false(), 1727 whens=exp.Whens( 1728 expressions=[when_not_matched_by_source, when_not_matched_by_target] 1729 ), 1730 ) 1731 else: 1732 insert_exp = exp.insert( 1733 query, 1734 table, 1735 columns=( 1736 list(target_columns_to_types) 1737 if not insert_overwrite_strategy.is_replace_where 1738 else None 1739 ), 1740 overwrite=insert_overwrite_strategy.is_insert_overwrite, 1741 ) 1742 if insert_overwrite_strategy.is_replace_where: 1743 insert_exp.set("where", where or exp.true()) 1744 self.execute(insert_exp, track_rows_processed=True) 1745 1746 def update_table( 1747 self, 1748 table_name: TableName, 1749 properties: t.Dict[str, t.Any], 1750 where: t.Optional[str | exp.Condition] = None, 1751 ) -> None: 1752 self.execute(exp.update(table_name, properties, where=where)) 1753 1754 def _merge( 1755 self, 1756 target_table: TableName, 1757 query: Query, 1758 on: exp.Expr, 1759 whens: exp.Whens, 1760 ) -> None: 1761 this = exp.alias_(exp.to_table(target_table), alias=MERGE_TARGET_ALIAS, table=True) 1762 using = exp.alias_( 1763 exp.Subquery(this=query), alias=MERGE_SOURCE_ALIAS, copy=False, table=True 1764 ) 1765 self.execute( 1766 exp.Merge(this=this, using=using, on=on, whens=whens), track_rows_processed=True 1767 ) 1768 1769 def scd_type_2_by_time( 1770 self, 1771 target_table: TableName, 1772 source_table: QueryOrDF, 1773 unique_key: t.Sequence[exp.Expr], 1774 valid_from_col: exp.Column, 1775 valid_to_col: exp.Column, 1776 execution_time: t.Union[TimeLike, exp.Column], 1777 updated_at_col: exp.Column, 1778 invalidate_hard_deletes: bool = True, 1779 updated_at_as_valid_from: bool = False, 1780 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1781 table_description: t.Optional[str] = None, 1782 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1783 truncate: bool = False, 1784 source_columns: t.Optional[t.List[str]] = None, 1785 **kwargs: t.Any, 1786 ) -> None: 1787 self._scd_type_2( 1788 target_table=target_table, 1789 source_table=source_table, 1790 unique_key=unique_key, 1791 valid_from_col=valid_from_col, 1792 valid_to_col=valid_to_col, 1793 execution_time=execution_time, 1794 updated_at_col=updated_at_col, 1795 invalidate_hard_deletes=invalidate_hard_deletes, 1796 updated_at_as_valid_from=updated_at_as_valid_from, 1797 target_columns_to_types=target_columns_to_types, 1798 table_description=table_description, 1799 column_descriptions=column_descriptions, 1800 truncate=truncate, 1801 source_columns=source_columns, 1802 **kwargs, 1803 ) 1804 1805 def scd_type_2_by_column( 1806 self, 1807 target_table: TableName, 1808 source_table: QueryOrDF, 1809 unique_key: t.Sequence[exp.Expr], 1810 valid_from_col: exp.Column, 1811 valid_to_col: exp.Column, 1812 execution_time: t.Union[TimeLike, exp.Column], 1813 check_columns: t.Union[exp.Star, t.Sequence[exp.Expr]], 1814 invalidate_hard_deletes: bool = True, 1815 execution_time_as_valid_from: bool = False, 1816 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1817 table_description: t.Optional[str] = None, 1818 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1819 truncate: bool = False, 1820 source_columns: t.Optional[t.List[str]] = None, 1821 **kwargs: t.Any, 1822 ) -> None: 1823 self._scd_type_2( 1824 target_table=target_table, 1825 source_table=source_table, 1826 unique_key=unique_key, 1827 valid_from_col=valid_from_col, 1828 valid_to_col=valid_to_col, 1829 execution_time=execution_time, 1830 check_columns=check_columns, 1831 target_columns_to_types=target_columns_to_types, 1832 invalidate_hard_deletes=invalidate_hard_deletes, 1833 execution_time_as_valid_from=execution_time_as_valid_from, 1834 table_description=table_description, 1835 column_descriptions=column_descriptions, 1836 truncate=truncate, 1837 source_columns=source_columns, 1838 **kwargs, 1839 ) 1840 1841 def _scd_type_2( 1842 self, 1843 target_table: TableName, 1844 source_table: QueryOrDF, 1845 unique_key: t.Sequence[exp.Expr], 1846 valid_from_col: exp.Column, 1847 valid_to_col: exp.Column, 1848 execution_time: t.Union[TimeLike, exp.Column], 1849 invalidate_hard_deletes: bool = True, 1850 updated_at_col: t.Optional[exp.Column] = None, 1851 check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Expr]]] = None, 1852 updated_at_as_valid_from: bool = False, 1853 execution_time_as_valid_from: bool = False, 1854 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1855 table_description: t.Optional[str] = None, 1856 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1857 truncate: bool = False, 1858 source_columns: t.Optional[t.List[str]] = None, 1859 **kwargs: t.Any, 1860 ) -> None: 1861 def remove_managed_columns( 1862 cols_to_types: t.Dict[str, exp.DataType], 1863 ) -> t.Dict[str, exp.DataType]: 1864 return { 1865 k: v for k, v in cols_to_types.items() if k not in {valid_from_name, valid_to_name} 1866 } 1867 1868 valid_from_name = valid_from_col.name 1869 valid_to_name = valid_to_col.name 1870 target_columns_to_types = target_columns_to_types or self.columns(target_table) 1871 if ( 1872 valid_from_name not in target_columns_to_types 1873 or valid_to_name not in target_columns_to_types 1874 or not columns_to_types_all_known(target_columns_to_types) 1875 ): 1876 target_columns_to_types = self.columns(target_table) 1877 unmanaged_columns_to_types = ( 1878 remove_managed_columns(target_columns_to_types) if target_columns_to_types else None 1879 ) 1880 source_queries, unmanaged_columns_to_types = self._get_source_queries_and_columns_to_types( 1881 source_table, 1882 unmanaged_columns_to_types, 1883 target_table=target_table, 1884 batch_size=0, 1885 source_columns=source_columns, 1886 ) 1887 updated_at_name = updated_at_col.name if updated_at_col else None 1888 if not target_columns_to_types: 1889 raise SQLMeshError(f"Could not get columns_to_types. Does {target_table} exist?") 1890 unmanaged_columns_to_types = unmanaged_columns_to_types or remove_managed_columns( 1891 target_columns_to_types 1892 ) 1893 if not unique_key: 1894 raise SQLMeshError("unique_key must be provided for SCD Type 2") 1895 if check_columns and updated_at_col: 1896 raise SQLMeshError( 1897 "Cannot use both `check_columns` and `updated_at_name` for SCD Type 2" 1898 ) 1899 if check_columns and updated_at_as_valid_from: 1900 raise SQLMeshError( 1901 "Cannot use both `check_columns` and `updated_at_as_valid_from` for SCD Type 2" 1902 ) 1903 if execution_time_as_valid_from and not check_columns: 1904 raise SQLMeshError( 1905 "Cannot use `execution_time_as_valid_from` without `check_columns` for SCD Type 2" 1906 ) 1907 if updated_at_name and updated_at_name not in target_columns_to_types: 1908 raise SQLMeshError( 1909 f"Column {updated_at_name} not found in {target_table}. Table must contain an `updated_at` timestamp for SCD Type 2" 1910 ) 1911 time_data_type = target_columns_to_types[valid_from_name] 1912 select_source_columns: t.List[t.Union[str, exp.Alias]] = [ 1913 col for col in unmanaged_columns_to_types if col != updated_at_name 1914 ] 1915 table_columns = [exp.column(c, quoted=True) for c in target_columns_to_types] 1916 if updated_at_name: 1917 select_source_columns.append( 1918 exp.cast(updated_at_col, time_data_type).as_(updated_at_col.this) # type: ignore 1919 ) 1920 1921 # If a star is provided, we include all unmanaged columns in the check. 1922 # This unnecessarily includes unique key columns but since they are used in the join, and therefore we know 1923 # they are equal or not, the extra check is not a problem and we gain simplified logic here. 1924 # If we want to change this, then we just need to check the expressions in unique_key and pull out the 1925 # column names and then remove them from the unmanaged_columns 1926 if check_columns: 1927 # Handle both Star directly and [Star()] (which can happen during serialization/deserialization) 1928 if isinstance(seq_get(ensure_list(check_columns), 0), exp.Star): 1929 check_columns = [exp.column(col) for col in unmanaged_columns_to_types] 1930 execution_ts = ( 1931 exp.cast(execution_time, time_data_type, dialect=self.dialect) 1932 if isinstance(execution_time, exp.Column) 1933 else to_time_column(execution_time, time_data_type, self.dialect, nullable=True) 1934 ) 1935 if updated_at_as_valid_from: 1936 if not updated_at_col: 1937 raise SQLMeshError( 1938 "Cannot use `updated_at_as_valid_from` without `updated_at_name` for SCD Type 2" 1939 ) 1940 update_valid_from_start: t.Union[str, exp.Expr] = updated_at_col 1941 # If using check_columns and the user doesn't always want execution_time for valid from 1942 # then we only use epoch 0 if we are truncating the table and loading rows for the first time. 1943 # All future new rows should have execution time. 1944 elif check_columns and (execution_time_as_valid_from or not truncate): 1945 update_valid_from_start = execution_ts 1946 else: 1947 update_valid_from_start = to_time_column( 1948 "1970-01-01 00:00:00+00:00", time_data_type, self.dialect, nullable=True 1949 ) 1950 insert_valid_from_start = execution_ts if check_columns else updated_at_col # type: ignore 1951 # joined._exists IS NULL is saying "if the row is deleted" 1952 delete_check = ( 1953 exp.column("_exists", "joined").is_(exp.Null()) if invalidate_hard_deletes else None 1954 ) 1955 prefixed_valid_to_col = valid_to_col.copy() 1956 prefixed_valid_to_col.this.set("this", f"t_{prefixed_valid_to_col.name}") 1957 prefixed_valid_from_col = valid_from_col.copy() 1958 prefixed_valid_from_col.this.set("this", f"t_{valid_from_col.name}") 1959 if check_columns: 1960 row_check_conditions = [] 1961 for col in check_columns: 1962 col_qualified = col.copy() 1963 col_qualified.set("table", exp.to_identifier("joined")) 1964 1965 t_col = col_qualified.copy() 1966 for column in t_col.find_all(exp.Column): 1967 column.this.set("this", f"t_{column.name}") 1968 1969 row_check_conditions.extend( 1970 [ 1971 col_qualified.neq(t_col), 1972 exp.and_(t_col.is_(exp.Null()), col_qualified.is_(exp.Null()).not_()), 1973 exp.and_(t_col.is_(exp.Null()).not_(), col_qualified.is_(exp.Null())), 1974 ] 1975 ) 1976 row_value_check = exp.or_(*row_check_conditions) 1977 unique_key_conditions = [] 1978 for key in unique_key: 1979 key_qualified = key.copy() 1980 key_qualified.set("table", exp.to_identifier("joined")) 1981 t_key = key_qualified.copy() 1982 for col in t_key.find_all(exp.Column): 1983 col.this.set("this", f"t_{col.name}") 1984 unique_key_conditions.extend( 1985 [t_key.is_(exp.Null()).not_(), key_qualified.is_(exp.Null()).not_()] 1986 ) 1987 unique_key_check = exp.and_(*unique_key_conditions) 1988 # unique_key_check is saying "if the row is updated" 1989 # row_value_check is saying "if the row has changed" 1990 updated_row_filter = exp.and_(unique_key_check, row_value_check) 1991 valid_to_case_stmt = ( 1992 exp.Case() 1993 .when( 1994 exp.and_( 1995 exp.or_( 1996 delete_check, 1997 updated_row_filter, 1998 ) 1999 ), 2000 execution_ts, 2001 ) 2002 .else_(prefixed_valid_to_col) 2003 .as_(valid_to_col.this) 2004 ) 2005 valid_from_case_stmt = exp.func( 2006 "COALESCE", 2007 prefixed_valid_from_col, 2008 update_valid_from_start, 2009 ).as_(valid_from_col.this) 2010 else: 2011 assert updated_at_col is not None 2012 updated_at_col_qualified = updated_at_col.copy() 2013 updated_at_col_qualified.set("table", exp.to_identifier("joined")) 2014 prefixed_updated_at_col = updated_at_col_qualified.copy() 2015 prefixed_updated_at_col.this.set("this", f"t_{updated_at_col_qualified.name}") 2016 updated_row_filter = updated_at_col_qualified > prefixed_updated_at_col 2017 2018 valid_to_case_stmt_builder = exp.Case().when( 2019 updated_row_filter, updated_at_col_qualified 2020 ) 2021 if delete_check: 2022 valid_to_case_stmt_builder = valid_to_case_stmt_builder.when( 2023 delete_check, execution_ts 2024 ) 2025 valid_to_case_stmt = valid_to_case_stmt_builder.else_(prefixed_valid_to_col).as_( 2026 valid_to_col.this 2027 ) 2028 2029 valid_from_case_stmt = ( 2030 exp.Case() 2031 .when( 2032 exp.and_( 2033 prefixed_valid_from_col.is_(exp.Null()), 2034 exp.column("_exists", "latest_deleted").is_(exp.Null()).not_(), 2035 ), 2036 exp.Case() 2037 .when( 2038 exp.column(valid_to_col.this, "latest_deleted") > updated_at_col, 2039 exp.column(valid_to_col.this, "latest_deleted"), 2040 ) 2041 .else_(updated_at_col), 2042 ) 2043 .when(prefixed_valid_from_col.is_(exp.Null()), update_valid_from_start) 2044 .else_(prefixed_valid_from_col) 2045 ).as_(valid_from_col.this) 2046 2047 existing_rows_query = exp.select(*table_columns, exp.true().as_("_exists")).from_( 2048 target_table 2049 ) 2050 if truncate: 2051 existing_rows_query = existing_rows_query.limit(0) 2052 2053 with source_queries[0] as source_query: 2054 prefixed_columns_to_types = [] 2055 for column in target_columns_to_types: 2056 prefixed_col = exp.column(column).copy() 2057 prefixed_col.this.set("this", f"t_{prefixed_col.name}") 2058 prefixed_columns_to_types.append(prefixed_col) 2059 prefixed_unmanaged_columns = [] 2060 for column in unmanaged_columns_to_types: 2061 prefixed_col = exp.column(column).copy() 2062 prefixed_col.this.set("this", f"t_{prefixed_col.name}") 2063 prefixed_unmanaged_columns.append(prefixed_col) 2064 query = ( 2065 exp.Select() # type: ignore 2066 .select(*table_columns) 2067 .from_("static") 2068 .union( 2069 exp.select(*table_columns).from_("updated_rows"), 2070 distinct=False, 2071 ) 2072 .union( 2073 exp.select(*table_columns).from_("inserted_rows"), 2074 distinct=False, 2075 ) 2076 .with_( 2077 "source", 2078 exp.select(exp.true().as_("_exists"), *select_source_columns) 2079 .distinct(*unique_key) 2080 .from_( 2081 self.use_server_nulls_for_unmatched_after_join(source_query).subquery( # type: ignore 2082 "raw_source" 2083 ) 2084 ), 2085 ) 2086 # Historical Records that Do Not Change 2087 .with_( 2088 "static", 2089 existing_rows_query.where(valid_to_col.is_(exp.Null()).not_()), 2090 ) 2091 # Latest Records that can be updated 2092 .with_( 2093 "latest", 2094 existing_rows_query.where(valid_to_col.is_(exp.Null())), 2095 ) 2096 # Deleted records which can be used to determine `valid_from` for undeleted source records 2097 .with_( 2098 "deleted", 2099 exp.select(*[exp.column(col, "static") for col in target_columns_to_types]) 2100 .from_("static") 2101 .join( 2102 "latest", 2103 on=exp.and_( 2104 *[ 2105 add_table(key, "static").eq(add_table(key, "latest")) 2106 for key in unique_key 2107 ] 2108 ), 2109 join_type="left", 2110 ) 2111 .where(exp.column(valid_to_col.this, "latest").is_(exp.Null())), 2112 ) 2113 # Get the latest `valid_to` deleted record for each unique key 2114 .with_( 2115 "latest_deleted", 2116 exp.select( 2117 exp.true().as_("_exists"), 2118 *(part.as_(f"_key{i}") for i, part in enumerate(unique_key)), 2119 exp.Max(this=valid_to_col).as_(valid_to_col.this), 2120 ) 2121 .from_("deleted") 2122 .group_by(*unique_key), 2123 ) 2124 # Do a full join between latest records and source table in order to combine them together 2125 # MySQL doesn't support full join so going to do a left then right join and remove dups with union 2126 # We do a left/right and filter right on only matching to remove the need to do union distinct 2127 # which allows scd type 2 to be compatible with unhashable data types 2128 .with_( 2129 "joined", 2130 exp.select( 2131 exp.column("_exists", table="source").as_("_exists"), 2132 *( 2133 exp.column(col, table="latest").as_(prefixed_columns_to_types[i].this) 2134 for i, col in enumerate(target_columns_to_types) 2135 ), 2136 *( 2137 exp.column(col, table="source").as_(col) 2138 for col in unmanaged_columns_to_types 2139 ), 2140 ) 2141 .from_("latest") 2142 .join( 2143 "source", 2144 on=exp.and_( 2145 *[ 2146 add_table(key, "latest").eq(add_table(key, "source")) 2147 for key in unique_key 2148 ] 2149 ), 2150 join_type="left", 2151 ) 2152 .union( 2153 exp.select( 2154 exp.column("_exists", table="source").as_("_exists"), 2155 *( 2156 exp.column(col, table="latest").as_( 2157 prefixed_columns_to_types[i].this 2158 ) 2159 for i, col in enumerate(target_columns_to_types) 2160 ), 2161 *( 2162 exp.column(col, table="source").as_(col) 2163 for col in unmanaged_columns_to_types 2164 ), 2165 ) 2166 .from_("latest") 2167 .join( 2168 "source", 2169 on=exp.and_( 2170 *[ 2171 add_table(key, "latest").eq(add_table(key, "source")) 2172 for key in unique_key 2173 ] 2174 ), 2175 join_type="right", 2176 ) 2177 .where(exp.column("_exists", table="latest").is_(exp.Null())), 2178 distinct=False, 2179 ), 2180 ) 2181 # Get deleted, new, no longer current, or unchanged records 2182 .with_( 2183 "updated_rows", 2184 exp.select( 2185 *( 2186 exp.func( 2187 "COALESCE", 2188 exp.column(prefixed_unmanaged_columns[i].this, table="joined"), 2189 exp.column(col, table="joined"), 2190 ).as_(col) 2191 for i, col in enumerate(unmanaged_columns_to_types) 2192 ), 2193 valid_from_case_stmt, 2194 valid_to_case_stmt, 2195 ) 2196 .from_("joined") 2197 .join( 2198 "latest_deleted", 2199 on=exp.and_( 2200 *[ 2201 add_table(part, "joined").eq( 2202 exp.column(f"_key{i}", "latest_deleted") 2203 ) 2204 for i, part in enumerate(unique_key) 2205 ] 2206 ), 2207 join_type="left", 2208 ), 2209 ) 2210 # Get records that have been "updated" which means inserting a new record with previous `valid_from` 2211 .with_( 2212 "inserted_rows", 2213 exp.select( 2214 *unmanaged_columns_to_types, 2215 insert_valid_from_start.as_(valid_from_col.this), # type: ignore 2216 to_time_column(exp.null(), time_data_type, self.dialect, nullable=True).as_( 2217 valid_to_col.this 2218 ), 2219 ) 2220 .from_("joined") 2221 .where(updated_row_filter), 2222 ) 2223 ) 2224 2225 self.replace_query( 2226 target_table, 2227 self.ensure_nulls_for_unmatched_after_join(query), 2228 target_columns_to_types=target_columns_to_types, 2229 table_description=table_description, 2230 column_descriptions=column_descriptions, 2231 **kwargs, 2232 ) 2233 2234 def merge( 2235 self, 2236 target_table: TableName, 2237 source_table: QueryOrDF, 2238 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 2239 unique_key: t.Sequence[exp.Expr], 2240 when_matched: t.Optional[exp.Whens] = None, 2241 merge_filter: t.Optional[exp.Expr] = None, 2242 source_columns: t.Optional[t.List[str]] = None, 2243 **kwargs: t.Any, 2244 ) -> None: 2245 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2246 source_table, 2247 target_columns_to_types, 2248 target_table=target_table, 2249 source_columns=source_columns, 2250 ) 2251 target_columns_to_types = target_columns_to_types or self.columns(target_table) 2252 on = exp.and_( 2253 *( 2254 add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS)) 2255 for part in unique_key 2256 ) 2257 ) 2258 if merge_filter: 2259 on = exp.and_(merge_filter, on) 2260 2261 if not when_matched: 2262 match_expressions = [ 2263 exp.When( 2264 matched=True, 2265 source=False, 2266 then=exp.Update( 2267 expressions=[ 2268 exp.column(col, MERGE_TARGET_ALIAS).eq( 2269 exp.column(col, MERGE_SOURCE_ALIAS) 2270 ) 2271 for col in target_columns_to_types 2272 ], 2273 ), 2274 ) 2275 ] 2276 else: 2277 match_expressions = when_matched.copy().expressions 2278 2279 match_expressions.append( 2280 exp.When( 2281 matched=False, 2282 source=False, 2283 then=exp.Insert( 2284 this=exp.Tuple( 2285 expressions=[exp.column(col) for col in target_columns_to_types] 2286 ), 2287 expression=exp.Tuple( 2288 expressions=[ 2289 exp.column(col, MERGE_SOURCE_ALIAS) for col in target_columns_to_types 2290 ] 2291 ), 2292 ), 2293 ) 2294 ) 2295 for source_query in source_queries: 2296 with source_query as query: 2297 self._merge( 2298 target_table=target_table, 2299 query=query, 2300 on=on, 2301 whens=exp.Whens(expressions=match_expressions), 2302 ) 2303 2304 def rename_table( 2305 self, 2306 old_table_name: TableName, 2307 new_table_name: TableName, 2308 ) -> None: 2309 new_table = exp.to_table(new_table_name) 2310 if new_table.catalog: 2311 old_table = exp.to_table(old_table_name) 2312 catalog = old_table.catalog or self.get_current_catalog() 2313 if catalog != new_table.catalog: 2314 raise UnsupportedCatalogOperationError( 2315 "Tried to rename table across catalogs which is not supported" 2316 ) 2317 self._rename_table(old_table_name, new_table_name) 2318 self._clear_data_object_cache(old_table_name) 2319 self._clear_data_object_cache(new_table_name) 2320 2321 def get_data_object( 2322 self, target_name: TableName, safe_to_cache: bool = False 2323 ) -> t.Optional[DataObject]: 2324 target_table = exp.to_table(target_name) 2325 existing_data_objects = self.get_data_objects( 2326 schema_(target_table.db, target_table.catalog), 2327 {target_table.name}, 2328 safe_to_cache=safe_to_cache, 2329 ) 2330 if existing_data_objects: 2331 return existing_data_objects[0] 2332 return None 2333 2334 def get_data_objects( 2335 self, 2336 schema_name: SchemaName, 2337 object_names: t.Optional[t.Set[str]] = None, 2338 safe_to_cache: bool = False, 2339 ) -> t.List[DataObject]: 2340 """Lists all data objects in the target schema. 2341 2342 Args: 2343 schema_name: The name of the schema to list data objects from. 2344 object_names: If provided, only return data objects with these names. 2345 safe_to_cache: Whether it is safe to cache the results of this call. 2346 2347 Returns: 2348 A list of data objects in the target schema. 2349 """ 2350 if object_names is not None: 2351 if not object_names: 2352 return [] 2353 2354 # Check cache for each object name 2355 target_schema = to_schema(schema_name) 2356 cached_objects = [] 2357 missing_names = set() 2358 2359 for name in object_names: 2360 cache_key = _get_data_object_cache_key( 2361 target_schema.catalog, target_schema.db, name 2362 ) 2363 if cache_key in self._data_object_cache: 2364 logger.debug("Data object cache hit: %s", cache_key) 2365 data_object = self._data_object_cache[cache_key] 2366 # If the object is none, then the table was previously looked for but not found 2367 if data_object: 2368 cached_objects.append(data_object) 2369 else: 2370 logger.debug("Data object cache miss: %s", cache_key) 2371 missing_names.add(name) 2372 2373 # Fetch missing objects from database 2374 if missing_names: 2375 object_names_list = list(missing_names) 2376 batches = [ 2377 object_names_list[i : i + self.DATA_OBJECT_FILTER_BATCH_SIZE] 2378 for i in range(0, len(object_names_list), self.DATA_OBJECT_FILTER_BATCH_SIZE) 2379 ] 2380 2381 fetched_objects = [] 2382 fetched_object_names = set() 2383 for batch in batches: 2384 objects = self._get_data_objects(schema_name, set(batch)) 2385 for obj in objects: 2386 if safe_to_cache: 2387 cache_key = _get_data_object_cache_key( 2388 obj.catalog, obj.schema_name, obj.name 2389 ) 2390 self._data_object_cache[cache_key] = obj 2391 fetched_objects.append(obj) 2392 fetched_object_names.add(obj.name) 2393 2394 if safe_to_cache: 2395 for missing_name in missing_names - fetched_object_names: 2396 cache_key = _get_data_object_cache_key( 2397 target_schema.catalog, target_schema.db, missing_name 2398 ) 2399 self._data_object_cache[cache_key] = None 2400 2401 return cached_objects + fetched_objects 2402 2403 return cached_objects 2404 2405 fetched_objects = self._get_data_objects(schema_name) 2406 if safe_to_cache: 2407 for obj in fetched_objects: 2408 cache_key = _get_data_object_cache_key(obj.catalog, obj.schema_name, obj.name) 2409 self._data_object_cache[cache_key] = obj 2410 return fetched_objects 2411 2412 def fetchone( 2413 self, 2414 query: t.Union[exp.Expr, str], 2415 ignore_unsupported_errors: bool = False, 2416 quote_identifiers: bool = False, 2417 ) -> t.Optional[t.Tuple]: 2418 with self.transaction(): 2419 self.execute( 2420 query, 2421 ignore_unsupported_errors=ignore_unsupported_errors, 2422 quote_identifiers=quote_identifiers, 2423 ) 2424 return self.cursor.fetchone() 2425 2426 def fetchall( 2427 self, 2428 query: t.Union[exp.Expr, str], 2429 ignore_unsupported_errors: bool = False, 2430 quote_identifiers: bool = False, 2431 ) -> t.List[t.Tuple]: 2432 with self.transaction(): 2433 self.execute( 2434 query, 2435 ignore_unsupported_errors=ignore_unsupported_errors, 2436 quote_identifiers=quote_identifiers, 2437 ) 2438 return self.cursor.fetchall() 2439 2440 def _fetch_native_df( 2441 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2442 ) -> DF: 2443 """Fetches a DataFrame that can be either Pandas or PySpark from the cursor""" 2444 with self.transaction(): 2445 self.execute(query, quote_identifiers=quote_identifiers) 2446 return self.cursor.fetchdf() 2447 2448 def _native_df_to_pandas_df( 2449 self, 2450 query_or_df: QueryOrDF, 2451 ) -> t.Union[Query, pd.DataFrame]: 2452 """ 2453 Take a "native" DataFrame (eg Pyspark, Bigframe, Snowpark etc) and convert it to Pandas 2454 """ 2455 import pandas as pd 2456 2457 if isinstance(query_or_df, (exp.Query, pd.DataFrame)): 2458 return query_or_df 2459 2460 # EngineAdapter subclasses that have native DataFrame types should override this 2461 raise NotImplementedError(f"Unable to convert {type(query_or_df)} to Pandas") 2462 2463 def fetchdf( 2464 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2465 ) -> pd.DataFrame: 2466 """Fetches a Pandas DataFrame from the cursor""" 2467 import pandas as pd 2468 2469 df = self._fetch_native_df(query, quote_identifiers=quote_identifiers) 2470 if not isinstance(df, pd.DataFrame): 2471 raise NotImplementedError( 2472 "The cursor's `fetch_native_df` method is not returning a pandas DataFrame. Need to update `fetchdf` so a Pandas DataFrame is returned" 2473 ) 2474 return df 2475 2476 def fetch_pyspark_df( 2477 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2478 ) -> PySparkDataFrame: 2479 """Fetches a PySpark DataFrame from the cursor""" 2480 raise NotImplementedError(f"Engine does not support PySpark DataFrames: {type(self)}") 2481 2482 @property 2483 def wap_enabled(self) -> bool: 2484 """Returns whether WAP is enabled for this engine.""" 2485 return self._extra_config.get("wap_enabled", False) 2486 2487 def wap_supported(self, table_name: TableName) -> bool: 2488 """Returns whether WAP for the target table is supported.""" 2489 return False 2490 2491 def wap_table_name(self, table_name: TableName, wap_id: str) -> str: 2492 """Returns the updated table name for the given WAP ID. 2493 2494 Args: 2495 table_name: The name of the target table. 2496 wap_id: The WAP ID to prepare. 2497 2498 Returns: 2499 The updated table name that should be used for writing. 2500 """ 2501 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2502 2503 def wap_prepare(self, table_name: TableName, wap_id: str) -> str: 2504 """Prepares the target table for WAP and returns the updated table name. 2505 2506 Args: 2507 table_name: The name of the target table. 2508 wap_id: The WAP ID to prepare. 2509 2510 Returns: 2511 The updated table name that should be used for writing. 2512 """ 2513 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2514 2515 def wap_publish(self, table_name: TableName, wap_id: str) -> None: 2516 """Publishes changes with the given WAP ID to the target table. 2517 2518 Args: 2519 table_name: The name of the target table. 2520 wap_id: The WAP ID to publish. 2521 """ 2522 raise NotImplementedError(f"Engine does not support WAP: {type(self)}") 2523 2524 def sync_grants_config( 2525 self, 2526 table: exp.Table, 2527 grants_config: GrantsConfig, 2528 table_type: DataObjectType = DataObjectType.TABLE, 2529 ) -> None: 2530 """Applies the grants_config to a table authoritatively. 2531 It first compares the specified grants against the current grants, and then 2532 applies the diffs to the table by revoking and granting privileges as needed. 2533 2534 Args: 2535 table: The table/view to apply grants to. 2536 grants_config: Dictionary mapping privileges to lists of grantees. 2537 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 2538 """ 2539 if not self.SUPPORTS_GRANTS: 2540 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 2541 2542 current_grants = self._get_current_grants_config(table) 2543 new_grants, revoked_grants = self._diff_grants_configs(grants_config, current_grants) 2544 revoke_exprs = self._revoke_grants_config_expr(table, revoked_grants, table_type) 2545 grant_exprs = self._apply_grants_config_expr(table, new_grants, table_type) 2546 dcl_exprs = revoke_exprs + grant_exprs 2547 2548 if dcl_exprs: 2549 self.execute(dcl_exprs) 2550 2551 @contextlib.contextmanager 2552 def transaction( 2553 self, 2554 condition: t.Optional[bool] = None, 2555 ) -> t.Iterator[None]: 2556 """A transaction context manager.""" 2557 if ( 2558 self._connection_pool.is_transaction_active 2559 or not self.SUPPORTS_TRANSACTIONS 2560 or (condition is not None and not condition) 2561 ): 2562 yield 2563 return 2564 2565 if self._pre_ping: 2566 try: 2567 logger.debug("Pinging the database to check the connection") 2568 self.ping() 2569 except Exception: 2570 logger.info("Connection to the database was lost. Reconnecting...") 2571 self._connection_pool.close() 2572 2573 self._connection_pool.begin() 2574 try: 2575 yield 2576 except Exception as e: 2577 self._connection_pool.rollback() 2578 raise e 2579 else: 2580 self._connection_pool.commit() 2581 2582 @contextlib.contextmanager 2583 def session(self, properties: SessionProperties) -> t.Iterator[None]: 2584 """A session context manager.""" 2585 if self._is_session_active(): 2586 yield 2587 return 2588 2589 self._begin_session(properties) 2590 try: 2591 yield 2592 finally: 2593 self._end_session() 2594 2595 def _begin_session(self, properties: SessionProperties) -> t.Any: 2596 """Begin a new session.""" 2597 2598 def _end_session(self) -> None: 2599 """End the existing session.""" 2600 2601 def _is_session_active(self) -> bool: 2602 """Indicates whether or not a session is active.""" 2603 return False 2604 2605 def execute( 2606 self, 2607 expressions: t.Union[str, exp.Expr, t.Sequence[exp.Expr]], 2608 ignore_unsupported_errors: bool = False, 2609 quote_identifiers: bool = True, 2610 track_rows_processed: bool = False, 2611 **kwargs: t.Any, 2612 ) -> None: 2613 """Execute a sql query.""" 2614 to_sql_kwargs = ( 2615 {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {} 2616 ) 2617 with self.transaction(): 2618 for e in ensure_list(expressions): 2619 if isinstance(e, exp.Expr): 2620 self._check_identifier_length(e) 2621 sql = self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs) 2622 else: 2623 sql = t.cast(str, e) 2624 2625 sql = self._attach_correlation_id(sql) 2626 2627 self._log_sql( 2628 sql, 2629 expression=e if isinstance(e, exp.Expr) else None, 2630 quote_identifiers=quote_identifiers, 2631 ) 2632 self._execute(sql, track_rows_processed, **kwargs) 2633 2634 def _attach_correlation_id(self, sql: str) -> str: 2635 if self.ATTACH_CORRELATION_ID and self.correlation_id: 2636 return f"/* {self.correlation_id} */ {sql}" 2637 return sql 2638 2639 def _log_sql( 2640 self, 2641 sql: str, 2642 expression: t.Optional[exp.Expr] = None, 2643 quote_identifiers: bool = True, 2644 ) -> None: 2645 if not logger.isEnabledFor(self._execute_log_level): 2646 return 2647 2648 sql_to_log = sql 2649 if expression is not None and not isinstance(expression, exp.Query): 2650 values = expression.find(exp.Values) 2651 if values: 2652 values.set("expressions", [exp.to_identifier("<REDACTED VALUES>")]) 2653 sql_to_log = self._to_sql(expression, quote=quote_identifiers) 2654 2655 logger.log(self._execute_log_level, "Executing SQL: %s", sql_to_log) 2656 2657 def _record_execution_stats( 2658 self, sql: str, rowcount: t.Optional[int] = None, bytes_processed: t.Optional[int] = None 2659 ) -> None: 2660 if self._query_execution_tracker: 2661 self._query_execution_tracker.record_execution(sql, rowcount, bytes_processed) 2662 2663 def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: 2664 self.cursor.execute(sql, **kwargs) 2665 2666 if ( 2667 self.SUPPORTS_QUERY_EXECUTION_TRACKING 2668 and track_rows_processed 2669 and self._query_execution_tracker 2670 and self._query_execution_tracker.is_tracking() 2671 ): 2672 if ( 2673 rowcount := getattr(self.cursor, "rowcount", None) 2674 ) is not None and rowcount is not None: 2675 try: 2676 self._record_execution_stats(sql, int(rowcount)) 2677 except (TypeError, ValueError): 2678 return 2679 2680 @contextlib.contextmanager 2681 def temp_table( 2682 self, 2683 query_or_df: QueryOrDF, 2684 name: TableName = "diff", 2685 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2686 source_columns: t.Optional[t.List[str]] = None, 2687 **kwargs: t.Any, 2688 ) -> t.Iterator[exp.Table]: 2689 """A context manager for working a temp table. 2690 2691 The table will be created with a random guid and cleaned up after the block. 2692 2693 Args: 2694 query_or_df: The query or df to create a temp table for. 2695 name: The base name of the temp table. 2696 target_columns_to_types: A mapping between the column name and its data type. 2697 2698 Yields: 2699 The table expression 2700 """ 2701 name = exp.to_table(name) 2702 # ensure that we use default catalog if none is not specified 2703 if isinstance(name, exp.Table) and not name.catalog and name.db and self.default_catalog: 2704 name.set("catalog", exp.parse_identifier(self.default_catalog)) 2705 2706 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2707 query_or_df, 2708 target_columns_to_types=target_columns_to_types, 2709 target_table=name, 2710 source_columns=source_columns, 2711 ) 2712 2713 with self.transaction(): 2714 table = self._get_temp_table(name) 2715 if table.db: 2716 self.create_schema(schema_(table.args["db"], table.args.get("catalog"))) 2717 self._create_table_from_source_queries( 2718 table, 2719 source_queries, 2720 target_columns_to_types, 2721 exists=True, 2722 table_description=None, 2723 column_descriptions=None, 2724 track_rows_processed=False, 2725 **kwargs, 2726 ) 2727 2728 try: 2729 yield table 2730 finally: 2731 self.drop_table(table) 2732 2733 def _table_or_view_properties_to_expressions( 2734 self, table_or_view_properties: t.Optional[t.Dict[str, exp.Expr]] = None 2735 ) -> t.List[exp.Property]: 2736 """Converts model properties (either physical or virtual) to a list of property expressions.""" 2737 if not table_or_view_properties: 2738 return [] 2739 return [ 2740 exp.Property(this=key, value=value.copy()) 2741 for key, value in table_or_view_properties.items() 2742 ] 2743 2744 def _build_partitioned_by_exp( 2745 self, 2746 partitioned_by: t.List[exp.Expr], 2747 *, 2748 partition_interval_unit: t.Optional[IntervalUnit] = None, 2749 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2750 catalog_name: t.Optional[str] = None, 2751 **kwargs: t.Any, 2752 ) -> t.Optional[t.Union[exp.PartitionedByProperty, exp.Property]]: 2753 return None 2754 2755 def _build_clustered_by_exp( 2756 self, 2757 clustered_by: t.List[exp.Expr], 2758 **kwargs: t.Any, 2759 ) -> t.Optional[exp.Cluster]: 2760 return None 2761 2762 def adjust_physical_properties_for_incremental( 2763 self, 2764 physical_properties: t.Dict[str, t.Any], 2765 *, 2766 requires_delete_capable_table: bool, 2767 unique_key: t.Optional[t.List[exp.Expr]], 2768 model_name: str, 2769 ) -> t.Dict[str, t.Any]: 2770 """Adjusts physical properties for an incremental model before the table is created. 2771 2772 Some engines require a specific physical table layout before they can run the DELETE/MERGE 2773 statements that incremental model kinds rely on (e.g. StarRocks only supports those on 2774 PRIMARY KEY tables). This hook lets each engine derive or validate the required properties 2775 while keeping the generic evaluator free of engine-specific branching. 2776 2777 Args: 2778 physical_properties: The model's physical properties. 2779 requires_delete_capable_table: Whether the model kind issues DELETE/MERGE statements 2780 (as opposed to append-only INSERTs), as determined by the generic evaluator. 2781 unique_key: The model's unique key, populated only when the kind allows promoting it to 2782 an engine-specific key (i.e. INCREMENTAL_BY_UNIQUE_KEY); otherwise None. 2783 model_name: The model name, for use in diagnostics. 2784 2785 Returns: 2786 The (possibly adjusted) physical properties. Implementations own the given mapping and 2787 may mutate it in place; the base implementation returns it unchanged. 2788 """ 2789 return physical_properties 2790 2791 def _build_table_properties_exp( 2792 self, 2793 catalog_name: t.Optional[str] = None, 2794 table_format: t.Optional[str] = None, 2795 storage_format: t.Optional[str] = None, 2796 partitioned_by: t.Optional[t.List[exp.Expr]] = None, 2797 partition_interval_unit: t.Optional[IntervalUnit] = None, 2798 clustered_by: t.Optional[t.List[exp.Expr]] = None, 2799 table_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 2800 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2801 table_description: t.Optional[str] = None, 2802 table_kind: t.Optional[str] = None, 2803 **kwargs: t.Any, 2804 ) -> t.Optional[exp.Properties]: 2805 """Creates a SQLGlot table properties expression for ddl.""" 2806 properties: t.List[exp.Expr] = [] 2807 2808 if table_description: 2809 properties.append( 2810 exp.SchemaCommentProperty( 2811 this=exp.Literal.string(self._truncate_table_comment(table_description)) 2812 ) 2813 ) 2814 2815 if table_properties: 2816 table_type = self._pop_creatable_type_from_properties(table_properties) 2817 properties.extend(ensure_list(table_type)) 2818 2819 if properties: 2820 return exp.Properties(expressions=properties) 2821 return None 2822 2823 def _build_view_properties_exp( 2824 self, 2825 view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 2826 table_description: t.Optional[str] = None, 2827 **kwargs: t.Any, 2828 ) -> t.Optional[exp.Properties]: 2829 """Creates a SQLGlot table properties expression for view""" 2830 properties: t.List[exp.Expr] = [] 2831 2832 if table_description: 2833 properties.append( 2834 exp.SchemaCommentProperty( 2835 this=exp.Literal.string(self._truncate_table_comment(table_description)) 2836 ) 2837 ) 2838 2839 if properties: 2840 return exp.Properties(expressions=properties) 2841 return None 2842 2843 def _truncate_comment(self, comment: str, length: t.Optional[int]) -> str: 2844 return comment[:length] if length else comment 2845 2846 def _truncate_table_comment(self, comment: str) -> str: 2847 return self._truncate_comment(comment, self.MAX_TABLE_COMMENT_LENGTH) 2848 2849 def _truncate_column_comment(self, comment: str) -> str: 2850 return self._truncate_comment(comment, self.MAX_COLUMN_COMMENT_LENGTH) 2851 2852 def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: 2853 """ 2854 Converts an expression to a SQL string. Has a set of default kwargs to apply, and then default 2855 kwargs defined for the given dialect, and then kwargs provided by the user when defining the engine 2856 adapter, and then finally kwargs provided by the user when calling this method. 2857 """ 2858 sql_gen_kwargs = { 2859 "dialect": self.dialect, 2860 "pretty": self._pretty_sql, 2861 "comments": False, 2862 **self._sql_gen_kwargs, 2863 **kwargs, 2864 } 2865 2866 expression = expression.copy() 2867 2868 if quote: 2869 quote_identifiers(expression) 2870 2871 return expression.sql(**sql_gen_kwargs, copy=False) # type: ignore 2872 2873 def _clear_data_object_cache(self, table_name: t.Optional[TableName] = None) -> None: 2874 """Clears the cache entry for the given table name, or clears the entire cache if table_name is None.""" 2875 if table_name is None: 2876 logger.debug("Clearing entire data object cache") 2877 self._data_object_cache.clear() 2878 else: 2879 table = exp.to_table(table_name) 2880 cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name) 2881 logger.debug("Clearing data object cache key: %s", cache_key) 2882 self._data_object_cache.pop(cache_key, None) 2883 2884 def _get_data_objects( 2885 self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None 2886 ) -> t.List[DataObject]: 2887 """ 2888 Returns all the data objects that exist in the given schema and optionally catalog. 2889 """ 2890 raise NotImplementedError() 2891 2892 def _get_temp_table( 2893 self, table: TableName, table_only: bool = False, quoted: bool = True 2894 ) -> exp.Table: 2895 """ 2896 Returns the name of the temp table that should be used for the given table name. 2897 """ 2898 table = t.cast(exp.Table, exp.to_table(table).copy()) 2899 table.set( 2900 "this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}", quoted=quoted) 2901 ) 2902 2903 if table_only: 2904 table.set("db", None) 2905 table.set("catalog", None) 2906 2907 return table 2908 2909 def _order_projections_and_filter( 2910 self, 2911 query: Query, 2912 target_columns_to_types: t.Dict[str, exp.DataType], 2913 where: t.Optional[exp.Expr] = None, 2914 coerce_types: bool = False, 2915 ) -> Query: 2916 if not isinstance(query, exp.Query) or ( 2917 not where and not coerce_types and query.named_selects == list(target_columns_to_types) 2918 ): 2919 return query 2920 2921 query = t.cast(exp.Query, query.copy()) 2922 with_ = query.args.pop("with_", None) 2923 2924 select_exprs: t.List[exp.Expr] = [ 2925 exp.column(c, quoted=True) for c in target_columns_to_types 2926 ] 2927 if coerce_types and columns_to_types_all_known(target_columns_to_types): 2928 select_exprs = [ 2929 exp.cast(select_exprs[i], col_tpe).as_(col, quoted=True) 2930 for i, (col, col_tpe) in enumerate(target_columns_to_types.items()) 2931 ] 2932 2933 query = exp.select(*select_exprs).from_(query.subquery("_subquery", copy=False), copy=False) 2934 if where: 2935 query = query.where(where, copy=False) 2936 2937 if with_: 2938 query.set("with_", with_) 2939 2940 return query 2941 2942 def _truncate_table(self, table_name: TableName) -> None: 2943 table = exp.to_table(table_name) 2944 self.execute(f"TRUNCATE TABLE {table.sql(dialect=self.dialect, identify=True)}") 2945 2946 def drop_data_object_on_type_mismatch( 2947 self, data_object: t.Optional[DataObject], expected_type: DataObjectType 2948 ) -> bool: 2949 """Drops a data object if it exists and is not of the expected type. 2950 2951 Args: 2952 data_object: The data object to check. 2953 expected_type: The expected type of the data object. 2954 2955 Returns: 2956 True if the data object was dropped, False otherwise. 2957 """ 2958 if data_object is None or data_object.type == expected_type: 2959 return False 2960 2961 logger.warning( 2962 "Target data object '%s' is a %s and not a %s, dropping it", 2963 data_object.to_table().sql(dialect=self.dialect), 2964 data_object.type.value, 2965 expected_type.value, 2966 ) 2967 self.drop_data_object(data_object) 2968 return True 2969 2970 def _replace_by_key( 2971 self, 2972 target_table: TableName, 2973 source_table: QueryOrDF, 2974 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 2975 key: t.Sequence[exp.Expr], 2976 is_unique_key: bool, 2977 source_columns: t.Optional[t.List[str]] = None, 2978 ) -> None: 2979 if target_columns_to_types is None: 2980 target_columns_to_types = self.columns(target_table) 2981 2982 temp_table = self._get_temp_table(target_table) 2983 key_exp = ( 2984 exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key, dialect=self.dialect) 2985 if len(key) > 1 2986 else key[0] 2987 ) 2988 column_names = list(target_columns_to_types or []) 2989 2990 with self.transaction(): 2991 self.ctas( 2992 temp_table, 2993 source_table, 2994 target_columns_to_types=target_columns_to_types, 2995 exists=False, 2996 source_columns=source_columns, 2997 ) 2998 2999 try: 3000 delete_query = exp.select(key_exp).from_(temp_table) 3001 insert_query = self._select_columns(target_columns_to_types).from_(temp_table) 3002 if not is_unique_key: 3003 delete_query = delete_query.distinct() 3004 else: 3005 insert_query = insert_query.distinct(*key) 3006 3007 insert_statement = exp.insert( 3008 insert_query, 3009 target_table, 3010 columns=column_names, 3011 ) 3012 delete_filter = key_exp.isin(query=delete_query) 3013 3014 if not self.INSERT_OVERWRITE_STRATEGY.is_replace_where: 3015 self.delete_from(target_table, delete_filter) 3016 else: 3017 insert_statement.set("where", delete_filter) 3018 insert_statement.set("this", exp.to_table(target_table)) 3019 3020 self.execute(insert_statement, track_rows_processed=True) 3021 finally: 3022 self.drop_table(temp_table) 3023 3024 def _build_create_comment_table_exp( 3025 self, table: exp.Table, table_comment: str, table_kind: str 3026 ) -> exp.Comment | str: 3027 return exp.Comment( 3028 this=table, 3029 kind=table_kind, 3030 expression=exp.Literal.string(self._truncate_table_comment(table_comment)), 3031 ) 3032 3033 def _create_table_comment( 3034 self, table_name: TableName, table_comment: str, table_kind: str = "TABLE" 3035 ) -> None: 3036 table = exp.to_table(table_name) 3037 3038 try: 3039 self.execute(self._build_create_comment_table_exp(table, table_comment, table_kind)) 3040 except Exception: 3041 logger.warning( 3042 f"Table comment for '{table.alias_or_name}' not registered - this may be due to limited permissions", 3043 exc_info=True, 3044 ) 3045 3046 def _build_create_comment_column_exp( 3047 self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE" 3048 ) -> exp.Comment | str: 3049 return exp.Comment( 3050 this=exp.column(column_name, *reversed(table.parts)), # type: ignore 3051 kind="COLUMN", 3052 expression=exp.Literal.string(self._truncate_column_comment(column_comment)), 3053 ) 3054 3055 def _create_column_comments( 3056 self, 3057 table_name: TableName, 3058 column_comments: t.Dict[str, str], 3059 table_kind: str = "TABLE", 3060 materialized_view: bool = False, 3061 ) -> None: 3062 table = exp.to_table(table_name) 3063 3064 for col, comment in column_comments.items(): 3065 try: 3066 self.execute(self._build_create_comment_column_exp(table, col, comment, table_kind)) 3067 except Exception: 3068 logger.warning( 3069 f"Column comments for column '{col}' in table '{table.alias_or_name}' not registered - this may be due to limited permissions", 3070 exc_info=True, 3071 ) 3072 3073 def _create_table_like( 3074 self, 3075 target_table_name: TableName, 3076 source_table_name: TableName, 3077 exists: bool, 3078 **kwargs: t.Any, 3079 ) -> None: 3080 self.create_table(target_table_name, self.columns(source_table_name), exists=exists) 3081 3082 def _rename_table( 3083 self, 3084 old_table_name: TableName, 3085 new_table_name: TableName, 3086 ) -> None: 3087 self.execute(exp.rename_table(old_table_name, new_table_name)) 3088 3089 def ensure_nulls_for_unmatched_after_join( 3090 self, 3091 query: Query, 3092 ) -> Query: 3093 return query 3094 3095 def use_server_nulls_for_unmatched_after_join( 3096 self, 3097 query: Query, 3098 ) -> Query: 3099 return query 3100 3101 def ping(self) -> None: 3102 try: 3103 self._execute(exp.select("1").sql(dialect=self.dialect)) 3104 finally: 3105 self._connection_pool.close_cursor() 3106 3107 @classmethod 3108 def _select_columns( 3109 cls, columns: t.Iterable[str], source_columns: t.Optional[t.List[str]] = None 3110 ) -> exp.Select: 3111 return exp.select( 3112 *( 3113 exp.column(c, quoted=True) 3114 if c in (source_columns or columns) 3115 else exp.alias_(exp.Null(), c, quoted=True) 3116 for c in columns 3117 ) 3118 ) 3119 3120 def _check_identifier_length(self, expression: exp.Expr) -> None: 3121 if self.MAX_IDENTIFIER_LENGTH is None or not isinstance(expression, exp.DDL): 3122 return 3123 3124 for identifier in expression.find_all(exp.Identifier): 3125 name = identifier.name 3126 name_length = len(name) 3127 if name_length > self.MAX_IDENTIFIER_LENGTH: 3128 raise SQLMeshError( 3129 f"Identifier name '{name}' (length {name_length}) exceeds {self.dialect.capitalize()}'s max identifier limit of {self.MAX_IDENTIFIER_LENGTH} characters" 3130 ) 3131 3132 def get_table_last_modified_ts(self, table_names: t.List[TableName]) -> t.List[int]: 3133 raise NotImplementedError() 3134 3135 @classmethod 3136 def _diff_grants_configs( 3137 cls, new_config: GrantsConfig, old_config: GrantsConfig 3138 ) -> t.Tuple[GrantsConfig, GrantsConfig]: 3139 """Compute additions and removals between two grants configurations. 3140 3141 This method compares new (desired) and old (current) GrantsConfigs case-insensitively 3142 for both privilege keys and grantees, while preserving original casing 3143 in the output GrantsConfigs. 3144 3145 Args: 3146 new_config: Desired grants configuration (specified by the user). 3147 old_config: Current grants configuration (returned by the database). 3148 3149 Returns: 3150 A tuple of (additions, removals) GrantsConfig where: 3151 - additions contains privileges/grantees present in new_config but not in old_config 3152 - additions uses keys and grantee strings from new_config (user-specified casing) 3153 - removals contains privileges/grantees present in old_config but not in new_config 3154 - removals uses keys and grantee strings from old_config (database-returned casing) 3155 3156 Notes: 3157 - Comparison is case-insensitive using casefold(); original casing is preserved in results. 3158 - Overlapping grantees (case-insensitive) are excluded from the results. 3159 """ 3160 3161 def _diffs(config1: GrantsConfig, config2: GrantsConfig) -> GrantsConfig: 3162 diffs: GrantsConfig = {} 3163 cf_config2 = {k.casefold(): {g.casefold() for g in v} for k, v in config2.items()} 3164 for key, grantees in config1.items(): 3165 cf_key = key.casefold() 3166 3167 # Missing key (add all grantees) 3168 if cf_key not in cf_config2: 3169 diffs[key] = grantees.copy() 3170 continue 3171 3172 # Include only grantees not in config2 3173 cf_grantees2 = cf_config2[cf_key] 3174 diff_grantees = [] 3175 for grantee in grantees: 3176 if grantee.casefold() not in cf_grantees2: 3177 diff_grantees.append(grantee) 3178 if diff_grantees: 3179 diffs[key] = diff_grantees 3180 return diffs 3181 3182 return _diffs(new_config, old_config), _diffs(old_config, new_config) 3183 3184 def _get_current_grants_config(self, table: exp.Table) -> GrantsConfig: 3185 """Returns current grants for a table as a dictionary. 3186 3187 This method queries the database and returns the current grants/permissions 3188 for the given table, parsed into a dictionary format. The it handles 3189 case-insensitive comparison between these current grants and the desired 3190 grants from model configuration. 3191 3192 Args: 3193 table: The table/view to query grants for. 3194 3195 Returns: 3196 Dictionary mapping permissions to lists of grantees. Permission names 3197 should be returned as the database provides them (typically uppercase 3198 for standard SQL permissions, but engine-specific roles may vary). 3199 3200 Raises: 3201 NotImplementedError: If the engine does not support grants. 3202 """ 3203 if not self.SUPPORTS_GRANTS: 3204 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3205 raise NotImplementedError("Subclass must implement get_current_grants") 3206 3207 def _apply_grants_config_expr( 3208 self, 3209 table: exp.Table, 3210 grants_config: GrantsConfig, 3211 table_type: DataObjectType = DataObjectType.TABLE, 3212 ) -> t.List[exp.Expr]: 3213 """Returns SQLGlot Grant expressions to apply grants to a table. 3214 3215 Args: 3216 table: The table/view to grant permissions on. 3217 grants_config: Dictionary mapping permissions to lists of grantees. 3218 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 3219 3220 Returns: 3221 List of SQLGlot expressions for grant operations. 3222 3223 Raises: 3224 NotImplementedError: If the engine does not support grants. 3225 """ 3226 if not self.SUPPORTS_GRANTS: 3227 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3228 raise NotImplementedError("Subclass must implement _apply_grants_config_expr") 3229 3230 def _revoke_grants_config_expr( 3231 self, 3232 table: exp.Table, 3233 grants_config: GrantsConfig, 3234 table_type: DataObjectType = DataObjectType.TABLE, 3235 ) -> t.List[exp.Expr]: 3236 """Returns SQLGlot expressions to revoke grants from a table. 3237 3238 Args: 3239 table: The table/view to revoke permissions from. 3240 grants_config: Dictionary mapping permissions to lists of grantees. 3241 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 3242 3243 Returns: 3244 List of SQLGlot expressions for revoke operations. 3245 3246 Raises: 3247 NotImplementedError: If the engine does not support grants. 3248 """ 3249 if not self.SUPPORTS_GRANTS: 3250 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 3251 raise NotImplementedError("Subclass must implement _revoke_grants_config_expr")
Base class wrapping a Database API compliant connection.
The EngineAdapter is an easily-subclassable interface that interacts with the underlying engine and data store.
Arguments:
- connection_factory_or_pool: a callable which produces a new Database API-compliant connection on every call.
- dialect: The dialect with which this adapter is associated.
- multithreaded: Indicates whether this adapter will be used by more than one thread.
EngineAdapter( connection_factory_or_pool: Union[Callable[[], Any], sqlmesh.utils.connection_pool.ConnectionPool], dialect: str = '', sql_gen_kwargs: Optional[Dict[str, sqlglot.dialects.dialect.Dialect | bool | str]] = None, multithreaded: bool = False, cursor_init: Optional[Callable[[Any], NoneType]] = None, default_catalog: Optional[str] = None, execute_log_level: int = 10, register_comments: bool = True, pre_ping: bool = False, pretty_sql: bool = False, shared_connection: bool = False, correlation_id: Optional[sqlmesh.utils.CorrelationId] = None, schema_differ_overrides: Optional[Dict[str, Any]] = None, query_execution_tracker: Optional[sqlmesh.core.snapshot.execution_tracker.QueryExecutionTracker] = None, **kwargs: Any)
132 def __init__( 133 self, 134 connection_factory_or_pool: t.Union[t.Callable[[], t.Any], ConnectionPool], 135 dialect: str = "", 136 sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None, 137 multithreaded: bool = False, 138 cursor_init: t.Optional[t.Callable[[t.Any], None]] = None, 139 default_catalog: t.Optional[str] = None, 140 execute_log_level: int = logging.DEBUG, 141 register_comments: bool = True, 142 pre_ping: bool = False, 143 pretty_sql: bool = False, 144 shared_connection: bool = False, 145 correlation_id: t.Optional[CorrelationId] = None, 146 schema_differ_overrides: t.Optional[t.Dict[str, t.Any]] = None, 147 query_execution_tracker: t.Optional[QueryExecutionTracker] = None, 148 **kwargs: t.Any, 149 ): 150 self.dialect = dialect.lower() or self.DIALECT 151 self._connection_pool = ( 152 connection_factory_or_pool 153 if isinstance(connection_factory_or_pool, ConnectionPool) 154 else create_connection_pool( 155 connection_factory_or_pool, 156 multithreaded, 157 shared_connection=shared_connection, 158 cursor_init=cursor_init, 159 ) 160 ) 161 self._sql_gen_kwargs = sql_gen_kwargs or {} 162 self._default_catalog = default_catalog 163 self._execute_log_level = execute_log_level 164 self._extra_config = kwargs 165 self._register_comments = register_comments 166 self._pre_ping = pre_ping 167 self._pretty_sql = pretty_sql 168 self._multithreaded = multithreaded 169 self.correlation_id = correlation_id 170 self._schema_differ_overrides = schema_differ_overrides 171 self._query_execution_tracker = query_execution_tracker 172 self._data_object_cache: t.Dict[str, t.Optional[DataObject]] = {}
RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES: FrozenSet[str] =
frozenset()
Physical property keys whose values may contain logical model references that should be resolved to physical table names during property rendering. Engines that need such resolution (e.g. StarRocks' excluded_trigger_tables) override this set.
174 def with_settings(self, **kwargs: t.Any) -> EngineAdapter: 175 extra_kwargs = { 176 "null_connection": True, 177 "execute_log_level": kwargs.pop("execute_log_level", self._execute_log_level), 178 "correlation_id": kwargs.pop("correlation_id", self.correlation_id), 179 "query_execution_tracker": kwargs.pop( 180 "query_execution_tracker", self._query_execution_tracker 181 ), 182 **self._extra_config, 183 **kwargs, 184 } 185 186 adapter = self.__class__( 187 self._connection_pool, 188 dialect=self.dialect, 189 sql_gen_kwargs=self._sql_gen_kwargs, 190 default_catalog=self._default_catalog, 191 register_comments=self._register_comments, 192 multithreaded=self._multithreaded, 193 pretty_sql=self._pretty_sql, 194 **extra_kwargs, 195 ) 196 197 return adapter
catalog_support: sqlmesh.core.engine_adapter.shared.CatalogSupport
def
supports_virtual_catalog(self) -> bool:
227 def supports_virtual_catalog(self) -> bool: 228 """Return True if this adapter can accept a virtual catalog for multi-gateway nesting alignment. 229 230 When a project mixes catalog-aware gateways (e.g. DuckDB) with catalog-unsupported gateways 231 (e.g. ClickHouse), all adapters need a uniform 3-level FQN so MappingSchema nesting stays 232 consistent. Adapters that return True here opt in to receiving an injected virtual catalog 233 via inject_virtual_catalog(), which causes the set_catalog decorator to strip the catalog 234 from DDL expressions rather than raising UnsupportedCatalogOperationError. 235 """ 236 return False
Return True if this adapter can accept a virtual catalog for multi-gateway nesting alignment.
When a project mixes catalog-aware gateways (e.g. DuckDB) with catalog-unsupported gateways (e.g. ClickHouse), all adapters need a uniform 3-level FQN so MappingSchema nesting stays consistent. Adapters that return True here opt in to receiving an injected virtual catalog via inject_virtual_catalog(), which causes the set_catalog decorator to strip the catalog from DDL expressions rather than raising UnsupportedCatalogOperationError.
def
inject_virtual_catalog(self, gateway: str) -> None:
238 def inject_virtual_catalog(self, gateway: str) -> None: 239 """Inject a gateway name to configure the adapter's virtual catalog. 240 241 The adapter determines the final catalog name from the gateway name (e.g. ClickHouse 242 wraps it as __{gateway}__). Only call this on adapters that return True from 243 supports_virtual_catalog(). After injection, catalog_support should return 244 SINGLE_CATALOG_ONLY so the set_catalog decorator strips the virtual catalog from DDL 245 expressions instead of raising an error. 246 """ 247 raise NotImplementedError( 248 f"{self.dialect} does not support virtual catalog injection. " 249 "Override supports_virtual_catalog() to return True and implement inject_virtual_catalog()." 250 )
Inject a gateway name to configure the adapter's virtual catalog.
The adapter determines the final catalog name from the gateway name (e.g. ClickHouse wraps it as __{gateway}__). Only call this on adapters that return True from supports_virtual_catalog(). After injection, catalog_support should return SINGLE_CATALOG_ONLY so the set_catalog decorator strips the virtual catalog from DDL expressions instead of raising an error.
schema_differ: sqlmesh.core.schema_diff.SchemaDiffer
default_catalog: Optional[str]
287 @property 288 def default_catalog(self) -> t.Optional[str]: 289 if self.catalog_support.is_unsupported: 290 return None 291 default_catalog = self._default_catalog or self.get_current_catalog() 292 if not default_catalog: 293 raise MissingDefaultCatalogError( 294 "Could not determine a default catalog despite it being supported." 295 ) 296 return default_catalog
engine_run_mode: sqlmesh.core.engine_adapter.shared.EngineRunMode
def
recycle(self) -> None:
447 def recycle(self) -> None: 448 """Closes all open connections and releases all allocated resources associated with any thread 449 except the calling one.""" 450 self._connection_pool.close_all(exclude_calling_thread=True)
Closes all open connections and releases all allocated resources associated with any thread except the calling one.
def
close(self) -> Any:
452 def close(self) -> t.Any: 453 """Closes all open connections and releases all allocated resources.""" 454 self._connection_pool.close_all()
Closes all open connections and releases all allocated resources.
def
get_current_catalog(self) -> Optional[str]:
456 def get_current_catalog(self) -> t.Optional[str]: 457 """Returns the catalog name of the current connection.""" 458 raise NotImplementedError()
Returns the catalog name of the current connection.
def
set_current_catalog(self, catalog: str) -> None:
460 def set_current_catalog(self, catalog: str) -> None: 461 """Sets the catalog name of the current connection.""" 462 raise NotImplementedError()
Sets the catalog name of the current connection.
def
get_catalog_type(self, catalog: Optional[str]) -> str:
464 def get_catalog_type(self, catalog: t.Optional[str]) -> str: 465 """Intended to be overridden for data virtualization systems like Trino that, 466 depending on the target catalog, require slightly different properties to be set when creating / updating tables 467 """ 468 if self.catalog_support.is_unsupported: 469 raise UnsupportedCatalogOperationError( 470 f"{self.dialect} does not support catalogs and a catalog was provided: {catalog}" 471 ) 472 return ( 473 self._catalog_type_overrides.get(catalog, self.DEFAULT_CATALOG_TYPE) 474 if catalog 475 else self.DEFAULT_CATALOG_TYPE 476 )
Intended to be overridden for data virtualization systems like Trino that, depending on the target catalog, require slightly different properties to be set when creating / updating tables
def
get_catalog_type_from_table(self, table: Union[str, sqlglot.expressions.query.Table]) -> str:
478 def get_catalog_type_from_table(self, table: TableName) -> str: 479 """Get the catalog type from a table name if it has a catalog specified, otherwise return the current catalog type""" 480 catalog = exp.to_table(table).catalog or self.get_current_catalog() 481 return self.get_catalog_type(catalog)
Get the catalog type from a table name if it has a catalog specified, otherwise return the current catalog type
current_catalog_type: str
483 @property 484 def current_catalog_type(self) -> str: 485 # `get_catalog_type_from_table` should be used over this property. Reason is that the table that is the target 486 # of the operation is what matters and not the catalog type of the connection. 487 # This still remains for legacy reasons and should be refactored out. 488 return self.get_catalog_type(self.get_current_catalog())
def
replace_query( self, table_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, source_columns: Optional[List[str]] = None, supports_replace_table_override: Optional[bool] = None, **kwargs: Any) -> None:
490 def replace_query( 491 self, 492 table_name: TableName, 493 query_or_df: QueryOrDF, 494 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 495 table_description: t.Optional[str] = None, 496 column_descriptions: t.Optional[t.Dict[str, str]] = None, 497 source_columns: t.Optional[t.List[str]] = None, 498 supports_replace_table_override: t.Optional[bool] = None, 499 **kwargs: t.Any, 500 ) -> None: 501 """Replaces an existing table with a query. 502 503 For partition based engines (hive, spark), insert override is used. For other systems, create or replace is used. 504 505 Args: 506 table_name: The name of the table (eg. prod.table) 507 query_or_df: The SQL query to run or a dataframe. 508 target_columns_to_types: Only used if a dataframe is provided. A mapping between the column name and its data type. 509 Expected to be ordered to match the order of values in the dataframe. 510 kwargs: Optional create table properties. 511 """ 512 target_table = exp.to_table(table_name) 513 514 target_data_object = self.get_data_object(target_table) 515 table_exists = target_data_object is not None 516 if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE): 517 table_exists = False 518 519 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 520 query_or_df, 521 target_columns_to_types, 522 target_table=target_table, 523 source_columns=source_columns, 524 ) 525 if not target_columns_to_types and table_exists: 526 target_columns_to_types = self.columns(target_table) 527 query = source_queries[0].query_factory() 528 self_referencing = any( 529 quote_identifiers(table) == quote_identifiers(target_table) 530 for table in query.find_all(exp.Table) 531 ) 532 # If a query references itself then it must have a table created regardless of approach used. 533 if self_referencing: 534 if not target_columns_to_types: 535 raise SQLMeshError( 536 f"Cannot create a self-referencing table {target_table.sql(dialect=self.dialect)} without knowing the column types. " 537 "Try casting the columns to an expected type or defining the columns in the model metadata. " 538 ) 539 self._create_table_from_columns( 540 target_table, 541 target_columns_to_types, 542 exists=True, 543 table_description=table_description, 544 column_descriptions=column_descriptions, 545 **kwargs, 546 ) 547 # All engines support `CREATE TABLE AS` so we use that if the table doesn't already exist and we 548 # use `CREATE OR REPLACE TABLE AS` if the engine supports it 549 supports_replace_table = ( 550 self.SUPPORTS_REPLACE_TABLE 551 if supports_replace_table_override is None 552 else supports_replace_table_override 553 ) 554 if supports_replace_table or not table_exists: 555 return self._create_table_from_source_queries( 556 target_table, 557 source_queries, 558 target_columns_to_types, 559 replace=supports_replace_table, 560 table_description=table_description, 561 column_descriptions=column_descriptions, 562 **kwargs, 563 ) 564 if self_referencing: 565 assert target_columns_to_types is not None 566 with self.temp_table( 567 self._select_columns(target_columns_to_types).from_(target_table), 568 name=target_table, 569 target_columns_to_types=target_columns_to_types, 570 **kwargs, 571 ) as temp_table: 572 for source_query in source_queries: 573 source_query.add_transform( 574 lambda node: ( # type: ignore 575 temp_table # type: ignore 576 if isinstance(node, exp.Table) 577 and quote_identifiers(node) == quote_identifiers(target_table) 578 else node 579 ) 580 ) 581 return self._insert_overwrite_by_condition( 582 target_table, 583 source_queries, 584 target_columns_to_types, 585 **kwargs, 586 ) 587 return self._insert_overwrite_by_condition( 588 target_table, 589 source_queries, 590 target_columns_to_types, 591 **kwargs, 592 )
Replaces an existing table with a query.
For partition based engines (hive, spark), insert override is used. For other systems, create or replace is used.
Arguments:
- table_name: The name of the table (eg. prod.table)
- query_or_df: The SQL query to run or a dataframe.
- target_columns_to_types: Only used if a dataframe is provided. A mapping between the column name and its data type. Expected to be ordered to match the order of values in the dataframe.
- kwargs: Optional create table properties.
def
create_index( self, table_name: Union[str, sqlglot.expressions.query.Table], index_name: str, columns: Tuple[str, ...], exists: bool = True) -> None:
594 def create_index( 595 self, 596 table_name: TableName, 597 index_name: str, 598 columns: t.Tuple[str, ...], 599 exists: bool = True, 600 ) -> None: 601 """Creates a new index for the given table if supported 602 603 Args: 604 table_name: The name of the target table. 605 index_name: The name of the index. 606 columns: The list of columns that constitute the index. 607 exists: Indicates whether to include the IF NOT EXISTS check. 608 """ 609 if not self.SUPPORTS_INDEXES: 610 return 611 612 expression = exp.Create( 613 this=exp.Index( 614 this=exp.to_identifier(index_name), 615 table=exp.to_table(table_name), 616 params=exp.IndexParameters(columns=[exp.to_column(c) for c in columns]), 617 ), 618 kind="INDEX", 619 exists=exists, 620 ) 621 self.execute(expression)
Creates a new index for the given table if supported
Arguments:
- table_name: The name of the target table.
- index_name: The name of the index.
- columns: The list of columns that constitute the index.
- exists: Indicates whether to include the IF NOT EXISTS check.
def
create_table( self, table_name: Union[str, sqlglot.expressions.query.Table], target_columns_to_types: Dict[str, sqlglot.expressions.datatypes.DataType], primary_key: Optional[Tuple[str, ...]] = None, exists: bool = True, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, **kwargs: Any) -> None:
650 def create_table( 651 self, 652 table_name: TableName, 653 target_columns_to_types: t.Dict[str, exp.DataType], 654 primary_key: t.Optional[t.Tuple[str, ...]] = None, 655 exists: bool = True, 656 table_description: t.Optional[str] = None, 657 column_descriptions: t.Optional[t.Dict[str, str]] = None, 658 **kwargs: t.Any, 659 ) -> None: 660 """Create a table using a DDL statement 661 662 Args: 663 table_name: The name of the table to create. Can be fully qualified or just table name. 664 target_columns_to_types: A mapping between the column name and its data type. 665 primary_key: Determines the table primary key. 666 exists: Indicates whether to include the IF NOT EXISTS check. 667 table_description: Optional table description from MODEL DDL. 668 column_descriptions: Optional column descriptions from model query. 669 kwargs: Optional create table properties. 670 """ 671 self._create_table_from_columns( 672 table_name, 673 target_columns_to_types, 674 primary_key, 675 exists, 676 table_description, 677 column_descriptions, 678 **kwargs, 679 )
Create a table using a DDL statement
Arguments:
- table_name: The name of the table to create. Can be fully qualified or just table name.
- target_columns_to_types: A mapping between the column name and its data type.
- primary_key: Determines the table primary key.
- exists: Indicates whether to include the IF NOT EXISTS check.
- table_description: Optional table description from MODEL DDL.
- column_descriptions: Optional column descriptions from model query.
- kwargs: Optional create table properties.
def
create_managed_table( self, table_name: Union[str, sqlglot.expressions.query.Table], query: <MagicMock id='130969805323856'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, partitioned_by: Optional[List[sqlglot.expressions.core.Expr]] = None, clustered_by: Optional[List[sqlglot.expressions.core.Expr]] = None, table_properties: Optional[Dict[str, sqlglot.expressions.core.Expr]] = None, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
681 def create_managed_table( 682 self, 683 table_name: TableName, 684 query: Query, 685 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 686 partitioned_by: t.Optional[t.List[exp.Expr]] = None, 687 clustered_by: t.Optional[t.List[exp.Expr]] = None, 688 table_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 689 table_description: t.Optional[str] = None, 690 column_descriptions: t.Optional[t.Dict[str, str]] = None, 691 source_columns: t.Optional[t.List[str]] = None, 692 **kwargs: t.Any, 693 ) -> None: 694 """Create a managed table using a query. 695 696 "Managed" means that once the table is created, the data is kept up to date by the underlying database engine and not SQLMesh. 697 698 Args: 699 table_name: The name of the table to create. Can be fully qualified or just table name. 700 query: The SQL query for the engine to base the managed table on 701 target_columns_to_types: A mapping between the column name and its data type. 702 partitioned_by: The partition columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 703 clustered_by: The cluster columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour)) 704 table_properties: Optional mapping of engine-specific properties to be set on the managed table 705 table_description: Optional table description from MODEL DDL. 706 column_descriptions: Optional column descriptions from model query. 707 kwargs: Optional create table properties. 708 """ 709 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}")
Create a managed table using a query.
"Managed" means that once the table is created, the data is kept up to date by the underlying database engine and not SQLMesh.
Arguments:
- table_name: The name of the table to create. Can be fully qualified or just table name.
- query: The SQL query for the engine to base the managed table on
- target_columns_to_types: A mapping between the column name and its data type.
- partitioned_by: The partition columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour))
- clustered_by: The cluster columns or engine specific expressions, only applicable in certain engines. (eg. (ds, hour))
- table_properties: Optional mapping of engine-specific properties to be set on the managed table
- table_description: Optional table description from MODEL DDL.
- column_descriptions: Optional column descriptions from model query.
- kwargs: Optional create table properties.
def
ctas( self, table_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, exists: bool = True, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
711 def ctas( 712 self, 713 table_name: TableName, 714 query_or_df: QueryOrDF, 715 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 716 exists: bool = True, 717 table_description: t.Optional[str] = None, 718 column_descriptions: t.Optional[t.Dict[str, str]] = None, 719 source_columns: t.Optional[t.List[str]] = None, 720 **kwargs: t.Any, 721 ) -> None: 722 """Create a table using a CTAS statement 723 724 Args: 725 table_name: The name of the table to create. Can be fully qualified or just table name. 726 query_or_df: The SQL query to run or a dataframe for the CTAS. 727 target_columns_to_types: A mapping between the column name and its data type. Required if using a DataFrame. 728 exists: Indicates whether to include the IF NOT EXISTS check. 729 table_description: Optional table description from MODEL DDL. 730 column_descriptions: Optional column descriptions from model query. 731 kwargs: Optional create table properties. 732 """ 733 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 734 query_or_df, 735 target_columns_to_types, 736 target_table=table_name, 737 source_columns=source_columns, 738 ) 739 return self._create_table_from_source_queries( 740 table_name, 741 source_queries, 742 target_columns_to_types, 743 exists, 744 table_description=table_description, 745 column_descriptions=column_descriptions, 746 **kwargs, 747 )
Create a table using a CTAS statement
Arguments:
- table_name: The name of the table to create. Can be fully qualified or just table name.
- query_or_df: The SQL query to run or a dataframe for the CTAS.
- target_columns_to_types: A mapping between the column name and its data type. Required if using a DataFrame.
- exists: Indicates whether to include the IF NOT EXISTS check.
- table_description: Optional table description from MODEL DDL.
- column_descriptions: Optional column descriptions from model query.
- kwargs: Optional create table properties.
def
create_state_table( self, table_name: str, target_columns_to_types: Dict[str, sqlglot.expressions.datatypes.DataType], primary_key: Optional[Tuple[str, ...]] = None) -> None:
749 def create_state_table( 750 self, 751 table_name: str, 752 target_columns_to_types: t.Dict[str, exp.DataType], 753 primary_key: t.Optional[t.Tuple[str, ...]] = None, 754 ) -> None: 755 """Create a table to store SQLMesh internal state. 756 757 Args: 758 table_name: The name of the table to create. Can be fully qualified or just table name. 759 target_columns_to_types: A mapping between the column name and its data type. 760 primary_key: Determines the table primary key. 761 """ 762 self.create_table( 763 table_name, 764 target_columns_to_types, 765 primary_key=primary_key, 766 )
Create a table to store SQLMesh internal state.
Arguments:
- table_name: The name of the table to create. Can be fully qualified or just table name.
- target_columns_to_types: A mapping between the column name and its data type.
- primary_key: Determines the table primary key.
def
create_table_like( self, target_table_name: Union[str, sqlglot.expressions.query.Table], source_table_name: Union[str, sqlglot.expressions.query.Table], exists: bool = True, **kwargs: Any) -> None:
1073 def create_table_like( 1074 self, 1075 target_table_name: TableName, 1076 source_table_name: TableName, 1077 exists: bool = True, 1078 **kwargs: t.Any, 1079 ) -> None: 1080 """Create a table to store SQLMesh internal state based on the definition of another table, including any 1081 column attributes and indexes defined in the original table. 1082 1083 Args: 1084 target_table_name: The name of the table to create. Can be fully qualified or just table name. 1085 source_table_name: The name of the table to base the new table on. 1086 """ 1087 self._create_table_like(target_table_name, source_table_name, exists=exists, **kwargs) 1088 self._clear_data_object_cache(target_table_name)
Create a table to store SQLMesh internal state based on the definition of another table, including any column attributes and indexes defined in the original table.
Arguments:
- target_table_name: The name of the table to create. Can be fully qualified or just table name.
- source_table_name: The name of the table to base the new table on.
def
clone_table( self, target_table_name: Union[str, sqlglot.expressions.query.Table], source_table_name: Union[str, sqlglot.expressions.query.Table], replace: bool = False, exists: bool = True, clone_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None:
1090 def clone_table( 1091 self, 1092 target_table_name: TableName, 1093 source_table_name: TableName, 1094 replace: bool = False, 1095 exists: bool = True, 1096 clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, 1097 **kwargs: t.Any, 1098 ) -> None: 1099 """Creates a table with the target name by cloning the source table. 1100 1101 Args: 1102 target_table_name: The name of the table that should be created. 1103 source_table_name: The name of the source table that should be cloned. 1104 replace: Whether or not to replace an existing table. 1105 exists: Indicates whether to include the IF NOT EXISTS check. 1106 """ 1107 if not self.SUPPORTS_CLONING: 1108 raise NotImplementedError(f"Engine does not support cloning: {type(self)}") 1109 1110 kwargs.pop("rendered_physical_properties", None) 1111 self.execute( 1112 exp.Create( 1113 this=exp.to_table(target_table_name), 1114 kind="TABLE", 1115 replace=replace, 1116 exists=exists, 1117 clone=exp.Clone( 1118 this=exp.to_table(source_table_name), 1119 **(clone_kwargs or {}), 1120 ), 1121 **kwargs, 1122 ) 1123 ) 1124 self._clear_data_object_cache(target_table_name)
Creates a table with the target name by cloning the source table.
Arguments:
- target_table_name: The name of the table that should be created.
- source_table_name: The name of the source table that should be cloned.
- replace: Whether or not to replace an existing table.
- exists: Indicates whether to include the IF NOT EXISTS check.
def
drop_data_object( self, data_object: sqlmesh.core.engine_adapter.shared.DataObject, ignore_if_not_exists: bool = True) -> None:
1126 def drop_data_object(self, data_object: DataObject, ignore_if_not_exists: bool = True) -> None: 1127 """Drops a data object of arbitrary type. 1128 1129 Args: 1130 data_object: The data object to drop. 1131 ignore_if_not_exists: If True, no error will be raised if the data object does not exist. 1132 """ 1133 if data_object.type.is_view: 1134 self.drop_view(data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists) 1135 elif data_object.type.is_materialized_view: 1136 self.drop_view( 1137 data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists, materialized=True 1138 ) 1139 elif data_object.type.is_table: 1140 self.drop_table(data_object.to_table(), exists=ignore_if_not_exists) 1141 elif data_object.type.is_managed_table: 1142 self.drop_managed_table(data_object.to_table(), exists=ignore_if_not_exists) 1143 else: 1144 raise SQLMeshError( 1145 f"Can't drop data object '{data_object.to_table().sql(dialect=self.dialect)}' of type '{data_object.type.value}'" 1146 )
Drops a data object of arbitrary type.
Arguments:
- data_object: The data object to drop.
- ignore_if_not_exists: If True, no error will be raised if the data object does not exist.
def
drop_table( self, table_name: Union[str, sqlglot.expressions.query.Table], exists: bool = True, **kwargs: Any) -> None:
1148 def drop_table(self, table_name: TableName, exists: bool = True, **kwargs: t.Any) -> None: 1149 """Drops a table. 1150 1151 Args: 1152 table_name: The name of the table to drop. 1153 exists: If exists, defaults to True. 1154 """ 1155 self._drop_object(name=table_name, exists=exists, **kwargs)
Drops a table.
Arguments:
- table_name: The name of the table to drop.
- exists: If exists, defaults to True.
def
drop_managed_table( self, table_name: Union[str, sqlglot.expressions.query.Table], exists: bool = True) -> None:
1157 def drop_managed_table(self, table_name: TableName, exists: bool = True) -> None: 1158 """Drops a managed table. 1159 1160 Args: 1161 table_name: The name of the table to drop. 1162 exists: If exists, defaults to True. 1163 """ 1164 raise NotImplementedError(f"Engine does not support managed tables: {type(self)}")
Drops a managed table.
Arguments:
- table_name: The name of the table to drop.
- exists: If exists, defaults to True.
def
get_alter_operations( self, current_table_name: Union[str, sqlglot.expressions.query.Table], target_table_name: Union[str, sqlglot.expressions.query.Table], *, ignore_destructive: bool = False, ignore_additive: bool = False) -> List[sqlmesh.core.schema_diff.TableAlterOperation]:
1192 def get_alter_operations( 1193 self, 1194 current_table_name: TableName, 1195 target_table_name: TableName, 1196 *, 1197 ignore_destructive: bool = False, 1198 ignore_additive: bool = False, 1199 ) -> t.List[TableAlterOperation]: 1200 """ 1201 Determines the alter statements needed to change the current table into the structure of the target table. 1202 """ 1203 return t.cast( 1204 t.List[TableAlterOperation], 1205 self.schema_differ.compare_columns( 1206 current_table_name, 1207 self.columns(current_table_name), 1208 self.columns(target_table_name), 1209 ignore_destructive=ignore_destructive, 1210 ignore_additive=ignore_additive, 1211 ), 1212 )
Determines the alter statements needed to change the current table into the structure of the target table.
def
alter_table( self, alter_expressions: Union[List[sqlglot.expressions.ddl.Alter], List[sqlmesh.core.schema_diff.TableAlterOperation]]) -> None:
1214 def alter_table( 1215 self, 1216 alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], 1217 ) -> None: 1218 """ 1219 Performs the alter statements to change the current table into the structure of the target table. 1220 """ 1221 with self.transaction(): 1222 for alter_expression in [ 1223 x.expression if isinstance(x, TableAlterOperation) else x for x in alter_expressions 1224 ]: 1225 self.execute(alter_expression)
Performs the alter statements to change the current table into the structure of the target table.
def
create_view( self, view_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, replace: bool = True, materialized: bool = False, materialized_properties: Optional[Dict[str, Any]] = None, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, view_properties: Optional[Dict[str, sqlglot.expressions.core.Expr]] = None, source_columns: Optional[List[str]] = None, **create_kwargs: Any) -> None:
1227 def create_view( 1228 self, 1229 view_name: TableName, 1230 query_or_df: QueryOrDF, 1231 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1232 replace: bool = True, 1233 materialized: bool = False, 1234 materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, 1235 table_description: t.Optional[str] = None, 1236 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1237 view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, 1238 source_columns: t.Optional[t.List[str]] = None, 1239 **create_kwargs: t.Any, 1240 ) -> None: 1241 """Create a view with a query or dataframe. 1242 1243 If a dataframe is passed in, it will be converted into a literal values statement. 1244 This should only be done if the dataframe is very small! 1245 1246 Args: 1247 view_name: The view name. 1248 query_or_df: A query or dataframe. 1249 target_columns_to_types: Columns to use in the view statement. 1250 replace: Whether or not to replace an existing view defaults to True. 1251 materialized: Whether to create a a materialized view. Only used for engines that support this feature. 1252 materialized_properties: Optional materialized view properties to add to the view. 1253 table_description: Optional table description from MODEL DDL. 1254 column_descriptions: Optional column descriptions from model query. 1255 view_properties: Optional view properties to add to the view. 1256 create_kwargs: Additional kwargs to pass into the Create expression 1257 """ 1258 import pandas as pd 1259 1260 if materialized_properties and not materialized: 1261 raise SQLMeshError("Materialized properties are only supported for materialized views") 1262 1263 query_or_df = self._native_df_to_pandas_df(query_or_df) 1264 1265 if isinstance(query_or_df, pd.DataFrame): 1266 values: t.List[t.Tuple[t.Any, ...]] = list( 1267 query_or_df.itertuples(index=False, name=None) 1268 ) 1269 target_columns_to_types, source_columns = self._columns_to_types( 1270 query_or_df, target_columns_to_types, source_columns 1271 ) 1272 if not target_columns_to_types: 1273 raise SQLMeshError("columns_to_types must be provided for dataframes") 1274 source_columns_to_types = get_source_columns_to_types( 1275 target_columns_to_types, source_columns 1276 ) 1277 query_or_df = self._values_to_sql( 1278 values, 1279 source_columns_to_types, 1280 batch_start=0, 1281 batch_end=len(values), 1282 ) 1283 1284 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1285 query_or_df, 1286 target_columns_to_types, 1287 batch_size=0, 1288 target_table=view_name, 1289 source_columns=source_columns, 1290 ) 1291 if len(source_queries) != 1: 1292 raise SQLMeshError("Only one source query is supported for creating views") 1293 1294 schema: t.Union[exp.Table, exp.Schema] = exp.to_table(view_name) 1295 if target_columns_to_types: 1296 schema = self._build_schema_exp( 1297 exp.to_table(view_name), 1298 target_columns_to_types, 1299 column_descriptions, 1300 is_view=True, 1301 materialized=materialized, 1302 ) 1303 1304 properties = create_kwargs.pop("properties", None) 1305 if not properties: 1306 properties = exp.Properties(expressions=[]) 1307 1308 if view_properties: 1309 table_type = self._pop_creatable_type_from_properties(view_properties) 1310 if table_type: 1311 properties.append("expressions", table_type) 1312 1313 if materialized and self.SUPPORTS_MATERIALIZED_VIEWS: 1314 properties.append("expressions", exp.MaterializedProperty()) 1315 1316 if not self.SUPPORTS_MATERIALIZED_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1317 schema = schema.this 1318 1319 if not self.SUPPORTS_VIEW_SCHEMA and isinstance(schema, exp.Schema): 1320 schema = schema.this 1321 1322 if materialized_properties: 1323 partitioned_by = materialized_properties.pop("partitioned_by", None) 1324 clustered_by = materialized_properties.pop("clustered_by", None) 1325 if ( 1326 partitioned_by 1327 and ( 1328 partitioned_by_prop := self._build_partitioned_by_exp( 1329 partitioned_by, **materialized_properties 1330 ) 1331 ) 1332 is not None 1333 ): 1334 materialized_properties["catalog_name"] = exp.to_table(view_name).catalog 1335 properties.append("expressions", partitioned_by_prop) 1336 if ( 1337 clustered_by 1338 and ( 1339 clustered_by_prop := self._build_clustered_by_exp( 1340 clustered_by, **materialized_properties 1341 ) 1342 ) 1343 is not None 1344 ): 1345 properties.append("expressions", clustered_by_prop) 1346 1347 create_view_properties = self._build_view_properties_exp( 1348 view_properties, 1349 ( 1350 table_description 1351 if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled 1352 else None 1353 ), 1354 physical_cluster=create_kwargs.pop("physical_cluster", None), 1355 ) 1356 if create_view_properties: 1357 for view_property in create_view_properties.expressions: 1358 # Small hack to make sure SECURE goes at the beginning before materialized as required by Snowflake 1359 if isinstance(view_property, exp.SecureProperty): 1360 properties.set("expressions", view_property, index=0, overwrite=False) 1361 else: 1362 properties.append("expressions", view_property) 1363 1364 if properties.expressions: 1365 create_kwargs["properties"] = properties 1366 1367 if replace: 1368 self.drop_data_object_on_type_mismatch( 1369 self.get_data_object(view_name), 1370 DataObjectType.VIEW if not materialized else DataObjectType.MATERIALIZED_VIEW, 1371 ) 1372 1373 with source_queries[0] as query: 1374 self.execute( 1375 exp.Create( 1376 this=schema, 1377 kind="VIEW", 1378 replace=replace, 1379 expression=query, 1380 **create_kwargs, 1381 ), 1382 quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, 1383 ) 1384 1385 self._clear_data_object_cache(view_name) 1386 1387 # Register table comment with commands if the engine doesn't support doing it in CREATE 1388 if ( 1389 table_description 1390 and self.COMMENT_CREATION_VIEW.is_comment_command_only 1391 and self.comments_enabled 1392 ): 1393 self._create_table_comment(view_name, table_description, "VIEW") 1394 # Register column comments with commands if the engine doesn't support doing it in 1395 # CREATE or we couldn't do it in the CREATE schema definition because we don't have 1396 # columns_to_types 1397 if ( 1398 column_descriptions 1399 and ( 1400 self.COMMENT_CREATION_VIEW.is_comment_command_only 1401 or ( 1402 self.COMMENT_CREATION_VIEW.is_in_schema_def_and_commands 1403 and not target_columns_to_types 1404 ) 1405 ) 1406 and self.comments_enabled 1407 ): 1408 self._create_column_comments(view_name, column_descriptions, "VIEW", materialized)
Create a view with a query or dataframe.
If a dataframe is passed in, it will be converted into a literal values statement. This should only be done if the dataframe is very small!
Arguments:
- view_name: The view name.
- query_or_df: A query or dataframe.
- target_columns_to_types: Columns to use in the view statement.
- replace: Whether or not to replace an existing view defaults to True.
- materialized: Whether to create a a materialized view. Only used for engines that support this feature.
- materialized_properties: Optional materialized view properties to add to the view.
- table_description: Optional table description from MODEL DDL.
- column_descriptions: Optional column descriptions from model query.
- view_properties: Optional view properties to add to the view.
- create_kwargs: Additional kwargs to pass into the Create expression
@set_catalog()
def
create_schema( self, schema_name: Union[str, sqlglot.expressions.query.Table], ignore_if_exists: bool = True, warn_on_error: bool = True, properties: Optional[List[sqlglot.expressions.core.Expr]] = None) -> None:
1410 @set_catalog() 1411 def create_schema( 1412 self, 1413 schema_name: SchemaName, 1414 ignore_if_exists: bool = True, 1415 warn_on_error: bool = True, 1416 properties: t.Optional[t.List[exp.Expr]] = None, 1417 ) -> None: 1418 properties = properties or [] 1419 return self._create_schema( 1420 schema_name=schema_name, 1421 ignore_if_exists=ignore_if_exists, 1422 warn_on_error=warn_on_error, 1423 properties=properties, 1424 kind="SCHEMA", 1425 )
def
drop_schema( self, schema_name: Union[str, sqlglot.expressions.query.Table], ignore_if_not_exists: bool = True, cascade: bool = False, **drop_args: Dict[str, sqlglot.expressions.core.Expr]) -> None:
1452 def drop_schema( 1453 self, 1454 schema_name: SchemaName, 1455 ignore_if_not_exists: bool = True, 1456 cascade: bool = False, 1457 **drop_args: t.Dict[str, exp.Expr], 1458 ) -> None: 1459 return self._drop_object( 1460 name=schema_name, 1461 exists=ignore_if_not_exists, 1462 kind="SCHEMA", 1463 cascade=cascade, 1464 **drop_args, 1465 )
def
drop_view( self, view_name: Union[str, sqlglot.expressions.query.Table], ignore_if_not_exists: bool = True, materialized: bool = False, **kwargs: Any) -> None:
1467 def drop_view( 1468 self, 1469 view_name: TableName, 1470 ignore_if_not_exists: bool = True, 1471 materialized: bool = False, 1472 **kwargs: t.Any, 1473 ) -> None: 1474 """Drop a view.""" 1475 self._drop_object( 1476 name=view_name, 1477 exists=ignore_if_not_exists, 1478 kind="VIEW", 1479 materialized=materialized and self.SUPPORTS_MATERIALIZED_VIEWS, 1480 **kwargs, 1481 )
Drop a view.
def
columns( self, table_name: Union[str, sqlglot.expressions.query.Table], include_pseudo_columns: bool = False) -> Dict[str, sqlglot.expressions.datatypes.DataType]:
1499 def columns( 1500 self, table_name: TableName, include_pseudo_columns: bool = False 1501 ) -> t.Dict[str, exp.DataType]: 1502 """Fetches column names and types for the target table.""" 1503 self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) 1504 describe_output = self.cursor.fetchall() 1505 return { 1506 # Note: MySQL returns the column type as bytes. 1507 column_name: exp.DataType.build(_decoded_str(column_type), dialect=self.dialect) 1508 for column_name, column_type, *_ in itertools.takewhile( 1509 lambda t: not t[0].startswith("#"), 1510 describe_output, 1511 ) 1512 if column_name and column_name.strip() and column_type and column_type.strip() 1513 }
Fetches column names and types for the target table.
def
table_exists(self, table_name: Union[str, sqlglot.expressions.query.Table]) -> bool:
1515 def table_exists(self, table_name: TableName) -> bool: 1516 table = exp.to_table(table_name) 1517 data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name) 1518 if data_object_cache_key in self._data_object_cache: 1519 logger.debug("Table existence cache hit: %s", data_object_cache_key) 1520 return self._data_object_cache[data_object_cache_key] is not None 1521 1522 try: 1523 self.execute(exp.Describe(this=table, kind="TABLE")) 1524 return True 1525 except Exception: 1526 return False
def
delete_from( self, table_name: Union[str, sqlglot.expressions.query.Table], where: Union[str, sqlglot.expressions.core.Expr]) -> None:
def
insert_append( self, table_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, track_rows_processed: bool = True, source_columns: Optional[List[str]] = None) -> None:
1531 def insert_append( 1532 self, 1533 table_name: TableName, 1534 query_or_df: QueryOrDF, 1535 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1536 track_rows_processed: bool = True, 1537 source_columns: t.Optional[t.List[str]] = None, 1538 ) -> None: 1539 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1540 query_or_df, 1541 target_columns_to_types, 1542 target_table=table_name, 1543 source_columns=source_columns, 1544 ) 1545 self._insert_append_source_queries( 1546 table_name, source_queries, target_columns_to_types, track_rows_processed 1547 )
def
insert_overwrite_by_partition( self, table_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, partitioned_by: List[sqlglot.expressions.core.Expr], target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, source_columns: Optional[List[str]] = None) -> None:
1582 def insert_overwrite_by_partition( 1583 self, 1584 table_name: TableName, 1585 query_or_df: QueryOrDF, 1586 partitioned_by: t.List[exp.Expr], 1587 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1588 source_columns: t.Optional[t.List[str]] = None, 1589 ) -> None: 1590 if self.INSERT_OVERWRITE_STRATEGY.is_insert_overwrite: 1591 target_table = exp.to_table(table_name) 1592 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1593 query_or_df, 1594 target_columns_to_types, 1595 target_table=target_table, 1596 source_columns=source_columns, 1597 ) 1598 self._insert_overwrite_by_condition( 1599 table_name, source_queries, target_columns_to_types=target_columns_to_types 1600 ) 1601 else: 1602 self._replace_by_key( 1603 table_name, 1604 query_or_df, 1605 target_columns_to_types, 1606 partitioned_by, 1607 is_unique_key=False, 1608 source_columns=source_columns, 1609 )
def
insert_overwrite_by_time_partition( self, table_name: Union[str, sqlglot.expressions.query.Table], query_or_df: <MagicMock id='130969804860832'>, start: Union[datetime.date, datetime.datetime, str, int, float], end: Union[datetime.date, datetime.datetime, str, int, float], time_formatter: Callable[[Union[datetime.date, datetime.datetime, str, int, float], Optional[Dict[str, sqlglot.expressions.datatypes.DataType]]], sqlglot.expressions.core.Expr], time_column: sqlmesh.core.model.kind.TimeColumn | sqlglot.expressions.core.Expr | str, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
1611 def insert_overwrite_by_time_partition( 1612 self, 1613 table_name: TableName, 1614 query_or_df: QueryOrDF, 1615 start: TimeLike, 1616 end: TimeLike, 1617 time_formatter: t.Callable[[TimeLike, t.Optional[t.Dict[str, exp.DataType]]], exp.Expr], 1618 time_column: TimeColumn | exp.Expr | str, 1619 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1620 source_columns: t.Optional[t.List[str]] = None, 1621 **kwargs: t.Any, 1622 ) -> None: 1623 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 1624 query_or_df, 1625 target_columns_to_types, 1626 target_table=table_name, 1627 source_columns=source_columns, 1628 ) 1629 if not target_columns_to_types or not columns_to_types_all_known(target_columns_to_types): 1630 target_columns_to_types = self.columns(table_name) 1631 low, high = [ 1632 time_formatter(dt, target_columns_to_types) 1633 for dt in make_inclusive(start, end, self.dialect) 1634 ] 1635 if isinstance(time_column, TimeColumn): 1636 time_column = time_column.column 1637 where = exp.Between( 1638 this=exp.to_column(time_column) if isinstance(time_column, str) else time_column, 1639 low=low, 1640 high=high, 1641 ) 1642 return self._insert_overwrite_by_time_partition( 1643 table_name, source_queries, target_columns_to_types, where, **kwargs 1644 )
def
update_table( self, table_name: Union[str, sqlglot.expressions.query.Table], properties: Dict[str, Any], where: Union[str, sqlglot.expressions.core.Condition, NoneType] = None) -> None:
def
scd_type_2_by_time( self, target_table: Union[str, sqlglot.expressions.query.Table], source_table: <MagicMock id='130969804860832'>, unique_key: Sequence[sqlglot.expressions.core.Expr], valid_from_col: sqlglot.expressions.core.Column, valid_to_col: sqlglot.expressions.core.Column, execution_time: Union[datetime.date, datetime.datetime, str, int, float, sqlglot.expressions.core.Column], updated_at_col: sqlglot.expressions.core.Column, invalidate_hard_deletes: bool = True, updated_at_as_valid_from: bool = False, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, truncate: bool = False, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
1769 def scd_type_2_by_time( 1770 self, 1771 target_table: TableName, 1772 source_table: QueryOrDF, 1773 unique_key: t.Sequence[exp.Expr], 1774 valid_from_col: exp.Column, 1775 valid_to_col: exp.Column, 1776 execution_time: t.Union[TimeLike, exp.Column], 1777 updated_at_col: exp.Column, 1778 invalidate_hard_deletes: bool = True, 1779 updated_at_as_valid_from: bool = False, 1780 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1781 table_description: t.Optional[str] = None, 1782 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1783 truncate: bool = False, 1784 source_columns: t.Optional[t.List[str]] = None, 1785 **kwargs: t.Any, 1786 ) -> None: 1787 self._scd_type_2( 1788 target_table=target_table, 1789 source_table=source_table, 1790 unique_key=unique_key, 1791 valid_from_col=valid_from_col, 1792 valid_to_col=valid_to_col, 1793 execution_time=execution_time, 1794 updated_at_col=updated_at_col, 1795 invalidate_hard_deletes=invalidate_hard_deletes, 1796 updated_at_as_valid_from=updated_at_as_valid_from, 1797 target_columns_to_types=target_columns_to_types, 1798 table_description=table_description, 1799 column_descriptions=column_descriptions, 1800 truncate=truncate, 1801 source_columns=source_columns, 1802 **kwargs, 1803 )
def
scd_type_2_by_column( self, target_table: Union[str, sqlglot.expressions.query.Table], source_table: <MagicMock id='130969804860832'>, unique_key: Sequence[sqlglot.expressions.core.Expr], valid_from_col: sqlglot.expressions.core.Column, valid_to_col: sqlglot.expressions.core.Column, execution_time: Union[datetime.date, datetime.datetime, str, int, float, sqlglot.expressions.core.Column], check_columns: Union[sqlglot.expressions.core.Star, Sequence[sqlglot.expressions.core.Expr]], invalidate_hard_deletes: bool = True, execution_time_as_valid_from: bool = False, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, table_description: Optional[str] = None, column_descriptions: Optional[Dict[str, str]] = None, truncate: bool = False, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
1805 def scd_type_2_by_column( 1806 self, 1807 target_table: TableName, 1808 source_table: QueryOrDF, 1809 unique_key: t.Sequence[exp.Expr], 1810 valid_from_col: exp.Column, 1811 valid_to_col: exp.Column, 1812 execution_time: t.Union[TimeLike, exp.Column], 1813 check_columns: t.Union[exp.Star, t.Sequence[exp.Expr]], 1814 invalidate_hard_deletes: bool = True, 1815 execution_time_as_valid_from: bool = False, 1816 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 1817 table_description: t.Optional[str] = None, 1818 column_descriptions: t.Optional[t.Dict[str, str]] = None, 1819 truncate: bool = False, 1820 source_columns: t.Optional[t.List[str]] = None, 1821 **kwargs: t.Any, 1822 ) -> None: 1823 self._scd_type_2( 1824 target_table=target_table, 1825 source_table=source_table, 1826 unique_key=unique_key, 1827 valid_from_col=valid_from_col, 1828 valid_to_col=valid_to_col, 1829 execution_time=execution_time, 1830 check_columns=check_columns, 1831 target_columns_to_types=target_columns_to_types, 1832 invalidate_hard_deletes=invalidate_hard_deletes, 1833 execution_time_as_valid_from=execution_time_as_valid_from, 1834 table_description=table_description, 1835 column_descriptions=column_descriptions, 1836 truncate=truncate, 1837 source_columns=source_columns, 1838 **kwargs, 1839 )
def
merge( self, target_table: Union[str, sqlglot.expressions.query.Table], source_table: <MagicMock id='130969804860832'>, target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]], unique_key: Sequence[sqlglot.expressions.core.Expr], when_matched: Optional[sqlglot.expressions.dml.Whens] = None, merge_filter: Optional[sqlglot.expressions.core.Expr] = None, source_columns: Optional[List[str]] = None, **kwargs: Any) -> None:
2234 def merge( 2235 self, 2236 target_table: TableName, 2237 source_table: QueryOrDF, 2238 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], 2239 unique_key: t.Sequence[exp.Expr], 2240 when_matched: t.Optional[exp.Whens] = None, 2241 merge_filter: t.Optional[exp.Expr] = None, 2242 source_columns: t.Optional[t.List[str]] = None, 2243 **kwargs: t.Any, 2244 ) -> None: 2245 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2246 source_table, 2247 target_columns_to_types, 2248 target_table=target_table, 2249 source_columns=source_columns, 2250 ) 2251 target_columns_to_types = target_columns_to_types or self.columns(target_table) 2252 on = exp.and_( 2253 *( 2254 add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS)) 2255 for part in unique_key 2256 ) 2257 ) 2258 if merge_filter: 2259 on = exp.and_(merge_filter, on) 2260 2261 if not when_matched: 2262 match_expressions = [ 2263 exp.When( 2264 matched=True, 2265 source=False, 2266 then=exp.Update( 2267 expressions=[ 2268 exp.column(col, MERGE_TARGET_ALIAS).eq( 2269 exp.column(col, MERGE_SOURCE_ALIAS) 2270 ) 2271 for col in target_columns_to_types 2272 ], 2273 ), 2274 ) 2275 ] 2276 else: 2277 match_expressions = when_matched.copy().expressions 2278 2279 match_expressions.append( 2280 exp.When( 2281 matched=False, 2282 source=False, 2283 then=exp.Insert( 2284 this=exp.Tuple( 2285 expressions=[exp.column(col) for col in target_columns_to_types] 2286 ), 2287 expression=exp.Tuple( 2288 expressions=[ 2289 exp.column(col, MERGE_SOURCE_ALIAS) for col in target_columns_to_types 2290 ] 2291 ), 2292 ), 2293 ) 2294 ) 2295 for source_query in source_queries: 2296 with source_query as query: 2297 self._merge( 2298 target_table=target_table, 2299 query=query, 2300 on=on, 2301 whens=exp.Whens(expressions=match_expressions), 2302 )
def
rename_table( self, old_table_name: Union[str, sqlglot.expressions.query.Table], new_table_name: Union[str, sqlglot.expressions.query.Table]) -> None:
2304 def rename_table( 2305 self, 2306 old_table_name: TableName, 2307 new_table_name: TableName, 2308 ) -> None: 2309 new_table = exp.to_table(new_table_name) 2310 if new_table.catalog: 2311 old_table = exp.to_table(old_table_name) 2312 catalog = old_table.catalog or self.get_current_catalog() 2313 if catalog != new_table.catalog: 2314 raise UnsupportedCatalogOperationError( 2315 "Tried to rename table across catalogs which is not supported" 2316 ) 2317 self._rename_table(old_table_name, new_table_name) 2318 self._clear_data_object_cache(old_table_name) 2319 self._clear_data_object_cache(new_table_name)
def
get_data_object( self, target_name: Union[str, sqlglot.expressions.query.Table], safe_to_cache: bool = False) -> Optional[sqlmesh.core.engine_adapter.shared.DataObject]:
2321 def get_data_object( 2322 self, target_name: TableName, safe_to_cache: bool = False 2323 ) -> t.Optional[DataObject]: 2324 target_table = exp.to_table(target_name) 2325 existing_data_objects = self.get_data_objects( 2326 schema_(target_table.db, target_table.catalog), 2327 {target_table.name}, 2328 safe_to_cache=safe_to_cache, 2329 ) 2330 if existing_data_objects: 2331 return existing_data_objects[0] 2332 return None
def
get_data_objects( self, schema_name: Union[str, sqlglot.expressions.query.Table], object_names: Optional[Set[str]] = None, safe_to_cache: bool = False) -> List[sqlmesh.core.engine_adapter.shared.DataObject]:
2334 def get_data_objects( 2335 self, 2336 schema_name: SchemaName, 2337 object_names: t.Optional[t.Set[str]] = None, 2338 safe_to_cache: bool = False, 2339 ) -> t.List[DataObject]: 2340 """Lists all data objects in the target schema. 2341 2342 Args: 2343 schema_name: The name of the schema to list data objects from. 2344 object_names: If provided, only return data objects with these names. 2345 safe_to_cache: Whether it is safe to cache the results of this call. 2346 2347 Returns: 2348 A list of data objects in the target schema. 2349 """ 2350 if object_names is not None: 2351 if not object_names: 2352 return [] 2353 2354 # Check cache for each object name 2355 target_schema = to_schema(schema_name) 2356 cached_objects = [] 2357 missing_names = set() 2358 2359 for name in object_names: 2360 cache_key = _get_data_object_cache_key( 2361 target_schema.catalog, target_schema.db, name 2362 ) 2363 if cache_key in self._data_object_cache: 2364 logger.debug("Data object cache hit: %s", cache_key) 2365 data_object = self._data_object_cache[cache_key] 2366 # If the object is none, then the table was previously looked for but not found 2367 if data_object: 2368 cached_objects.append(data_object) 2369 else: 2370 logger.debug("Data object cache miss: %s", cache_key) 2371 missing_names.add(name) 2372 2373 # Fetch missing objects from database 2374 if missing_names: 2375 object_names_list = list(missing_names) 2376 batches = [ 2377 object_names_list[i : i + self.DATA_OBJECT_FILTER_BATCH_SIZE] 2378 for i in range(0, len(object_names_list), self.DATA_OBJECT_FILTER_BATCH_SIZE) 2379 ] 2380 2381 fetched_objects = [] 2382 fetched_object_names = set() 2383 for batch in batches: 2384 objects = self._get_data_objects(schema_name, set(batch)) 2385 for obj in objects: 2386 if safe_to_cache: 2387 cache_key = _get_data_object_cache_key( 2388 obj.catalog, obj.schema_name, obj.name 2389 ) 2390 self._data_object_cache[cache_key] = obj 2391 fetched_objects.append(obj) 2392 fetched_object_names.add(obj.name) 2393 2394 if safe_to_cache: 2395 for missing_name in missing_names - fetched_object_names: 2396 cache_key = _get_data_object_cache_key( 2397 target_schema.catalog, target_schema.db, missing_name 2398 ) 2399 self._data_object_cache[cache_key] = None 2400 2401 return cached_objects + fetched_objects 2402 2403 return cached_objects 2404 2405 fetched_objects = self._get_data_objects(schema_name) 2406 if safe_to_cache: 2407 for obj in fetched_objects: 2408 cache_key = _get_data_object_cache_key(obj.catalog, obj.schema_name, obj.name) 2409 self._data_object_cache[cache_key] = obj 2410 return fetched_objects
Lists all data objects in the target schema.
Arguments:
- schema_name: The name of the schema to list data objects from.
- object_names: If provided, only return data objects with these names.
- safe_to_cache: Whether it is safe to cache the results of this call.
Returns:
A list of data objects in the target schema.
def
fetchone( self, query: Union[sqlglot.expressions.core.Expr, str], ignore_unsupported_errors: bool = False, quote_identifiers: bool = False) -> Optional[Tuple]:
2412 def fetchone( 2413 self, 2414 query: t.Union[exp.Expr, str], 2415 ignore_unsupported_errors: bool = False, 2416 quote_identifiers: bool = False, 2417 ) -> t.Optional[t.Tuple]: 2418 with self.transaction(): 2419 self.execute( 2420 query, 2421 ignore_unsupported_errors=ignore_unsupported_errors, 2422 quote_identifiers=quote_identifiers, 2423 ) 2424 return self.cursor.fetchone()
def
fetchall( self, query: Union[sqlglot.expressions.core.Expr, str], ignore_unsupported_errors: bool = False, quote_identifiers: bool = False) -> List[Tuple]:
2426 def fetchall( 2427 self, 2428 query: t.Union[exp.Expr, str], 2429 ignore_unsupported_errors: bool = False, 2430 quote_identifiers: bool = False, 2431 ) -> t.List[t.Tuple]: 2432 with self.transaction(): 2433 self.execute( 2434 query, 2435 ignore_unsupported_errors=ignore_unsupported_errors, 2436 quote_identifiers=quote_identifiers, 2437 ) 2438 return self.cursor.fetchall()
def
fetchdf( self, query: Union[sqlglot.expressions.core.Expr, str], quote_identifiers: bool = False) -> pandas.core.frame.DataFrame:
2463 def fetchdf( 2464 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2465 ) -> pd.DataFrame: 2466 """Fetches a Pandas DataFrame from the cursor""" 2467 import pandas as pd 2468 2469 df = self._fetch_native_df(query, quote_identifiers=quote_identifiers) 2470 if not isinstance(df, pd.DataFrame): 2471 raise NotImplementedError( 2472 "The cursor's `fetch_native_df` method is not returning a pandas DataFrame. Need to update `fetchdf` so a Pandas DataFrame is returned" 2473 ) 2474 return df
Fetches a Pandas DataFrame from the cursor
def
fetch_pyspark_df( self, query: Union[sqlglot.expressions.core.Expr, str], quote_identifiers: bool = False) -> <MagicMock id='130969803882896'>:
2476 def fetch_pyspark_df( 2477 self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False 2478 ) -> PySparkDataFrame: 2479 """Fetches a PySpark DataFrame from the cursor""" 2480 raise NotImplementedError(f"Engine does not support PySpark DataFrames: {type(self)}")
Fetches a PySpark DataFrame from the cursor
wap_enabled: bool
2482 @property 2483 def wap_enabled(self) -> bool: 2484 """Returns whether WAP is enabled for this engine.""" 2485 return self._extra_config.get("wap_enabled", False)
Returns whether WAP is enabled for this engine.
def
wap_supported(self, table_name: Union[str, sqlglot.expressions.query.Table]) -> bool:
2487 def wap_supported(self, table_name: TableName) -> bool: 2488 """Returns whether WAP for the target table is supported.""" 2489 return False
Returns whether WAP for the target table is supported.
def
wap_table_name( self, table_name: Union[str, sqlglot.expressions.query.Table], wap_id: str) -> str:
2491 def wap_table_name(self, table_name: TableName, wap_id: str) -> str: 2492 """Returns the updated table name for the given WAP ID. 2493 2494 Args: 2495 table_name: The name of the target table. 2496 wap_id: The WAP ID to prepare. 2497 2498 Returns: 2499 The updated table name that should be used for writing. 2500 """ 2501 raise NotImplementedError(f"Engine does not support WAP: {type(self)}")
Returns the updated table name for the given WAP ID.
Arguments:
- table_name: The name of the target table.
- wap_id: The WAP ID to prepare.
Returns:
The updated table name that should be used for writing.
def
wap_prepare( self, table_name: Union[str, sqlglot.expressions.query.Table], wap_id: str) -> str:
2503 def wap_prepare(self, table_name: TableName, wap_id: str) -> str: 2504 """Prepares the target table for WAP and returns the updated table name. 2505 2506 Args: 2507 table_name: The name of the target table. 2508 wap_id: The WAP ID to prepare. 2509 2510 Returns: 2511 The updated table name that should be used for writing. 2512 """ 2513 raise NotImplementedError(f"Engine does not support WAP: {type(self)}")
Prepares the target table for WAP and returns the updated table name.
Arguments:
- table_name: The name of the target table.
- wap_id: The WAP ID to prepare.
Returns:
The updated table name that should be used for writing.
def
wap_publish( self, table_name: Union[str, sqlglot.expressions.query.Table], wap_id: str) -> None:
2515 def wap_publish(self, table_name: TableName, wap_id: str) -> None: 2516 """Publishes changes with the given WAP ID to the target table. 2517 2518 Args: 2519 table_name: The name of the target table. 2520 wap_id: The WAP ID to publish. 2521 """ 2522 raise NotImplementedError(f"Engine does not support WAP: {type(self)}")
Publishes changes with the given WAP ID to the target table.
Arguments:
- table_name: The name of the target table.
- wap_id: The WAP ID to publish.
def
sync_grants_config( self, table: sqlglot.expressions.query.Table, grants_config: <MagicMock id='130969803989408'>, table_type: sqlmesh.core.engine_adapter.shared.DataObjectType = <DataObjectType.TABLE: 'table'>) -> None:
2524 def sync_grants_config( 2525 self, 2526 table: exp.Table, 2527 grants_config: GrantsConfig, 2528 table_type: DataObjectType = DataObjectType.TABLE, 2529 ) -> None: 2530 """Applies the grants_config to a table authoritatively. 2531 It first compares the specified grants against the current grants, and then 2532 applies the diffs to the table by revoking and granting privileges as needed. 2533 2534 Args: 2535 table: The table/view to apply grants to. 2536 grants_config: Dictionary mapping privileges to lists of grantees. 2537 table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW). 2538 """ 2539 if not self.SUPPORTS_GRANTS: 2540 raise NotImplementedError(f"Engine does not support grants: {type(self)}") 2541 2542 current_grants = self._get_current_grants_config(table) 2543 new_grants, revoked_grants = self._diff_grants_configs(grants_config, current_grants) 2544 revoke_exprs = self._revoke_grants_config_expr(table, revoked_grants, table_type) 2545 grant_exprs = self._apply_grants_config_expr(table, new_grants, table_type) 2546 dcl_exprs = revoke_exprs + grant_exprs 2547 2548 if dcl_exprs: 2549 self.execute(dcl_exprs)
Applies the grants_config to a table authoritatively. It first compares the specified grants against the current grants, and then applies the diffs to the table by revoking and granting privileges as needed.
Arguments:
- table: The table/view to apply grants to.
- grants_config: Dictionary mapping privileges to lists of grantees.
- table_type: The type of database object (TABLE, VIEW, MATERIALIZED_VIEW).
@contextlib.contextmanager
def
transaction(self, condition: Optional[bool] = None) -> Iterator[NoneType]:
2551 @contextlib.contextmanager 2552 def transaction( 2553 self, 2554 condition: t.Optional[bool] = None, 2555 ) -> t.Iterator[None]: 2556 """A transaction context manager.""" 2557 if ( 2558 self._connection_pool.is_transaction_active 2559 or not self.SUPPORTS_TRANSACTIONS 2560 or (condition is not None and not condition) 2561 ): 2562 yield 2563 return 2564 2565 if self._pre_ping: 2566 try: 2567 logger.debug("Pinging the database to check the connection") 2568 self.ping() 2569 except Exception: 2570 logger.info("Connection to the database was lost. Reconnecting...") 2571 self._connection_pool.close() 2572 2573 self._connection_pool.begin() 2574 try: 2575 yield 2576 except Exception as e: 2577 self._connection_pool.rollback() 2578 raise e 2579 else: 2580 self._connection_pool.commit()
A transaction context manager.
@contextlib.contextmanager
def
session( self, properties: Dict[str, sqlglot.expressions.core.Expr | str | int | float | bool]) -> Iterator[NoneType]:
2582 @contextlib.contextmanager 2583 def session(self, properties: SessionProperties) -> t.Iterator[None]: 2584 """A session context manager.""" 2585 if self._is_session_active(): 2586 yield 2587 return 2588 2589 self._begin_session(properties) 2590 try: 2591 yield 2592 finally: 2593 self._end_session()
A session context manager.
def
execute( self, expressions: Union[str, sqlglot.expressions.core.Expr, Sequence[sqlglot.expressions.core.Expr]], ignore_unsupported_errors: bool = False, quote_identifiers: bool = True, track_rows_processed: bool = False, **kwargs: Any) -> None:
2605 def execute( 2606 self, 2607 expressions: t.Union[str, exp.Expr, t.Sequence[exp.Expr]], 2608 ignore_unsupported_errors: bool = False, 2609 quote_identifiers: bool = True, 2610 track_rows_processed: bool = False, 2611 **kwargs: t.Any, 2612 ) -> None: 2613 """Execute a sql query.""" 2614 to_sql_kwargs = ( 2615 {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {} 2616 ) 2617 with self.transaction(): 2618 for e in ensure_list(expressions): 2619 if isinstance(e, exp.Expr): 2620 self._check_identifier_length(e) 2621 sql = self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs) 2622 else: 2623 sql = t.cast(str, e) 2624 2625 sql = self._attach_correlation_id(sql) 2626 2627 self._log_sql( 2628 sql, 2629 expression=e if isinstance(e, exp.Expr) else None, 2630 quote_identifiers=quote_identifiers, 2631 ) 2632 self._execute(sql, track_rows_processed, **kwargs)
Execute a sql query.
@contextlib.contextmanager
def
temp_table( self, query_or_df: <MagicMock id='130969804860832'>, name: Union[str, sqlglot.expressions.query.Table] = 'diff', target_columns_to_types: Optional[Dict[str, sqlglot.expressions.datatypes.DataType]] = None, source_columns: Optional[List[str]] = None, **kwargs: Any) -> Iterator[sqlglot.expressions.query.Table]:
2680 @contextlib.contextmanager 2681 def temp_table( 2682 self, 2683 query_or_df: QueryOrDF, 2684 name: TableName = "diff", 2685 target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, 2686 source_columns: t.Optional[t.List[str]] = None, 2687 **kwargs: t.Any, 2688 ) -> t.Iterator[exp.Table]: 2689 """A context manager for working a temp table. 2690 2691 The table will be created with a random guid and cleaned up after the block. 2692 2693 Args: 2694 query_or_df: The query or df to create a temp table for. 2695 name: The base name of the temp table. 2696 target_columns_to_types: A mapping between the column name and its data type. 2697 2698 Yields: 2699 The table expression 2700 """ 2701 name = exp.to_table(name) 2702 # ensure that we use default catalog if none is not specified 2703 if isinstance(name, exp.Table) and not name.catalog and name.db and self.default_catalog: 2704 name.set("catalog", exp.parse_identifier(self.default_catalog)) 2705 2706 source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( 2707 query_or_df, 2708 target_columns_to_types=target_columns_to_types, 2709 target_table=name, 2710 source_columns=source_columns, 2711 ) 2712 2713 with self.transaction(): 2714 table = self._get_temp_table(name) 2715 if table.db: 2716 self.create_schema(schema_(table.args["db"], table.args.get("catalog"))) 2717 self._create_table_from_source_queries( 2718 table, 2719 source_queries, 2720 target_columns_to_types, 2721 exists=True, 2722 table_description=None, 2723 column_descriptions=None, 2724 track_rows_processed=False, 2725 **kwargs, 2726 ) 2727 2728 try: 2729 yield table 2730 finally: 2731 self.drop_table(table)
A context manager for working a temp table.
The table will be created with a random guid and cleaned up after the block.
Arguments:
- query_or_df: The query or df to create a temp table for.
- name: The base name of the temp table.
- target_columns_to_types: A mapping between the column name and its data type.
Yields:
The table expression
def
adjust_physical_properties_for_incremental( self, physical_properties: Dict[str, Any], *, requires_delete_capable_table: bool, unique_key: Optional[List[sqlglot.expressions.core.Expr]], model_name: str) -> Dict[str, Any]:
2762 def adjust_physical_properties_for_incremental( 2763 self, 2764 physical_properties: t.Dict[str, t.Any], 2765 *, 2766 requires_delete_capable_table: bool, 2767 unique_key: t.Optional[t.List[exp.Expr]], 2768 model_name: str, 2769 ) -> t.Dict[str, t.Any]: 2770 """Adjusts physical properties for an incremental model before the table is created. 2771 2772 Some engines require a specific physical table layout before they can run the DELETE/MERGE 2773 statements that incremental model kinds rely on (e.g. StarRocks only supports those on 2774 PRIMARY KEY tables). This hook lets each engine derive or validate the required properties 2775 while keeping the generic evaluator free of engine-specific branching. 2776 2777 Args: 2778 physical_properties: The model's physical properties. 2779 requires_delete_capable_table: Whether the model kind issues DELETE/MERGE statements 2780 (as opposed to append-only INSERTs), as determined by the generic evaluator. 2781 unique_key: The model's unique key, populated only when the kind allows promoting it to 2782 an engine-specific key (i.e. INCREMENTAL_BY_UNIQUE_KEY); otherwise None. 2783 model_name: The model name, for use in diagnostics. 2784 2785 Returns: 2786 The (possibly adjusted) physical properties. Implementations own the given mapping and 2787 may mutate it in place; the base implementation returns it unchanged. 2788 """ 2789 return physical_properties
Adjusts physical properties for an incremental model before the table is created.
Some engines require a specific physical table layout before they can run the DELETE/MERGE statements that incremental model kinds rely on (e.g. StarRocks only supports those on PRIMARY KEY tables). This hook lets each engine derive or validate the required properties while keeping the generic evaluator free of engine-specific branching.
Arguments:
- physical_properties: The model's physical properties.
- requires_delete_capable_table: Whether the model kind issues DELETE/MERGE statements (as opposed to append-only INSERTs), as determined by the generic evaluator.
- unique_key: The model's unique key, populated only when the kind allows promoting it to an engine-specific key (i.e. INCREMENTAL_BY_UNIQUE_KEY); otherwise None.
- model_name: The model name, for use in diagnostics.
Returns:
The (possibly adjusted) physical properties. Implementations own the given mapping and may mutate it in place; the base implementation returns it unchanged.
def
drop_data_object_on_type_mismatch( self, data_object: Optional[sqlmesh.core.engine_adapter.shared.DataObject], expected_type: sqlmesh.core.engine_adapter.shared.DataObjectType) -> bool:
2946 def drop_data_object_on_type_mismatch( 2947 self, data_object: t.Optional[DataObject], expected_type: DataObjectType 2948 ) -> bool: 2949 """Drops a data object if it exists and is not of the expected type. 2950 2951 Args: 2952 data_object: The data object to check. 2953 expected_type: The expected type of the data object. 2954 2955 Returns: 2956 True if the data object was dropped, False otherwise. 2957 """ 2958 if data_object is None or data_object.type == expected_type: 2959 return False 2960 2961 logger.warning( 2962 "Target data object '%s' is a %s and not a %s, dropping it", 2963 data_object.to_table().sql(dialect=self.dialect), 2964 data_object.type.value, 2965 expected_type.value, 2966 ) 2967 self.drop_data_object(data_object) 2968 return True
Drops a data object if it exists and is not of the expected type.
Arguments:
- data_object: The data object to check.
- expected_type: The expected type of the data object.
Returns:
True if the data object was dropped, False otherwise.
def
ensure_nulls_for_unmatched_after_join( self, query: <MagicMock id='130969805323856'>) -> <MagicMock id='130969805323856'>:
def
use_server_nulls_for_unmatched_after_join( self, query: <MagicMock id='130969805323856'>) -> <MagicMock id='130969805323856'>:
Base class wrapping a Database API compliant connection.
The EngineAdapter is an easily-subclassable interface that interacts with the underlying engine and data store.
Arguments:
- connection_factory_or_pool: a callable which produces a new Database API-compliant connection on every call.
- dialect: The dialect with which this adapter is associated.
- multithreaded: Indicates whether this adapter will be used by more than one thread.
Inherited Members
- EngineAdapter
- EngineAdapter
- DIALECT
- DEFAULT_BATCH_SIZE
- DATA_OBJECT_FILTER_BATCH_SIZE
- SUPPORTS_TRANSACTIONS
- COMMENT_CREATION_TABLE
- COMMENT_CREATION_VIEW
- MAX_TABLE_COMMENT_LENGTH
- MAX_COLUMN_COMMENT_LENGTH
- INSERT_OVERWRITE_STRATEGY
- SUPPORTS_MATERIALIZED_VIEWS
- SUPPORTS_MATERIALIZED_VIEW_SCHEMA
- SUPPORTS_VIEW_SCHEMA
- SUPPORTS_CLONING
- SUPPORTS_MANAGED_MODELS
- SUPPORTS_CREATE_DROP_CATALOG
- SUPPORTED_DROP_CASCADE_OBJECT_KINDS
- SCHEMA_DIFFER_KWARGS
- SUPPORTS_TUPLE_IN
- HAS_VIEW_BINDING
- RECREATE_MATERIALIZED_VIEW_ON_EVALUATION
- SUPPORTS_REPLACE_TABLE
- SUPPORTS_GRANTS
- DEFAULT_CATALOG_TYPE
- QUOTE_IDENTIFIERS_IN_VIEWS
- MAX_IDENTIFIER_LENGTH
- ATTACH_CORRELATION_ID
- SUPPORTS_QUERY_EXECUTION_TRACKING
- SUPPORTS_METADATA_TABLE_LAST_MODIFIED_TS
- RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES
- dialect
- correlation_id
- with_settings
- cursor
- connection
- spark
- snowpark
- bigframe
- comments_enabled
- catalog_support
- supports_virtual_catalog
- inject_virtual_catalog
- schema_differ
- default_catalog
- engine_run_mode
- recycle
- close
- get_current_catalog
- set_current_catalog
- get_catalog_type
- get_catalog_type_from_table
- current_catalog_type
- replace_query
- create_index
- create_table
- create_managed_table
- ctas
- create_state_table
- create_table_like
- clone_table
- drop_data_object
- drop_table
- drop_managed_table
- get_alter_operations
- alter_table
- create_view
- create_schema
- drop_schema
- drop_view
- create_catalog
- drop_catalog
- columns
- table_exists
- delete_from
- insert_append
- insert_overwrite_by_partition
- insert_overwrite_by_time_partition
- update_table
- scd_type_2_by_time
- scd_type_2_by_column
- merge
- rename_table
- get_data_object
- get_data_objects
- fetchone
- fetchall
- fetchdf
- fetch_pyspark_df
- wap_enabled
- wap_supported
- wap_table_name
- wap_prepare
- wap_publish
- sync_grants_config
- transaction
- session
- execute
- temp_table
- adjust_physical_properties_for_incremental
- drop_data_object_on_type_mismatch
- ensure_nulls_for_unmatched_after_join
- use_server_nulls_for_unmatched_after_join
- ping
- get_table_last_modified_ts