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