meerschaum.utils.sql
Flavor-specific SQL tools.
Only sql_item_name and build_where are used by the Rust engine (they delegate to it). The rest
(get_update_queries, get_create_table_queries, dateadd_str, table_exists, ...) are SDK helpers
for plugin and custom-connector authors: they keep their v4 implementations and are not used by
Pipe syncs. With a SQLConnector they run through its Rust connection pool (raw SQL, rc15 Q53:
table_exists, get_table_cols_types, get_table_cols_indices need no SQLAlchemy); a SQLAlchemy
session, engine or connection keeps working as in v4 (the sql extra).
1#! /usr/bin/env python 2# -*- coding: utf-8 -*- 3# vim:fenc=utf-8 4 5""" 6Flavor-specific SQL tools. 7 8Only `sql_item_name` and `build_where` are used by the Rust engine (they delegate to it). The rest 9(`get_update_queries`, `get_create_table_queries`, `dateadd_str`, `table_exists`, ...) are SDK helpers 10for plugin and custom-connector authors: they keep their v4 implementations and are not used by 11`Pipe` syncs. With a `SQLConnector` they run through its Rust connection pool (raw SQL, rc15 Q53: 12`table_exists`, `get_table_cols_types`, `get_table_cols_indices` need no SQLAlchemy); a SQLAlchemy 13session, engine or connection keeps working as in v4 (the `sql` extra). 14""" 15 16from __future__ import annotations 17 18from datetime import datetime, timezone 19import meerschaum as mrsm 20from meerschaum.utils.typing import Optional, Dict, Any, Union, List, Tuple, Iterable 21### Preserve legacy imports. 22from meerschaum.utils.dtypes.sql import ( 23 DB_TO_PD_DTYPES, 24 PD_TO_DB_DTYPES_FLAVORS, 25 get_pd_type_from_db_type as get_pd_type, 26 get_db_type_from_pd_type as get_db_type, 27 TIMEZONE_NAIVE_FLAVORS, 28) 29from meerschaum.utils.warnings import warn 30from meerschaum.utils.debug import dprint 31from meerschaum import _core 32 33test_queries = { 34 'default' : 'SELECT 1', 35 'oracle' : 'SELECT 1 FROM DUAL', 36 'informix' : 'SELECT COUNT(*) FROM systables', 37 'hsqldb' : 'SELECT 1 FROM INFORMATION_SCHEMA.SYSTEM_USERS', 38} 39version_queries = { 40 'default': "SELECT VERSION() AS {version_name}", 41 'sqlite': "SELECT SQLITE_VERSION() AS {version_name}", 42 'geopackage': "SELECT SQLITE_VERSION() AS {version_name}", 43 'mssql': "SELECT @@version", 44 'oracle': "SELECT version from PRODUCT_COMPONENT_VERSION WHERE rownum = 1", 45} 46SKIP_IF_EXISTS_FLAVORS = {'mssql', 'oracle'} 47DROP_IF_EXISTS_FLAVORS = { 48 'timescaledb', 49 'timescaledb-ha', 50 'postgresql', 51 'postgis', 52 'citus', 53 'mssql', 54 'mysql', 55 'mariadb', 56 'sqlite', 57 'geopackage', 58} 59SKIP_AUTO_INCREMENT_FLAVORS = {'citus', 'duckdb'} 60columns_types_queries = { 61 'default': """ 62 SELECT 63 table_catalog AS database, 64 table_schema AS schema, 65 table_name AS table, 66 column_name AS column, 67 data_type AS type, 68 numeric_precision, 69 numeric_scale 70 FROM information_schema.columns 71 WHERE table_name IN ('{table}', '{table_trunc}') 72 """, 73 'sqlite': """ 74 SELECT 75 '' "database", 76 '' "schema", 77 m.name "table", 78 p.name "column", 79 p.type "type" 80 FROM sqlite_master m 81 LEFT OUTER JOIN pragma_table_info(m.name) p 82 ON m.name <> p.name 83 WHERE m.type = 'table' 84 AND m.name IN ('{table}', '{table_trunc}') 85 """, 86 'geopackage': """ 87 SELECT 88 '' "database", 89 '' "schema", 90 m.name "table", 91 p.name "column", 92 p.type "type" 93 FROM sqlite_master m 94 LEFT OUTER JOIN pragma_table_info(m.name) p 95 ON m.name <> p.name 96 WHERE m.type = 'table' 97 AND m.name IN ('{table}', '{table_trunc}') 98 """, 99 'mssql': """ 100 SELECT 101 TABLE_CATALOG AS [database], 102 TABLE_SCHEMA AS [schema], 103 TABLE_NAME AS [table], 104 COLUMN_NAME AS [column], 105 CASE WHEN CHARACTER_MAXIMUM_LENGTH = -1 THEN DATA_TYPE + '(max)' ELSE DATA_TYPE END AS [type], 106 NUMERIC_PRECISION AS [numeric_precision], 107 NUMERIC_SCALE AS [numeric_scale] 108 FROM {db_prefix}INFORMATION_SCHEMA.COLUMNS WITH (NOLOCK) 109 WHERE TABLE_NAME IN ( 110 '{table}', 111 '{table_trunc}' 112 ) 113 114 """, 115 'mysql': """ 116 SELECT 117 TABLE_SCHEMA `database`, 118 TABLE_SCHEMA `schema`, 119 TABLE_NAME `table`, 120 COLUMN_NAME `column`, 121 DATA_TYPE `type`, 122 NUMERIC_PRECISION `numeric_precision`, 123 NUMERIC_SCALE `numeric_scale` 124 FROM INFORMATION_SCHEMA.COLUMNS 125 WHERE TABLE_NAME IN ('{table}', '{table_trunc}') 126 """, 127 'mariadb': """ 128 SELECT 129 TABLE_SCHEMA `database`, 130 TABLE_SCHEMA `schema`, 131 TABLE_NAME `table`, 132 COLUMN_NAME `column`, 133 DATA_TYPE `type`, 134 NUMERIC_PRECISION `numeric_precision`, 135 NUMERIC_SCALE `numeric_scale` 136 FROM INFORMATION_SCHEMA.COLUMNS 137 WHERE TABLE_NAME IN ('{table}', '{table_trunc}') 138 """, 139 'oracle': """ 140 SELECT 141 NULL AS "database", 142 NULL AS "schema", 143 TABLE_NAME AS "table", 144 COLUMN_NAME AS "column", 145 DATA_TYPE AS "type", 146 DATA_PRECISION AS "numeric_precision", 147 DATA_SCALE AS "numeric_scale" 148 FROM all_tab_columns 149 WHERE TABLE_NAME IN ( 150 '{table}', 151 '{table_trunc}', 152 '{table_lower}', 153 '{table_lower_trunc}', 154 '{table_upper}', 155 '{table_upper_trunc}' 156 ) 157 """, 158} 159columns_indices_queries = { 160 'default': """ 161 SELECT 162 current_database() AS "database", 163 n.nspname AS "schema", 164 t.relname AS "table", 165 c.column_name AS "column", 166 i.relname AS "index", 167 CASE WHEN con.contype = 'p' THEN 'PRIMARY KEY' ELSE 'INDEX' END AS "index_type" 168 FROM pg_class t 169 INNER JOIN pg_index AS ix 170 ON t.oid = ix.indrelid 171 INNER JOIN pg_class AS i 172 ON i.oid = ix.indexrelid 173 INNER JOIN pg_namespace AS n 174 ON n.oid = t.relnamespace 175 INNER JOIN pg_attribute AS a 176 ON a.attnum = ANY(ix.indkey) 177 AND a.attrelid = t.oid 178 INNER JOIN information_schema.columns AS c 179 ON c.column_name = a.attname 180 AND c.table_name = t.relname 181 AND c.table_schema = n.nspname 182 LEFT JOIN pg_constraint AS con 183 ON con.conindid = i.oid 184 AND con.contype = 'p' 185 WHERE 186 t.relname IN ('{table}', '{table_trunc}') 187 AND n.nspname = '{schema}' 188 """, 189 'sqlite': """ 190 WITH indexed_columns AS ( 191 SELECT 192 '{table}' AS table_name, 193 pi.name AS column_name, 194 i.name AS index_name, 195 'INDEX' AS index_type 196 FROM 197 sqlite_master AS i, 198 pragma_index_info(i.name) AS pi 199 WHERE 200 i.type = 'index' 201 AND i.tbl_name = '{table}' 202 ), 203 primary_key_columns AS ( 204 SELECT 205 '{table}' AS table_name, 206 ti.name AS column_name, 207 'PRIMARY_KEY' AS index_name, 208 'PRIMARY KEY' AS index_type 209 FROM 210 pragma_table_info('{table}') AS ti 211 WHERE 212 ti.pk > 0 213 ) 214 SELECT 215 NULL AS "database", 216 NULL AS "schema", 217 "table_name" AS "table", 218 "column_name" AS "column", 219 "index_name" AS "index", 220 "index_type" 221 FROM indexed_columns 222 UNION ALL 223 SELECT 224 NULL AS "database", 225 NULL AS "schema", 226 table_name AS "table", 227 column_name AS "column", 228 index_name AS "index", 229 index_type 230 FROM primary_key_columns 231 """, 232 'geopackage': """ 233 WITH indexed_columns AS ( 234 SELECT 235 '{table}' AS table_name, 236 pi.name AS column_name, 237 i.name AS index_name, 238 'INDEX' AS index_type 239 FROM 240 sqlite_master AS i, 241 pragma_index_info(i.name) AS pi 242 WHERE 243 i.type = 'index' 244 AND i.tbl_name = '{table}' 245 ), 246 primary_key_columns AS ( 247 SELECT 248 '{table}' AS table_name, 249 ti.name AS column_name, 250 'PRIMARY_KEY' AS index_name, 251 'PRIMARY KEY' AS index_type 252 FROM 253 pragma_table_info('{table}') AS ti 254 WHERE 255 ti.pk > 0 256 ) 257 SELECT 258 NULL AS "database", 259 NULL AS "schema", 260 "table_name" AS "table", 261 "column_name" AS "column", 262 "index_name" AS "index", 263 "index_type" 264 FROM indexed_columns 265 UNION ALL 266 SELECT 267 NULL AS "database", 268 NULL AS "schema", 269 table_name AS "table", 270 column_name AS "column", 271 index_name AS "index", 272 index_type 273 FROM primary_key_columns 274 """, 275 'mssql': """ 276 SELECT 277 NULL AS [database], 278 s.name AS [schema], 279 t.name AS [table], 280 c.name AS [column], 281 i.name AS [index], 282 CASE 283 WHEN kc.type = 'PK' THEN 'PRIMARY KEY' 284 ELSE 'INDEX' 285 END AS [index_type], 286 CASE 287 WHEN i.type = 1 THEN CAST(1 AS BIT) 288 ELSE CAST(0 AS BIT) 289 END AS [clustered] 290 FROM 291 sys.schemas s WITH (NOLOCK) 292 INNER JOIN sys.tables t WITH (NOLOCK) 293 ON s.schema_id = t.schema_id 294 INNER JOIN sys.indexes i WITH (NOLOCK) 295 ON t.object_id = i.object_id 296 INNER JOIN sys.index_columns ic WITH (NOLOCK) 297 ON i.object_id = ic.object_id 298 AND i.index_id = ic.index_id 299 INNER JOIN sys.columns c WITH (NOLOCK) 300 ON ic.object_id = c.object_id 301 AND ic.column_id = c.column_id 302 LEFT JOIN sys.key_constraints kc WITH (NOLOCK) 303 ON kc.parent_object_id = i.object_id 304 AND kc.type = 'PK' 305 AND kc.name = i.name 306 WHERE 307 t.name IN ('{table}', '{table_trunc}') 308 AND s.name = '{schema}' 309 AND i.type IN (1, 2) 310 """, 311 'oracle': """ 312 SELECT 313 NULL AS "database", 314 ic.table_owner AS "schema", 315 ic.table_name AS "table", 316 ic.column_name AS "column", 317 i.index_name AS "index", 318 CASE 319 WHEN c.constraint_type = 'P' THEN 'PRIMARY KEY' 320 WHEN i.uniqueness = 'UNIQUE' THEN 'UNIQUE INDEX' 321 ELSE 'INDEX' 322 END AS index_type 323 FROM 324 all_ind_columns ic 325 INNER JOIN all_indexes i 326 ON ic.index_name = i.index_name 327 AND ic.table_owner = i.owner 328 LEFT JOIN all_constraints c 329 ON i.index_name = c.constraint_name 330 AND i.table_owner = c.owner 331 AND c.constraint_type = 'P' 332 WHERE ic.table_name IN ( 333 '{table}', 334 '{table_trunc}', 335 '{table_upper}', 336 '{table_upper_trunc}' 337 ) 338 """, 339 'mysql': """ 340 SELECT 341 TABLE_SCHEMA AS `database`, 342 TABLE_SCHEMA AS `schema`, 343 TABLE_NAME AS `table`, 344 COLUMN_NAME AS `column`, 345 INDEX_NAME AS `index`, 346 CASE 347 WHEN NON_UNIQUE = 0 THEN 'PRIMARY KEY' 348 ELSE 'INDEX' 349 END AS `index_type` 350 FROM 351 information_schema.STATISTICS 352 WHERE 353 TABLE_NAME IN ('{table}', '{table_trunc}') 354 """, 355 'mariadb': """ 356 SELECT 357 TABLE_SCHEMA AS `database`, 358 TABLE_SCHEMA AS `schema`, 359 TABLE_NAME AS `table`, 360 COLUMN_NAME AS `column`, 361 INDEX_NAME AS `index`, 362 CASE 363 WHEN NON_UNIQUE = 0 THEN 'PRIMARY KEY' 364 ELSE 'INDEX' 365 END AS `index_type` 366 FROM 367 information_schema.STATISTICS 368 WHERE 369 TABLE_NAME IN ('{table}', '{table_trunc}') 370 """, 371} 372reset_autoincrement_queries: Dict[str, Union[str, List[str]]] = { 373 'default': """ 374 SELECT SETVAL(pg_get_serial_sequence('{table_name}', '{column}'), {val}) 375 FROM {table_name} 376 """, 377 'mssql': """ 378 DBCC CHECKIDENT ('{table_name}', RESEED, {val}) 379 """, 380 'mysql': """ 381 ALTER TABLE {table_name} AUTO_INCREMENT = {val} 382 """, 383 'mariadb': """ 384 ALTER TABLE {table_name} AUTO_INCREMENT = {val} 385 """, 386 'sqlite': """ 387 UPDATE sqlite_sequence 388 SET seq = {val} 389 WHERE name = '{table}' 390 """, 391 'geopackage': """ 392 UPDATE sqlite_sequence 393 SET seq = {val} 394 WHERE name = '{table}' 395 """, 396 'oracle': ( 397 "ALTER TABLE {table_name} MODIFY {column_name} " 398 "GENERATED BY DEFAULT ON NULL AS IDENTITY (START WITH {val_plus_1})" 399 ), 400} 401max_name_lens = { 402 'default' : 64, 403 'mssql' : 128, 404 'oracle' : 30, 405 'postgresql' : 64, 406 'postgis' : 64, 407 'timescaledb' : 64, 408 'timescaledb-ha': 64, 409 'citus' : 64, 410 'cockroachdb' : 64, 411 'sqlite' : 1024, 412 'geopackage' : 1024, 413 'mysql' : 64, 414 'mariadb' : 64, 415} 416json_flavors = { 417 'postgresql', 418 'postgis', 419 'timescaledb', 420 'timescaledb-ha', 421 'citus', 422 'cockroachdb', 423} 424NO_SCHEMA_FLAVORS = { 425 'oracle', 426 'sqlite', 427 'geopackage', 428 'mysql', 429 'mariadb', 430 'duckdb', 431} 432DEFAULT_SCHEMA_FLAVORS = { 433 'postgresql': 'public', 434 'postgis': 'public', 435 'timescaledb': 'public', 436 'timescaledb-ha': 'public', 437 'citus': 'public', 438 'cockroachdb': 'public', 439 'mysql': 'mysql', 440 'mariadb': 'mysql', 441 'mssql': 'dbo', 442} 443NO_CTE_FLAVORS = { 444 'mysql', 445 'mariadb', 446} 447def clean(substring: str) -> None: 448 """ 449 Ensure a substring is clean enough to be inserted into a SQL query. 450 Raises an exception when banned words are used. 451 """ 452 from meerschaum.utils.warnings import error 453 banned_symbols = [ 454 ';', '--', 'drop ', '/*', '*/', 455 'union ', 'exec ', 'execute ', 456 'insert ', 'update ', 'delete ', 457 'create ', 'alter ', 'truncate ', 458 ] 459 for symbol in banned_symbols: 460 if symbol in str(substring).lower(): 461 error(f"Invalid string: '{substring}'") 462 463 464_VALID_DATEPARTS = frozenset({'year', 'month', 'day', 'hour', 'minute', 'second'}) 465 466 467def dateadd_str( 468 flavor: str = 'postgresql', 469 datepart: str = 'day', 470 number: Union[int, float] = 0, 471 begin: Union[str, datetime, int] = 'now', 472 db_type: Optional[str] = None, 473) -> str: 474 """ 475 Generate a `DATEADD` clause depending on database flavor. 476 477 Parameters 478 ---------- 479 flavor: str, default `'postgresql'` 480 SQL database flavor, e.g. `'postgresql'`, `'sqlite'`. 481 482 Currently supported flavors: 483 484 - `'postgresql'` 485 - `'postgis'` 486 - `'timescaledb'` 487 - `'timescaledb-ha'` 488 - `'citus'` 489 - `'cockroachdb'` 490 - `'duckdb'` 491 - `'mssql'` 492 - `'mysql'` 493 - `'mariadb'` 494 - `'sqlite'` 495 - `'geopackage'` 496 - `'oracle'` 497 498 datepart: str, default `'day'` 499 Which part of the date to modify. Supported values: 500 501 - `'year'` 502 - `'month'` 503 - `'day'` 504 - `'hour'` 505 - `'minute'` 506 - `'second'` 507 508 number: Union[int, float], default `0` 509 How many units to add to the date part. 510 511 begin: Union[str, datetime], default `'now'` 512 Base datetime to which to add dateparts. 513 514 db_type: Optional[str], default None 515 If provided, cast the datetime string as the type. 516 Otherwise, infer this from the input datetime value. 517 518 Returns 519 ------- 520 The appropriate `DATEADD` string for the corresponding database flavor. 521 522 Examples 523 -------- 524 >>> dateadd_str( 525 ... flavor='mssql', 526 ... begin=datetime(2022, 1, 1, 0, 0), 527 ... number=1, 528 ... ) 529 "DATEADD(day, 1, CAST('2022-01-01 00:00:00' AS DATETIME2))" 530 >>> dateadd_str( 531 ... flavor='postgresql', 532 ... begin=datetime(2022, 1, 1, 0, 0), 533 ... number=1, 534 ... ) 535 "CAST('2022-01-01 00:00:00' AS TIMESTAMP) + INTERVAL '1 day'" 536 537 """ 538 from meerschaum.utils.packages import attempt_import 539 from meerschaum.utils.dtypes.sql import get_db_type_from_pd_type, get_pd_type_from_db_type 540 dateutil_parser = attempt_import('dateutil.parser') 541 if 'int' in str(type(begin)).lower(): 542 num_str = str(begin) 543 if number is not None and number != 0: 544 num_str += ( 545 f' + {number}' 546 if number > 0 547 else f" - {number * -1}" 548 ) 549 return num_str 550 if datepart not in _VALID_DATEPARTS: 551 raise ValueError( 552 f"Invalid datepart '{datepart}'. Must be one of: {sorted(_VALID_DATEPARTS)}" 553 ) 554 if not begin: 555 return '' 556 557 _original_begin = begin 558 begin_time = None 559 ### Sanity check: make sure `begin` is a valid datetime before we inject anything. 560 if not isinstance(begin, datetime): 561 try: 562 begin_time = dateutil_parser.parse(begin) 563 except Exception: 564 begin_time = None 565 else: 566 begin_time = begin 567 568 ### Unable to parse into a datetime. 569 if begin_time is None: 570 ### Throw an error if banned symbols are included in the `begin` string. 571 clean(str(begin)) 572 ### If begin is a valid datetime, wrap it in quotes. 573 else: 574 if isinstance(begin, datetime) and begin.tzinfo is not None: 575 begin = begin.astimezone(timezone.utc) 576 begin = ( 577 f"'{begin.replace(tzinfo=None)}'" 578 if isinstance(begin, datetime) and flavor in TIMEZONE_NAIVE_FLAVORS 579 else f"'{begin}'" 580 ) 581 582 dt_is_utc = ( 583 begin_time.tzinfo is not None 584 if begin_time is not None 585 else ('+' in str(begin) or '-' in str(begin).split(':', maxsplit=1)[-1]) 586 ) 587 if db_type: 588 db_type_is_utc = 'utc' in get_pd_type_from_db_type(db_type).lower() 589 dt_is_utc = dt_is_utc or db_type_is_utc 590 db_type = db_type or get_db_type_from_pd_type( 591 ('datetime64[ns, UTC]' if dt_is_utc else 'datetime64[ns]'), 592 flavor=flavor, 593 ) 594 595 da = "" 596 if flavor in ( 597 'postgresql', 598 'postgis', 599 'timescaledb', 600 'timescaledb-ha', 601 'cockroachdb', 602 'citus', 603 ): 604 begin = ( 605 f"CAST({begin} AS {db_type})" if begin != 'now' 606 else f"CAST(NOW() AT TIME ZONE 'utc' AS {db_type})" 607 ) 608 da = begin + (f" + INTERVAL '{number} {datepart}'" if number != 0 else '') 609 610 elif flavor == 'duckdb': 611 begin = f"CAST({begin} AS {db_type})" if begin != 'now' else 'NOW()' 612 da = begin + (f" + INTERVAL '{number} {datepart}'" if number != 0 else '') 613 614 elif flavor in ('mssql',): 615 if begin_time and begin_time.microsecond != 0 and not dt_is_utc: 616 begin = begin[:-4] + "'" 617 begin = f"CAST({begin} AS {db_type})" if begin != 'now' else 'GETUTCDATE()' 618 if dt_is_utc: 619 begin += " AT TIME ZONE 'UTC'" 620 da = f"DATEADD({datepart}, {number}, {begin})" if number != 0 else begin 621 622 elif flavor in ('mysql', 'mariadb'): 623 begin = ( 624 f"CAST({begin} AS DATETIME(6))" 625 if begin != 'now' 626 else 'UTC_TIMESTAMP(6)' 627 ) 628 da = (f"DATE_ADD({begin}, INTERVAL {number} {datepart})" if number != 0 else begin) 629 630 elif flavor in ('sqlite', 'geopackage'): 631 da = f"datetime({begin}, '{number} {datepart}')" 632 633 elif flavor == 'oracle': 634 if begin == 'now': 635 begin = str( 636 datetime.now(timezone.utc).replace(tzinfo=None).strftime(r'%Y:%m:%d %M:%S.%f') 637 ) 638 elif begin_time: 639 begin = str(begin_time.strftime(r'%Y-%m-%d %H:%M:%S.%f')) 640 dt_format = 'YYYY-MM-DD HH24:MI:SS.FF' 641 _begin = f"'{begin}'" if begin_time else begin 642 da = ( 643 (f"TO_TIMESTAMP({_begin}, '{dt_format}')" if begin_time else _begin) 644 + (f" + INTERVAL '{number}' {datepart}" if number != 0 else "") 645 ) 646 return da 647 648 649def test_connection( 650 self, 651 **kw: Any 652) -> Union[bool, None]: 653 """ 654 Test if a successful connection to the database may be made. 655 656 Parameters 657 ---------- 658 **kw: 659 `max_retries` (default `1`), `retry_wait` seconds between attempts (default `0`) 660 and `warn` (print the error of a failed attempt, default `False`). 661 662 Returns 663 ------- 664 `True` if a connection is made, otherwise `False`. 665 666 """ 667 import time 668 max_retries = max(int(kw.get('max_retries', 1)), 1) 669 retry_wait = kw.get('retry_wait', 0) 670 warn = kw.get('warn', False) 671 for attempt in range(max_retries): 672 try: 673 ### The Rust connector's `SELECT 1`: no SQLAlchemy engine, so the base wheel 674 ### can check a reachable database. 675 if self._core_connector().test_connection(): 676 return True 677 except Exception as e: 678 if warn: 679 print(e) 680 if attempt + 1 < max_retries and retry_wait: 681 time.sleep(retry_wait) 682 return False 683 684 685def get_distinct_col_count( 686 col: str, 687 query: str, 688 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 689 debug: bool = False 690) -> Optional[int]: 691 """ 692 Returns the number of distinct items in a column of a SQL query. 693 694 Parameters 695 ---------- 696 col: str: 697 The column in the query to count. 698 699 query: str: 700 The SQL query to count from. 701 702 connector: Optional[mrsm.connectors.sql.SQLConnector], default None: 703 The SQLConnector to execute the query. 704 705 debug: bool, default False: 706 Verbosity toggle. 707 708 Returns 709 ------- 710 An `int` of the number of columns in the query or `None` if the query fails. 711 712 """ 713 if connector is None: 714 connector = mrsm.get_connector('sql') 715 716 _col_name = sql_item_name(col, connector.flavor, None) 717 718 _meta_query = ( 719 f""" 720 WITH src AS ( {query} ), 721 dist AS ( SELECT DISTINCT {_col_name} FROM src ) 722 SELECT COUNT(*) FROM dist""" 723 ) if connector.flavor not in ('mysql', 'mariadb') else ( 724 f""" 725 SELECT COUNT(*) 726 FROM ( 727 SELECT DISTINCT {_col_name} 728 FROM ({query}) AS src 729 ) AS dist""" 730 ) 731 732 result = connector.value(_meta_query, debug=debug, silent=True) 733 try: 734 return int(result) 735 except Exception: 736 return None 737 738 739def sql_item_name(item: str, flavor: str, schema: Optional[str] = None) -> str: 740 """ 741 Parse SQL items depending on the flavor. 742 743 Parameters 744 ---------- 745 item: str 746 The database item (table, view, etc.) in need of quotes. 747 748 flavor: str 749 The database flavor (`'postgresql'`, `'mssql'`, `'sqllite'`, etc.). 750 751 schema: Optional[str], default None 752 If provided, prefix the table name with the schema. 753 754 Returns 755 ------- 756 A `str` which contains the input `item` wrapped in the corresponding escape characters. 757 758 Examples 759 -------- 760 >>> sql_item_name('table', 'sqlite') 761 '"table"' 762 >>> sql_item_name('table', 'mssql') 763 "[table]" 764 >>> sql_item_name('table', 'postgresql', schema='abc') 765 '"abc"."table"' 766 767 """ 768 _rust = _core.sql_item_name 769 return _rust(item, flavor, schema) 770 771 772def pg_capital(s: str, quote_capitals: bool = True) -> str: 773 """ 774 If string contains a capital letter, wrap it in double quotes. 775 776 Parameters 777 ---------- 778 s: str 779 The string to be escaped. 780 781 quote_capitals: bool, default True 782 If `False`, do not quote strings with contain only a mix of capital and lower-case letters. 783 784 Returns 785 ------- 786 The input string wrapped in quotes only if it needs them. 787 788 Examples 789 -------- 790 >>> pg_capital("My Table") 791 '"My Table"' 792 >>> pg_capital('my_table') 793 'my_table' 794 795 """ 796 if s.startswith('"') and s.endswith('"'): 797 return s 798 799 s = s.replace('"', '') 800 801 needs_quotes = s.startswith('_') 802 if not needs_quotes: 803 for c in s: 804 if c == '_': 805 continue 806 807 if not c.isalnum() or (quote_capitals and c.isupper()): 808 needs_quotes = True 809 break 810 811 if needs_quotes: 812 return '"' + s + '"' 813 814 return s 815 816 817def oracle_capital(s: str) -> str: 818 """ 819 Capitalize the string of an item on an Oracle database. 820 """ 821 return s 822 823 824def truncate_item_name(item: str, flavor: str) -> str: 825 """ 826 Truncate item names to stay within the database flavor's character limit. 827 828 Parameters 829 ---------- 830 item: str 831 The database item being referenced. This string is the "canonical" name internally. 832 833 flavor: str 834 The flavor of the database on which `item` resides. 835 836 Returns 837 ------- 838 The truncated string. 839 """ 840 from meerschaum.utils.misc import truncate_string_sections 841 return truncate_string_sections( 842 item, max_len=max_name_lens.get(flavor, max_name_lens['default']) 843 ) 844 845 846def build_where( 847 params: Dict[str, Any], 848 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 849 with_where: bool = True, 850 flavor: str = 'postgresql', 851) -> str: 852 """ 853 Build the `WHERE` clause based on the input criteria. 854 855 Parameters 856 ---------- 857 params: Dict[str, Any]: 858 The keywords dictionary to convert into a WHERE clause. 859 If a value is a string which begins with an underscore, negate that value 860 (e.g. `!=` instead of `=` or `NOT IN` instead of `IN`). 861 A value of `_None` will be interpreted as `IS NOT NULL`. 862 863 connector: Optional[meerschaum.connectors.sql.SQLConnector], default None: 864 The Meerschaum SQLConnector that will be executing the query. 865 The connector is used to extract the SQL dialect. 866 867 with_where: bool, default True: 868 If `True`, include the leading `'WHERE'` string. 869 870 flavor: str, default 'postgresql' 871 If `connector` is `None`, fall back to this flavor. 872 873 Returns 874 ------- 875 A `str` of the `WHERE` clause from the input `params` dictionary for the connector's flavor. 876 877 Examples 878 -------- 879 ``` 880 >>> print(build_where({'foo': [1, 2, 3]})) 881 882 WHERE 883 "foo" IN ('1', '2', '3') 884 ``` 885 """ 886 _rust = _core.build_where 887 _query_flavor = ( 888 getattr(connector, 'flavor', flavor) if connector is not None else flavor 889 ) 890 return _rust(params, _query_flavor, with_where) 891 892 893def table_exists( 894 table: str, 895 connector: mrsm.connectors.sql.SQLConnector, 896 schema: Optional[str] = None, 897 debug: bool = False, 898) -> bool: 899 """Check if a table exists. 900 901 Parameters 902 ---------- 903 table: str: 904 The name of the table in question. 905 906 connector: mrsm.connectors.sql.SQLConnector 907 The connector to the database which holds the table. 908 909 schema: Optional[str], default None 910 Optionally specify the table schema. 911 Defaults to `connector.schema`. 912 913 debug: bool, default False : 914 Verbosity toggle. 915 916 Returns 917 ------- 918 A `bool` indicating whether or not the table exists on the database. 919 """ 920 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 921 from meerschaum.connectors.sql._native import table_exists as native_table_exists 922 schema = schema or connector.schema 923 truncated_table_name = truncate_item_name(str(table), connector.flavor) 924 ### Raw SQL through the Rust connector: no SQLAlchemy (rc15 Q53). 925 exists = native_table_exists(connector, truncated_table_name, schema) 926 if exists is not None: 927 return exists 928 sqlalchemy = import_sqlalchemy() 929 insp = sqlalchemy.inspect(connector.engine) 930 return insp.has_table(truncated_table_name, schema=schema) 931 932 933def get_sqlalchemy_table( 934 table: str, 935 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 936 schema: Optional[str] = None, 937 refresh: bool = False, 938 debug: bool = False, 939) -> Union['sqlalchemy.Table', None]: 940 """ 941 Construct a SQLAlchemy table from its name. 942 943 Parameters 944 ---------- 945 table: str 946 The name of the table on the database. Does not need to be escaped. 947 948 connector: Optional[meerschaum.connectors.sql.SQLConnector], default None: 949 The connector to the database which holds the table. 950 951 schema: Optional[str], default None 952 Specify on which schema the table resides. 953 Defaults to the schema set in `connector`. 954 955 refresh: bool, default False 956 If `True`, rebuild the cached table object. 957 958 debug: bool, default False: 959 Verbosity toggle. 960 961 Returns 962 ------- 963 A `sqlalchemy.Table` object for the table. 964 965 """ 966 if connector is None: 967 from meerschaum import get_connector 968 connector = get_connector('sql') 969 970 if connector.flavor == 'duckdb': 971 return None 972 973 from meerschaum.connectors.sql.tables import get_tables 974 from meerschaum.utils.packages import attempt_import 975 from meerschaum.utils.warnings import warn 976 if refresh: 977 connector.metadata.clear() 978 tables = get_tables(mrsm_instance=connector, debug=debug, create=False) 979 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 980 sqlalchemy = import_sqlalchemy() 981 truncated_table_name = truncate_item_name(str(table), connector.flavor) 982 table_kwargs = { 983 'autoload_with': connector.engine, 984 } 985 if schema: 986 table_kwargs['schema'] = schema 987 988 if refresh or truncated_table_name not in tables: 989 try: 990 tables[truncated_table_name] = sqlalchemy.Table( 991 truncated_table_name, 992 connector.metadata, 993 **table_kwargs 994 ) 995 except sqlalchemy.exc.NoSuchTableError: 996 warn(f"Table '{truncated_table_name}' does not exist in '{connector}'.") 997 return None 998 return tables[truncated_table_name] 999 1000 1001def _read_records(connectable: Any, query: str, debug: bool = False) -> List[Dict[str, Any]]: 1002 """The rows of `connectable.read(query)` as dictionaries (a pandas or a Polars result).""" 1003 df = connectable.read(query, debug=debug) 1004 if df is None: 1005 return [] 1006 if hasattr(df, 'to_dicts'): 1007 return df.to_dicts() 1008 return df.to_dict(orient='records') 1009 1010 1011def get_table_cols_types( 1012 table: str, 1013 connectable: Union[ 1014 'mrsm.connectors.sql.SQLConnector', 1015 'sqlalchemy.orm.session.Session', 1016 'sqlalchemy.engine.base.Engine' 1017 ], 1018 flavor: Optional[str] = None, 1019 schema: Optional[str] = None, 1020 database: Optional[str] = None, 1021 debug: bool = False, 1022) -> Dict[str, str]: 1023 """ 1024 Return a dictionary mapping a table's columns to data types. 1025 This is useful for inspecting tables creating during a not-yet-committed session. 1026 1027 NOTE: This may return incorrect columns if the schema is not explicitly stated. 1028 Use this function if you are confident the table name is unique or if you have 1029 and explicit schema. 1030 To use the configured schema, get the columns from `get_sqlalchemy_table()` instead. 1031 1032 Parameters 1033 ---------- 1034 table: str 1035 The name of the table (unquoted). 1036 1037 connectable: Union[ 1038 'mrsm.connectors.sql.SQLConnector', 1039 'sqlalchemy.orm.session.Session', 1040 'sqlalchemy.engine.base.Engine' 1041 ] 1042 The connection object used to fetch the columns and types. 1043 1044 flavor: Optional[str], default None 1045 The database dialect flavor to use for the query. 1046 If omitted, default to `connectable.flavor`. 1047 1048 schema: Optional[str], default None 1049 If provided, restrict the query to this schema. 1050 1051 database: Optional[str]. default None 1052 If provided, restrict the query to this database. 1053 1054 Returns 1055 ------- 1056 A dictionary mapping column names to data types. 1057 """ 1058 import textwrap 1059 from meerschaum.connectors import SQLConnector 1060 from meerschaum.connectors.sql._native import sql_text, uses_native 1061 flavor = flavor or getattr(connectable, 'flavor', None) 1062 if not flavor: 1063 raise ValueError("Please provide a database flavor.") 1064 if flavor == 'duckdb' and not isinstance(connectable, SQLConnector): 1065 raise ValueError("You must provide a SQLConnector when using DuckDB.") 1066 native = isinstance(connectable, SQLConnector) and uses_native(connectable) 1067 if flavor in NO_SCHEMA_FLAVORS: 1068 schema = None 1069 if schema is None: 1070 schema = DEFAULT_SCHEMA_FLAVORS.get(flavor, None) 1071 if flavor in ('sqlite', 'duckdb', 'oracle', 'geopackage'): 1072 database = None 1073 table_trunc = truncate_item_name(table, flavor=flavor) 1074 table_lower = table.lower() 1075 table_upper = table.upper() 1076 table_lower_trunc = truncate_item_name(table_lower, flavor=flavor) 1077 table_upper_trunc = truncate_item_name(table_upper, flavor=flavor) 1078 db_prefix = ( 1079 "tempdb." 1080 if flavor == 'mssql' and table.startswith('#') 1081 else "" 1082 ) 1083 1084 def _esc(s: str) -> str: 1085 return s.replace("'", "''") 1086 1087 cols_types_query = sql_text( 1088 connectable, 1089 textwrap.dedent(columns_types_queries.get( 1090 flavor, 1091 columns_types_queries['default'] 1092 ).format( 1093 table=_esc(table), 1094 table_trunc=_esc(table_trunc), 1095 table_lower=_esc(table_lower), 1096 table_lower_trunc=_esc(table_lower_trunc), 1097 table_upper=_esc(table_upper), 1098 table_upper_trunc=_esc(table_upper_trunc), 1099 db_prefix=db_prefix, 1100 )).lstrip().rstrip() 1101 ) 1102 1103 cols = ['database', 'schema', 'table', 'column', 'type', 'numeric_precision', 'numeric_scale'] 1104 result_cols_ix = dict(enumerate(cols)) 1105 1106 debug_kwargs = {'debug': debug} if isinstance(connectable, SQLConnector) else {} 1107 if not debug_kwargs and debug: 1108 dprint(cols_types_query) 1109 1110 try: 1111 result_rows = ( 1112 [ 1113 row 1114 for row in connectable.execute(cols_types_query, **debug_kwargs).fetchall() 1115 ] 1116 if flavor != 'duckdb' or native 1117 else [ 1118 tuple([doc[col] for col in cols]) 1119 for doc in _read_records(connectable, cols_types_query, debug=debug) 1120 ] 1121 ) 1122 cols_types_docs = [ 1123 { 1124 result_cols_ix[i]: val 1125 for i, val in enumerate(row) 1126 } 1127 for row in result_rows 1128 ] 1129 cols_types_docs_filtered = [ 1130 doc 1131 for doc in cols_types_docs 1132 if ( 1133 ( 1134 not schema 1135 or doc['schema'] == schema 1136 ) 1137 and 1138 ( 1139 not database 1140 or doc['database'] == database 1141 ) 1142 ) 1143 ] 1144 1145 ### NOTE: This may return incorrect columns if the schema is not explicitly stated. 1146 if cols_types_docs and not cols_types_docs_filtered: 1147 cols_types_docs_filtered = cols_types_docs 1148 1149 ### NOTE: Check for PostGIS GEOMETRY columns. 1150 geometry_cols_types = {} 1151 user_defined_cols = [ 1152 doc 1153 for doc in cols_types_docs_filtered 1154 if str(doc.get('type', None)).upper() == 'USER-DEFINED' 1155 ] 1156 if user_defined_cols: 1157 geometry_cols_types.update( 1158 get_postgis_geo_columns_types( 1159 connectable, 1160 table, 1161 schema=schema, 1162 debug=debug, 1163 ) 1164 ) 1165 1166 cols_types = { 1167 ( 1168 doc['column'] 1169 if flavor != 'oracle' else ( 1170 ( 1171 doc['column'].lower() 1172 if (doc['column'].isupper() and doc['column'].replace('_', '').isalpha()) 1173 else doc['column'] 1174 ) 1175 ) 1176 ): doc['type'].upper() + ( 1177 f'({precision},{scale})' 1178 if ( 1179 (precision := doc.get('numeric_precision', None)) 1180 and 1181 (scale := doc.get('numeric_scale', None)) 1182 ) 1183 else '' 1184 ) 1185 for doc in cols_types_docs_filtered 1186 } 1187 cols_types.update(geometry_cols_types) 1188 return cols_types 1189 except Exception as e: 1190 warn(f"Failed to fetch columns for table '{table}':\n{e}") 1191 return {} 1192 1193 1194def get_table_cols_indices( 1195 table: str, 1196 connectable: Union[ 1197 'mrsm.connectors.sql.SQLConnector', 1198 'sqlalchemy.orm.session.Session', 1199 'sqlalchemy.engine.base.Engine' 1200 ], 1201 flavor: Optional[str] = None, 1202 schema: Optional[str] = None, 1203 database: Optional[str] = None, 1204 debug: bool = False, 1205) -> Dict[str, List[str]]: 1206 """ 1207 Return a dictionary mapping a table's columns to lists of indices. 1208 This is useful for inspecting tables creating during a not-yet-committed session. 1209 1210 NOTE: This may return incorrect columns if the schema is not explicitly stated. 1211 Use this function if you are confident the table name is unique or if you have 1212 and explicit schema. 1213 To use the configured schema, get the columns from `get_sqlalchemy_table()` instead. 1214 1215 Parameters 1216 ---------- 1217 table: str 1218 The name of the table (unquoted). 1219 1220 connectable: Union[ 1221 'mrsm.connectors.sql.SQLConnector', 1222 'sqlalchemy.orm.session.Session', 1223 'sqlalchemy.engine.base.Engine' 1224 ] 1225 The connection object used to fetch the columns and types. 1226 1227 flavor: Optional[str], default None 1228 The database dialect flavor to use for the query. 1229 If omitted, default to `connectable.flavor`. 1230 1231 schema: Optional[str], default None 1232 If provided, restrict the query to this schema. 1233 1234 database: Optional[str]. default None 1235 If provided, restrict the query to this database. 1236 1237 Returns 1238 ------- 1239 A dictionary mapping column names to a list of indices. 1240 """ 1241 import textwrap 1242 from collections import defaultdict 1243 from meerschaum.connectors import SQLConnector 1244 from meerschaum.connectors.sql._native import sql_text, uses_native 1245 flavor = flavor or getattr(connectable, 'flavor', None) 1246 if not flavor: 1247 raise ValueError("Please provide a database flavor.") 1248 if flavor == 'duckdb' and not isinstance(connectable, SQLConnector): 1249 raise ValueError("You must provide a SQLConnector when using DuckDB.") 1250 native = isinstance(connectable, SQLConnector) and uses_native(connectable) 1251 if flavor in NO_SCHEMA_FLAVORS: 1252 schema = None 1253 if schema is None: 1254 schema = DEFAULT_SCHEMA_FLAVORS.get(flavor, None) 1255 if flavor in ('sqlite', 'duckdb', 'oracle', 'geopackage'): 1256 database = None 1257 table_trunc = truncate_item_name(table, flavor=flavor) 1258 table_lower = table.lower() 1259 table_upper = table.upper() 1260 table_lower_trunc = truncate_item_name(table_lower, flavor=flavor) 1261 table_upper_trunc = truncate_item_name(table_upper, flavor=flavor) 1262 db_prefix = ( 1263 "tempdb." 1264 if flavor == 'mssql' and table.startswith('#') 1265 else "" 1266 ) 1267 1268 def _esc(s: str) -> str: 1269 return s.replace("'", "''") 1270 1271 cols_indices_query = sql_text( 1272 connectable, 1273 textwrap.dedent(columns_indices_queries.get( 1274 flavor, 1275 columns_indices_queries['default'] 1276 ).format( 1277 table=_esc(table), 1278 table_trunc=_esc(table_trunc), 1279 table_lower=_esc(table_lower), 1280 table_lower_trunc=_esc(table_lower_trunc), 1281 table_upper=_esc(table_upper), 1282 table_upper_trunc=_esc(table_upper_trunc), 1283 db_prefix=db_prefix, 1284 schema=_esc(schema) if schema else schema, 1285 )).lstrip().rstrip() 1286 ) 1287 1288 cols = ['database', 'schema', 'table', 'column', 'index', 'index_type'] 1289 if flavor == 'mssql': 1290 cols.append('clustered') 1291 result_cols_ix = dict(enumerate(cols)) 1292 1293 debug_kwargs = {'debug': debug} if isinstance(connectable, SQLConnector) else {} 1294 if not debug_kwargs and debug: 1295 dprint(cols_indices_query) 1296 1297 try: 1298 result_rows = ( 1299 [ 1300 row 1301 for row in connectable.execute(cols_indices_query, **debug_kwargs).fetchall() 1302 ] 1303 if flavor != 'duckdb' or native 1304 else [ 1305 tuple([doc[col] for col in cols]) 1306 for doc in _read_records(connectable, cols_indices_query, debug=debug) 1307 ] 1308 ) 1309 cols_types_docs = [ 1310 { 1311 result_cols_ix[i]: val 1312 for i, val in enumerate(row) 1313 } 1314 for row in result_rows 1315 ] 1316 cols_types_docs_filtered = [ 1317 doc 1318 for doc in cols_types_docs 1319 if ( 1320 ( 1321 not schema 1322 or doc['schema'] == schema 1323 ) 1324 and 1325 ( 1326 not database 1327 or doc['database'] == database 1328 ) 1329 ) 1330 ] 1331 ### NOTE: This may return incorrect columns if the schema is not explicitly stated. 1332 if cols_types_docs and not cols_types_docs_filtered: 1333 cols_types_docs_filtered = cols_types_docs 1334 1335 cols_indices = defaultdict(lambda: []) 1336 for doc in cols_types_docs_filtered: 1337 col = ( 1338 doc['column'] 1339 if flavor != 'oracle' 1340 else ( 1341 doc['column'].lower() 1342 if (doc['column'].isupper() and doc['column'].replace('_', '').isalpha()) 1343 else doc['column'] 1344 ) 1345 ) 1346 index_doc = { 1347 'name': doc.get('index', None), 1348 'type': doc.get('index_type', None) 1349 } 1350 if flavor == 'mssql': 1351 index_doc['clustered'] = doc.get('clustered', None) 1352 cols_indices[col].append(index_doc) 1353 1354 return dict(cols_indices) 1355 except Exception as e: 1356 warn(f"Failed to fetch columns for table '{table}':\n{e}") 1357 return {} 1358 1359 1360def get_db_version(conn: 'SQLConnector', debug: bool = False) -> Union[str, None]: 1361 """ 1362 Fetch the database version if possible. 1363 """ 1364 version_name = sql_item_name('version', conn.flavor, None) 1365 version_query = version_queries.get( 1366 conn.flavor, 1367 version_queries['default'] 1368 ).format(version_name=version_name) 1369 return conn.value(version_query, debug=debug, silent=True) 1370 1371 1372### Non-TimescaleDB flavors that support declarative range partitioning on a datetime column. 1373### PostgreSQL-style declares an empty parent (`PARTITION BY RANGE`) with children added 1374### separately; MySQL-style requires the initial partitions inline at `CREATE TABLE`. 1375PG_PARTITION_FLAVORS = {'postgresql', 'postgis'} 1376MYSQL_PARTITION_FLAVORS = {'mysql', 'mariadb'} 1377def find_top_level_select_index(query: str, last: bool = False) -> int: 1378 """ 1379 Return the index of the first (or last) top-level `SELECT` keyword in a query. 1380 1381 The scan skips `--` line comments, `/* */` block comments, string literals, 1382 quoted identifiers, and parenthesized subqueries, so a definition may safely 1383 contain comments (including comments mentioning "select"). 1384 1385 Parameters 1386 ---------- 1387 query: str 1388 The SQL query to scan. 1389 1390 last: bool, default False 1391 If `True`, return the index of the last top-level `SELECT` 1392 (e.g. the final statement of a CTE query) instead of the first. 1393 1394 Returns 1395 ------- 1396 The character index of the `SELECT` keyword, or `-1` if none exists 1397 at the top level. 1398 """ 1399 q_lower = query.lower() 1400 n = len(q_lower) 1401 depth = 0 1402 found_ix = -1 1403 i = 0 1404 while i < n: 1405 ch = q_lower[i] 1406 nxt = q_lower[i + 1] if (i + 1) < n else '' 1407 if ch == '-' and nxt == '-': 1408 eol = q_lower.find('\n', i) 1409 i = n if eol == -1 else (eol + 1) 1410 continue 1411 if ch == '/' and nxt == '*': 1412 end = q_lower.find('*/', i + 2) 1413 i = n if end == -1 else (end + 2) 1414 continue 1415 if ch == "'": 1416 i += 1 1417 while i < n: 1418 if q_lower[i] == "'": 1419 if (i + 1) < n and q_lower[i + 1] == "'": 1420 i += 2 1421 continue 1422 break 1423 i += 1 1424 i += 1 1425 continue 1426 if ch in ('"', '`'): 1427 end = q_lower.find(ch, i + 1) 1428 i = n if end == -1 else (end + 1) 1429 continue 1430 if ch == '[': 1431 end = q_lower.find(']', i + 1) 1432 i = n if end == -1 else (end + 1) 1433 continue 1434 if ch == '(': 1435 depth += 1 1436 i += 1 1437 continue 1438 if ch == ')': 1439 depth -= 1 1440 i += 1 1441 continue 1442 if depth == 0 and ch == 's' and q_lower[i:i + 6] == 'select': 1443 before_ok = i == 0 or not (q_lower[i - 1].isalnum() or q_lower[i - 1] == '_') 1444 after_ok = (i + 6) == n or not (q_lower[i + 6].isalnum() or q_lower[i + 6] == '_') 1445 if before_ok and after_ok: 1446 if not last: 1447 return i 1448 found_ix = i 1449 i += 6 1450 continue 1451 i += 1 1452 return found_ix 1453 1454 1455def get_leading_with_index(query: str) -> int: 1456 """ 1457 Return the index of a query's leading `WITH` keyword, 1458 skipping whitespace and comments, or `-1` if the query 1459 does not begin with `WITH`. 1460 """ 1461 q_lower = query.lower() 1462 n = len(q_lower) 1463 i = 0 1464 while i < n: 1465 ch = q_lower[i] 1466 nxt = q_lower[i + 1] if (i + 1) < n else '' 1467 if ch in (' ', '\t', '\r', '\n'): 1468 i += 1 1469 continue 1470 if ch == '-' and nxt == '-': 1471 eol = q_lower.find('\n', i) 1472 i = n if eol == -1 else (eol + 1) 1473 continue 1474 if ch == '/' and nxt == '*': 1475 end = q_lower.find('*/', i + 2) 1476 i = n if end == -1 else (end + 2) 1477 continue 1478 break 1479 if q_lower[i:i + 4] == 'with' and ( 1480 (i + 4) == n or not (q_lower[i + 4].isalnum() or q_lower[i + 4] == '_') 1481 ): 1482 return i 1483 return -1 1484 1485 1486def wrap_query_with_cte( 1487 sub_query: str, 1488 parent_query: str, 1489 flavor: str, 1490 cte_name: str = "src", 1491) -> str: 1492 """ 1493 Wrap a subquery in a CTE and append an encapsulating query. 1494 1495 Parameters 1496 ---------- 1497 sub_query: str 1498 The query to be referenced. This may itself contain CTEs and comments. 1499 Unless `cte_name` is provided, this will be aliased as `src`. 1500 1501 parent_query: str 1502 The larger query to append which references the subquery. 1503 If it contains CTEs of its own, they are hoisted into the 1504 enclosing `WITH` list. 1505 1506 flavor: str 1507 The database flavor, e.g. `'mssql'`. 1508 1509 cte_name: str, default 'src' 1510 The CTE alias, defaults to `src`. 1511 1512 Returns 1513 ------- 1514 An encapsulating query which allows you to treat `sub_query` as a temporary table. 1515 Raises an exception when a query's top-level `SELECT` cannot be located 1516 (e.g. an incomplete definition), rather than emitting invalid SQL. 1517 1518 Examples 1519 -------- 1520 1521 >>> from meerschaum.utils.sql import wrap_query_with_cte 1522 >>> sub_query = "WITH foo AS (SELECT 1 AS val) SELECT (val * 2) AS newval FROM foo" 1523 >>> parent_query = "SELECT newval * 3 FROM src" 1524 >>> query = wrap_query_with_cte(sub_query, parent_query, 'mssql') 1525 >>> print(query) 1526 >>> # WITH foo AS (SELECT 1 AS val) 1527 >>> # , [src] AS ( 1528 >>> # SELECT (val * 2) AS newval FROM foo 1529 >>> # ) 1530 >>> # SELECT newval * 3 FROM src 1531 1532 """ 1533 import textwrap 1534 from meerschaum.utils.warnings import error 1535 sub_query = sub_query.lstrip() 1536 parent_query = parent_query.lstrip() 1537 cte_name_quoted = sql_item_name(cte_name, flavor, None) 1538 1539 if flavor in NO_CTE_FLAVORS: 1540 return ( 1541 parent_query 1542 .replace(cte_name_quoted, '--MRSM_SUBQUERY--') 1543 .replace(cte_name, '--MRSM_SUBQUERY--') 1544 .replace('--MRSM_SUBQUERY--', f"(\n{sub_query}\n) AS {cte_name_quoted}") 1545 ) 1546 1547 if get_leading_with_index(sub_query) != -1: 1548 select_index = find_top_level_select_index(sub_query) 1549 if select_index == -1: 1550 error( 1551 "Could not locate the top-level SELECT of the following query:\n" 1552 + sub_query 1553 ) 1554 1555 ### Flatten: the definitions part keeps the sub-query's own CTEs 1556 ### (`WITH a AS (...) -- trailing comments included`), and the body 1557 ### (`SELECT ...`) becomes the new CTE. The joining comma goes on its 1558 ### own line so a trailing `--` comment cannot swallow it. 1559 definitions = sub_query[:select_index].rstrip() 1560 if definitions.endswith(','): 1561 definitions = definitions[:-1].rstrip() 1562 body = sub_query[select_index:].strip() 1563 1564 lead = ( 1565 f"{definitions}\n" 1566 f", {cte_name_quoted} AS (\n" 1567 f"{textwrap.indent(body, ' ')}\n" 1568 f")" 1569 ) 1570 else: 1571 lead = ( 1572 f"WITH {cte_name_quoted} AS (\n" 1573 f"{textwrap.indent(sub_query, ' ')}\n" 1574 f")" 1575 ) 1576 1577 parent_with_index = get_leading_with_index(parent_query) 1578 if parent_with_index == -1: 1579 return f"{lead}\n{parent_query}" 1580 1581 ### The parent query declares CTEs of its own: hoist them after the new 1582 ### definition instead of emitting an invalid second `WITH`. 1583 parent_prefix = parent_query[:parent_with_index].rstrip() 1584 parent_rest = parent_query[parent_with_index + len('with'):].lstrip() 1585 return ( 1586 (f"{parent_prefix}\n" if parent_prefix else "") 1587 + f"{lead}\n, {parent_rest}" 1588 ) 1589 1590 1591def format_cte_subquery( 1592 sub_query: str, 1593 flavor: str, 1594 sub_name: str = 'src', 1595 cols_to_select: Union[List[str], str] = '*', 1596) -> str: 1597 """ 1598 Given a subquery, build a wrapper query that selects from the CTE subquery. 1599 1600 Parameters 1601 ---------- 1602 sub_query: str 1603 The subquery to wrap. 1604 1605 flavor: str 1606 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`. 1607 1608 sub_name: str, default 'src' 1609 If possible, give this name to the CTE (must be unquoted). 1610 1611 cols_to_select: Union[List[str], str], default '' 1612 If specified, choose which columns to select from the CTE. 1613 If a list of strings is provided, each item will be quoted and joined with commas. 1614 If a string is given, assume it is quoted and insert it into the query. 1615 1616 Returns 1617 ------- 1618 A wrapper query that selects from the CTE. 1619 """ 1620 quoted_sub_name = sql_item_name(sub_name, flavor, None) 1621 cols_str = ( 1622 cols_to_select 1623 if isinstance(cols_to_select, str) 1624 else ', '.join([sql_item_name(col, flavor, None) for col in cols_to_select]) 1625 ) 1626 parent_query = ( 1627 f"SELECT {cols_str}\n" 1628 f"FROM {quoted_sub_name}" 1629 ) 1630 return wrap_query_with_cte(sub_query, parent_query, flavor, cte_name=sub_name) 1631 1632 1633def session_execute( 1634 session: 'sqlalchemy.orm.session.Session', 1635 queries: Union[List[str], str], 1636 with_results: bool = False, 1637 debug: bool = False, 1638) -> Union[mrsm.SuccessTuple, Tuple[mrsm.SuccessTuple, List['sqlalchemy.sql.ResultProxy']]]: 1639 """ 1640 Similar to `SQLConnector.exec_queries()`, execute a list of queries 1641 and roll back when one fails. 1642 1643 Parameters 1644 ---------- 1645 session: sqlalchemy.orm.session.Session 1646 A SQLAlchemy session representing a transaction. 1647 1648 queries: Union[List[str], str] 1649 A query or list of queries to be executed. 1650 If a query fails, roll back the session. 1651 1652 with_results: bool, default False 1653 If `True`, return a list of result objects. 1654 1655 Returns 1656 ------- 1657 A `SuccessTuple` indicating the queries were successfully executed. 1658 If `with_results`, return the `SuccessTuple` and a list of results. 1659 """ 1660 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 1661 sqlalchemy = import_sqlalchemy() 1662 if not isinstance(queries, list): 1663 queries = [queries] 1664 successes, msgs, results = [], [], [] 1665 for query in queries: 1666 if debug: 1667 dprint(query) 1668 query_text = sqlalchemy.text(query) 1669 fail_msg = "Failed to execute queries." 1670 try: 1671 result = session.execute(query_text) 1672 query_success = result is not None 1673 query_msg = "Success" if query_success else fail_msg 1674 except Exception as e: 1675 query_success = False 1676 query_msg = f"{fail_msg}\n{e}" 1677 result = None 1678 successes.append(query_success) 1679 msgs.append(query_msg) 1680 results.append(result) 1681 if not query_success: 1682 if debug: 1683 dprint("Rolling back session.") 1684 session.rollback() 1685 break 1686 success, msg = all(successes), '\n'.join(msgs) 1687 if with_results: 1688 return (success, msg), results 1689 return success, msg 1690 1691 1692def get_postgis_geo_columns_types( 1693 connectable: Union[ 1694 'mrsm.connectors.sql.SQLConnector', 1695 'sqlalchemy.orm.session.Session', 1696 'sqlalchemy.engine.base.Engine' 1697 ], 1698 table: str, 1699 schema: Optional[str] = 'public', 1700 debug: bool = False, 1701) -> Dict[str, str]: 1702 """ 1703 Return a dictionary mapping PostGIS geometry column names to geometry types. 1704 1705 Parameters 1706 ---------- 1707 connectable: Union[SQLConnector, sqlalchemy.orm.session.Session, sqlalchemy.engine.base.Engine] 1708 The SQLConnector, Session, or Engine. 1709 1710 table: str 1711 The table's name. 1712 1713 schema: Optional[str], default 'public' 1714 The schema to use for the fully qualified table name. 1715 1716 Returns 1717 ------- 1718 A dictionary mapping column names to geometry types. 1719 """ 1720 from meerschaum.utils.dtypes import get_geometry_type_srid 1721 from meerschaum.connectors.sql._native import sql_text 1722 default_type, default_srid = get_geometry_type_srid() 1723 default_type = default_type.upper() 1724 1725 clean(table) 1726 clean(str(schema)) 1727 schema = schema or 'public' 1728 truncated_schema_name = truncate_item_name(schema, flavor='postgis') 1729 truncated_table_name = truncate_item_name(table, flavor='postgis') 1730 query = sql_text( 1731 connectable, 1732 "SELECT \"f_geometry_column\" AS \"column\", 'GEOMETRY' AS \"func\", \"type\", geom.\"srid\"\n" 1733 "FROM \"geometry_columns\" AS geom\n" 1734 f"WHERE \"f_table_schema\" = '{truncated_schema_name}'\n" 1735 f" AND \"f_table_name\" = '{truncated_table_name}'\n" 1736 "UNION ALL\n" 1737 "SELECT \"f_geography_column\" AS \"column\", 'GEOGRAPHY' AS \"func\", \"type\", geog.\"srid\"\n" 1738 "FROM \"geography_columns\" AS geog\n" 1739 f"WHERE \"f_table_schema\" = '{truncated_schema_name}'\n" 1740 f" AND \"f_table_name\" = '{truncated_table_name}'\n" 1741 ) 1742 debug_kwargs = {'debug': debug} if isinstance(connectable, mrsm.connectors.SQLConnector) else {} 1743 result_rows = [ 1744 row 1745 for row in connectable.execute(query, **debug_kwargs).fetchall() 1746 ] 1747 cols_type_tuples = { 1748 row[0]: ( 1749 row[1], # func 1750 row[2], # type 1751 ( 1752 row[3] 1753 if row[3] 1754 else 0 1755 ) # srid 1756 ) 1757 for row in result_rows 1758 } 1759 1760 geometry_cols_types = { 1761 col: ( 1762 f"{func}({typ.upper()}, {srid})" 1763 if srid 1764 else ( 1765 func 1766 + ( 1767 f'({typ.upper()})' 1768 if typ.upper() not in ('GEOMETRY', 'GEOGRAPHY') 1769 else '' 1770 ) 1771 ) 1772 ) 1773 for col, (func, typ, srid) in cols_type_tuples.items() 1774 } 1775 return geometry_cols_types 1776 1777 1778UPDATE_QUERIES = { 1779 'default': """ 1780 UPDATE {target_table_name} AS f 1781 {sets_subquery_none} 1782 FROM {target_table_name} AS t 1783 INNER JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p 1784 ON 1785 {and_subquery_t} 1786 WHERE 1787 {and_subquery_f} 1788 AND 1789 {date_bounds_subquery} 1790 """, 1791 'timescaledb-upsert': """ 1792 INSERT INTO {target_table_name} ({patch_cols_str}) 1793 SELECT {patch_cols_str} 1794 FROM {patch_table_name} 1795 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1796 """, 1797 'timescaledb-ha-upsert': """ 1798 INSERT INTO {target_table_name} ({patch_cols_str}) 1799 SELECT {patch_cols_str} 1800 FROM {patch_table_name} 1801 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1802 """, 1803 'postgresql-upsert': """ 1804 INSERT INTO {target_table_name} ({patch_cols_str}) 1805 SELECT {patch_cols_str} 1806 FROM {patch_table_name} 1807 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1808 """, 1809 'postgis-upsert': """ 1810 INSERT INTO {target_table_name} ({patch_cols_str}) 1811 SELECT {patch_cols_str} 1812 FROM {patch_table_name} 1813 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1814 """, 1815 'citus-upsert': """ 1816 INSERT INTO {target_table_name} ({patch_cols_str}) 1817 SELECT {patch_cols_str} 1818 FROM {patch_table_name} 1819 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1820 """, 1821 'cockroachdb-upsert': """ 1822 INSERT INTO {target_table_name} ({patch_cols_str}) 1823 SELECT {patch_cols_str} 1824 FROM {patch_table_name} 1825 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1826 """, 1827 'mysql': """ 1828 UPDATE {target_table_name} AS f 1829 JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p 1830 ON 1831 {and_subquery_f} 1832 {sets_subquery_f} 1833 WHERE 1834 {date_bounds_subquery} 1835 """, 1836 'mysql-upsert': """ 1837 INSERT {ignore}INTO {target_table_name} ({patch_cols_str}) 1838 SELECT {patch_cols_str} 1839 FROM {patch_table_name} 1840 {on_duplicate_key_update} 1841 {cols_equal_values} 1842 """, 1843 'mariadb': """ 1844 UPDATE {target_table_name} AS f 1845 JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p 1846 ON 1847 {and_subquery_f} 1848 {sets_subquery_f} 1849 WHERE 1850 {date_bounds_subquery} 1851 """, 1852 'mariadb-upsert': """ 1853 INSERT {ignore}INTO {target_table_name} ({patch_cols_str}) 1854 SELECT {patch_cols_str} 1855 FROM {patch_table_name} 1856 {on_duplicate_key_update} 1857 {cols_equal_values} 1858 """, 1859 'mssql': """ 1860 {with_temp_date_bounds} 1861 MERGE {target_table_name} f 1862 USING (SELECT {patch_cols_str} FROM {patch_table_name}) p 1863 ON 1864 {and_subquery_f} 1865 AND 1866 {date_bounds_subquery} 1867 WHEN MATCHED THEN 1868 UPDATE 1869 {sets_subquery_none}; 1870 """, 1871 'mssql-upsert': [ 1872 "{identity_insert_on}", 1873 """ 1874 {with_temp_date_bounds} 1875 MERGE {target_table_name} f 1876 USING (SELECT {patch_cols_str} FROM {patch_table_name}) p 1877 ON 1878 {and_subquery_f} 1879 AND 1880 {date_bounds_subquery}{when_matched_update_sets_subquery_none} 1881 WHEN NOT MATCHED THEN 1882 INSERT ({patch_cols_str}) 1883 VALUES ({patch_cols_prefixed_str}); 1884 """, 1885 "{identity_insert_off}", 1886 ], 1887 'oracle': """ 1888 MERGE INTO {target_table_name} f 1889 USING (SELECT {patch_cols_str} FROM {patch_table_name}) p 1890 ON ( 1891 {and_subquery_f} 1892 AND 1893 {date_bounds_subquery} 1894 ) 1895 WHEN MATCHED THEN 1896 UPDATE 1897 {sets_subquery_none} 1898 """, 1899 'oracle-upsert': """ 1900 MERGE INTO {target_table_name} f 1901 USING (SELECT {patch_cols_str} FROM {patch_table_name}) p 1902 ON ( 1903 {and_subquery_f} 1904 AND 1905 {date_bounds_subquery} 1906 ){when_matched_update_sets_subquery_none} 1907 WHEN NOT MATCHED THEN 1908 INSERT ({patch_cols_str}) 1909 VALUES ({patch_cols_prefixed_str}) 1910 """, 1911 'sqlite-upsert': """ 1912 INSERT INTO {target_table_name} ({patch_cols_str}) 1913 SELECT {patch_cols_str} 1914 FROM {patch_table_name} 1915 WHERE true 1916 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1917 """, 1918 'sqlite_delete_insert': [ 1919 """ 1920 DELETE FROM {target_table_name} AS f 1921 WHERE ROWID IN ( 1922 SELECT t.ROWID 1923 FROM {target_table_name} AS t 1924 INNER JOIN (SELECT * FROM {patch_table_name}) AS p 1925 ON {and_subquery_t} 1926 ); 1927 """, 1928 """ 1929 INSERT INTO {target_table_name} AS f 1930 SELECT {patch_cols_str} FROM {patch_table_name} AS p 1931 """, 1932 ], 1933 'geopackage-upsert': """ 1934 INSERT INTO {target_table_name} ({patch_cols_str}) 1935 SELECT {patch_cols_str} 1936 FROM {patch_table_name} 1937 WHERE true 1938 ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded} 1939 """, 1940} 1941 1942 1943 1944def get_update_queries( 1945 target: str, 1946 patch: str, 1947 connectable: Union[ 1948 mrsm.connectors.sql.SQLConnector, 1949 'sqlalchemy.orm.session.Session' 1950 ], 1951 join_cols: Iterable[str], 1952 flavor: Optional[str] = None, 1953 upsert: bool = False, 1954 datetime_col: Optional[str] = None, 1955 schema: Optional[str] = None, 1956 patch_schema: Optional[str] = None, 1957 target_cols_types: Optional[Dict[str, str]] = None, 1958 patch_cols_types: Optional[Dict[str, str]] = None, 1959 identity_insert: bool = False, 1960 null_indices: bool = True, 1961 cast_columns: bool = True, 1962 debug: bool = False, 1963) -> List[str]: 1964 """ 1965 Build a list of `MERGE`, `UPDATE`, `DELETE`/`INSERT` queries to apply a patch to target table. 1966 1967 Parameters 1968 ---------- 1969 target: str 1970 The name of the target table. 1971 1972 patch: str 1973 The name of the patch table. This should have the same shape as the target. 1974 1975 connectable: Union[meerschaum.connectors.sql.SQLConnector, sqlalchemy.orm.session.Session] 1976 The `SQLConnector` or SQLAlchemy session which will later execute the queries. 1977 1978 join_cols: List[str] 1979 The columns to use to join the patch to the target. 1980 1981 flavor: Optional[str], default None 1982 If using a SQLAlchemy session, provide the expected database flavor. 1983 1984 upsert: bool, default False 1985 If `True`, return an upsert query rather than an update. 1986 1987 datetime_col: Optional[str], default None 1988 If provided, bound the join query using this column as the datetime index. 1989 This must be present on both tables. 1990 1991 schema: Optional[str], default None 1992 If provided, use this schema when quoting the target table. 1993 Defaults to `connector.schema`. 1994 1995 patch_schema: Optional[str], default None 1996 If provided, use this schema when quoting the patch table. 1997 Defaults to `schema`. 1998 1999 target_cols_types: Optional[Dict[str, Any]], default None 2000 If provided, use these as the columns-types dictionary for the target table. 2001 Default will infer from the database context. 2002 2003 patch_cols_types: Optional[Dict[str, Any]], default None 2004 If provided, use these as the columns-types dictionary for the target table. 2005 Default will infer from the database context. 2006 2007 identity_insert: bool, default False 2008 If `True`, include `SET IDENTITY_INSERT` queries before and after the update queries. 2009 Only applies for MSSQL upserts. 2010 2011 null_indices: bool, default True 2012 If `False`, do not coalesce index columns before joining. 2013 2014 cast_columns: bool, default True 2015 If `False`, do not cast update columns to the target table types. 2016 2017 debug: bool, default False 2018 Verbosity toggle. 2019 2020 Returns 2021 ------- 2022 A list of query strings to perform the update operation. 2023 """ 2024 import textwrap 2025 from meerschaum.connectors import SQLConnector 2026 from meerschaum.utils.debug import dprint 2027 from meerschaum.utils.dtypes import are_dtypes_equal 2028 from meerschaum.utils.dtypes.sql import DB_FLAVORS_CAST_DTYPES, get_pd_type_from_db_type 2029 flavor = flavor or getattr(connectable, 'flavor', None) 2030 if not flavor: 2031 raise ValueError("Provide a flavor if using a SQLAlchemy session.") 2032 if ( 2033 flavor in ('sqlite', 'geopackage') 2034 and isinstance(connectable, SQLConnector) 2035 and connectable.db_version < '3.33.0' 2036 ): 2037 flavor = 'sqlite_delete_insert' 2038 flavor_key = (f'{flavor}-upsert' if upsert else flavor) 2039 base_queries = UPDATE_QUERIES.get( 2040 flavor_key, 2041 UPDATE_QUERIES['default'] 2042 ) 2043 if not isinstance(base_queries, list): 2044 base_queries = [base_queries] 2045 schema = schema or (connectable.schema if isinstance(connectable, SQLConnector) else None) 2046 patch_schema = patch_schema or schema 2047 target_table_columns = get_table_cols_types( 2048 target, 2049 connectable, 2050 flavor=flavor, 2051 schema=schema, 2052 debug=debug, 2053 ) if not target_cols_types else target_cols_types 2054 patch_table_columns = get_table_cols_types( 2055 patch, 2056 connectable, 2057 flavor=flavor, 2058 schema=patch_schema, 2059 debug=debug, 2060 ) if not patch_cols_types else patch_cols_types 2061 2062 patch_cols_str = ', '.join( 2063 [ 2064 sql_item_name(col, flavor) 2065 for col in patch_table_columns 2066 ] 2067 ) 2068 patch_cols_prefixed_str = ', '.join( 2069 [ 2070 'p.' + sql_item_name(col, flavor) 2071 for col in patch_table_columns 2072 ] 2073 ) 2074 2075 join_cols_str = ', '.join( 2076 [ 2077 sql_item_name(col, flavor) 2078 for col in join_cols 2079 ] 2080 ) 2081 2082 value_cols = [] 2083 join_cols_types = [] 2084 if debug: 2085 dprint("target_table_columns:") 2086 mrsm.pprint(target_table_columns) 2087 for c_name, c_type in target_table_columns.items(): 2088 if c_name not in patch_table_columns: 2089 continue 2090 if flavor in DB_FLAVORS_CAST_DTYPES: 2091 c_type = DB_FLAVORS_CAST_DTYPES[flavor].get(c_type.upper(), c_type) 2092 ( 2093 join_cols_types 2094 if c_name in join_cols 2095 else value_cols 2096 ).append((c_name, c_type)) 2097 if debug: 2098 dprint(f"value_cols: {value_cols}") 2099 2100 if not join_cols_types: 2101 return [] 2102 if not value_cols and not upsert: 2103 return [] 2104 2105 coalesce_join_cols_str = ', '.join( 2106 [ 2107 ( 2108 ( 2109 'COALESCE(' 2110 + sql_item_name(c_name, flavor) 2111 + ', ' 2112 + get_null_replacement(c_type, flavor) 2113 + ')' 2114 ) 2115 if null_indices 2116 else sql_item_name(c_name, flavor) 2117 ) 2118 for c_name, c_type in join_cols_types 2119 ] 2120 ) 2121 2122 update_or_nothing = ('UPDATE' if value_cols else 'NOTHING') 2123 2124 def sets_subquery(l_prefix: str, r_prefix: str): 2125 if not value_cols: 2126 return '' 2127 2128 utc_value_cols = { 2129 c_name 2130 for c_name, c_type in value_cols 2131 if ('utc' in get_pd_type_from_db_type(c_type).lower()) 2132 } if flavor not in TIMEZONE_NAIVE_FLAVORS else set() 2133 2134 cast_func_cols = { 2135 c_name: ( 2136 ('', '', '') 2137 if not cast_columns or ( 2138 flavor == 'oracle' 2139 and are_dtypes_equal(get_pd_type_from_db_type(c_type), 'bytes') 2140 ) 2141 else ( 2142 ('CAST(', f" AS {c_type.replace('_', ' ')}", ')' + ( 2143 " AT TIME ZONE 'UTC'" 2144 if c_name in utc_value_cols 2145 else '' 2146 )) 2147 if flavor not in ('sqlite', 'geopackage') 2148 else ('', '', '') 2149 ) 2150 ) 2151 for c_name, c_type in value_cols 2152 } 2153 return 'SET ' + ',\n'.join([ 2154 ( 2155 l_prefix + sql_item_name(c_name, flavor, None) 2156 + ' = ' 2157 + cast_func_cols[c_name][0] 2158 + r_prefix + sql_item_name(c_name, flavor, None) 2159 + cast_func_cols[c_name][1] 2160 + cast_func_cols[c_name][2] 2161 ) for c_name, c_type in value_cols 2162 ]) 2163 2164 def and_subquery(l_prefix: str, r_prefix: str): 2165 return '\n AND\n '.join([ 2166 ( 2167 ( 2168 "COALESCE(" 2169 + l_prefix 2170 + sql_item_name(c_name, flavor, None) 2171 + ", " 2172 + get_null_replacement(c_type, flavor) 2173 + ")" 2174 + '\n =\n ' 2175 + "COALESCE(" 2176 + r_prefix 2177 + sql_item_name(c_name, flavor, None) 2178 + ", " 2179 + get_null_replacement(c_type, flavor) 2180 + ")" 2181 ) 2182 if null_indices 2183 else ( 2184 l_prefix 2185 + sql_item_name(c_name, flavor, None) 2186 + ' = ' 2187 + r_prefix 2188 + sql_item_name(c_name, flavor, None) 2189 ) 2190 ) for c_name, c_type in join_cols_types 2191 ]) 2192 2193 skip_query_val = "" 2194 target_table_name = sql_item_name(target, flavor, schema) 2195 patch_table_name = sql_item_name(patch, flavor, patch_schema) 2196 dt_col_name = sql_item_name(datetime_col, flavor, None) if datetime_col else None 2197 date_bounds_table = patch_table_name if flavor != 'mssql' else '[date_bounds]' 2198 min_dt_col_name = f"MIN({dt_col_name})" if flavor != 'mssql' else '[Min_dt]' 2199 max_dt_col_name = f"MAX({dt_col_name})" if flavor != 'mssql' else '[Max_dt]' 2200 date_bounds_subquery = ( 2201 f"""f.{dt_col_name} >= (SELECT {min_dt_col_name} FROM {date_bounds_table}) 2202 AND 2203 f.{dt_col_name} <= (SELECT {max_dt_col_name} FROM {date_bounds_table})""" 2204 if datetime_col 2205 else "1 = 1" 2206 ) 2207 with_temp_date_bounds = f"""WITH [date_bounds] AS ( 2208 SELECT MIN({dt_col_name}) AS {min_dt_col_name}, MAX({dt_col_name}) AS {max_dt_col_name} 2209 FROM {patch_table_name} 2210 )""" if datetime_col else "" 2211 identity_insert_on = ( 2212 f"SET IDENTITY_INSERT {target_table_name} ON" 2213 if identity_insert 2214 else skip_query_val 2215 ) 2216 identity_insert_off = ( 2217 f"SET IDENTITY_INSERT {target_table_name} OFF" 2218 if identity_insert 2219 else skip_query_val 2220 ) 2221 2222 ### NOTE: MSSQL upserts must exclude the update portion if only upserting indices. 2223 when_matched_update_sets_subquery_none = "" if not value_cols else ( 2224 "\n WHEN MATCHED THEN\n" 2225 f" UPDATE {sets_subquery('', 'p.')}" 2226 ) 2227 2228 cols_equal_values = '\n,'.join( 2229 [ 2230 f"{sql_item_name(c_name, flavor)} = VALUES({sql_item_name(c_name, flavor)})" 2231 for c_name, c_type in value_cols 2232 ] 2233 ) 2234 on_duplicate_key_update = ( 2235 "ON DUPLICATE KEY UPDATE" 2236 if value_cols 2237 else "" 2238 ) 2239 ignore = "IGNORE " if not value_cols else "" 2240 2241 formatted_queries = [ 2242 textwrap.dedent(base_query.format( 2243 sets_subquery_none=sets_subquery('', 'p.'), 2244 sets_subquery_none_excluded=sets_subquery('', 'EXCLUDED.'), 2245 sets_subquery_f=sets_subquery('f.', 'p.'), 2246 and_subquery_f=and_subquery('p.', 'f.'), 2247 and_subquery_t=and_subquery('p.', 't.'), 2248 target_table_name=target_table_name, 2249 patch_table_name=patch_table_name, 2250 patch_cols_str=patch_cols_str, 2251 patch_cols_prefixed_str=patch_cols_prefixed_str, 2252 date_bounds_subquery=date_bounds_subquery, 2253 join_cols_str=join_cols_str, 2254 coalesce_join_cols_str=coalesce_join_cols_str, 2255 update_or_nothing=update_or_nothing, 2256 when_matched_update_sets_subquery_none=when_matched_update_sets_subquery_none, 2257 cols_equal_values=cols_equal_values, 2258 on_duplicate_key_update=on_duplicate_key_update, 2259 ignore=ignore, 2260 with_temp_date_bounds=with_temp_date_bounds, 2261 identity_insert_on=identity_insert_on, 2262 identity_insert_off=identity_insert_off, 2263 )).lstrip().rstrip() 2264 for base_query in base_queries 2265 ] 2266 2267 ### NOTE: Allow for skipping some queries. 2268 return [query for query in formatted_queries if query] 2269 2270 2271 2272def get_null_replacement(typ: str, flavor: str) -> str: 2273 """ 2274 Return a value that may temporarily be used in place of NULL for this type. 2275 2276 Parameters 2277 ---------- 2278 typ: str 2279 The typ to be converted to NULL. 2280 2281 flavor: str 2282 The database flavor for which this value will be used. 2283 2284 Returns 2285 ------- 2286 A value which may stand in place of NULL for this type. 2287 `'None'` is returned if a value cannot be determined. 2288 """ 2289 from meerschaum.utils.dtypes.sql import DB_FLAVORS_CAST_DTYPES 2290 if 'geometry' in typ.lower(): 2291 return "'010100000000008058346FCDC100008058346FCDC1'" 2292 if 'int' in typ.lower() or typ.lower() in ('numeric', 'number'): 2293 return '-987654321' 2294 if 'bool' in typ.lower() or typ.lower() == 'bit': 2295 bool_typ = ( 2296 PD_TO_DB_DTYPES_FLAVORS 2297 .get('bool', {}) 2298 .get(flavor, PD_TO_DB_DTYPES_FLAVORS['bool']['default']) 2299 ) 2300 if flavor in DB_FLAVORS_CAST_DTYPES: 2301 bool_typ = DB_FLAVORS_CAST_DTYPES[flavor].get(bool_typ, bool_typ) 2302 val_to_cast = ( 2303 -987654321 2304 if flavor in ('mysql', 'mariadb') 2305 else 0 2306 ) 2307 return f'CAST({val_to_cast} AS {bool_typ})' 2308 if 'time' in typ.lower() or 'date' in typ.lower(): 2309 db_type = typ if typ.isupper() else None 2310 return dateadd_str(flavor=flavor, begin='1900-01-01', db_type=db_type) 2311 if 'float' in typ.lower() or 'double' in typ.lower() or typ.lower() in ('decimal',): 2312 return '-987654321.0' 2313 if flavor == 'oracle' and typ.lower().split('(', maxsplit=1)[0] == 'char': 2314 return "'-987654321'" 2315 if flavor == 'oracle' and typ.lower() in ('blob', 'bytes'): 2316 return '00' 2317 if typ.lower() in ('uniqueidentifier', 'guid', 'uuid'): 2318 magic_val = 'DEADBEEF-ABBA-BABE-CAFE-DECAFC0FFEE5' 2319 if flavor == 'mssql': 2320 return f"CAST('{magic_val}' AS UNIQUEIDENTIFIER)" 2321 return f"'{magic_val}'" 2322 return ('n' if flavor == 'oracle' else '') + "'-987654321'" 2323 2324 2325 2326def get_rename_table_queries( 2327 old_table: str, 2328 new_table: str, 2329 flavor: str, 2330 schema: Optional[str] = None, 2331 new_schema: Optional[str] = None, 2332) -> List[str]: 2333 """ 2334 Return queries to alter a table's name. 2335 2336 Parameters 2337 ---------- 2338 old_table: str 2339 The unquoted name of the old table. 2340 2341 new_table: str 2342 The unquoted name of the new table. 2343 2344 flavor: str 2345 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`. 2346 2347 schema: Optional[str], default None 2348 The schema on which the table resides. 2349 2350 new_schema: Optional[str], default None 2351 If provided, move the table to this schema. 2352 Defaults to `schema`. 2353 2354 Returns 2355 ------- 2356 A list of `ALTER TABLE` or equivalent queries for the database flavor. 2357 """ 2358 if new_schema is None: 2359 new_schema = schema 2360 2361 if flavor in ('sqlite', 'geopackage'): 2362 schema = None 2363 new_schema = None 2364 2365 old_table_name = sql_item_name(old_table, flavor, schema) 2366 new_table_name = sql_item_name(new_table, flavor, None) 2367 2368 if flavor == 'mssql': 2369 queries = [] 2370 if schema != new_schema: 2371 queries.append(f"ALTER SCHEMA {sql_item_name(new_schema, flavor)} TRANSFER {old_table_name}") 2372 schema = new_schema 2373 old_table_name = sql_item_name(old_table, flavor, schema) 2374 if old_table != new_table: 2375 queries.append(f"EXEC sp_rename '{old_table_name}', '{new_table}'") 2376 return queries 2377 2378 if flavor in ('mysql', 'mariadb'): 2379 new_table_qualified = sql_item_name(new_table, flavor, new_schema) 2380 return [f"RENAME TABLE {old_table_name} TO {new_table_qualified}"] 2381 2382 if_exists_str = "IF EXISTS" if flavor in DROP_IF_EXISTS_FLAVORS else "" 2383 if flavor == 'duckdb': 2384 tmp_table = '_tmp_rename_' + new_table 2385 tmp_table_name = sql_item_name(tmp_table, flavor, schema) 2386 return ( 2387 get_create_table_queries( 2388 f"SELECT * FROM {old_table_name}", 2389 tmp_table, 2390 'duckdb', 2391 schema, 2392 ) + get_create_table_queries( 2393 f"SELECT * FROM {tmp_table_name}", 2394 new_table, 2395 'duckdb', 2396 new_schema, 2397 ) + [ 2398 f"DROP TABLE {if_exists_str} {tmp_table_name}", 2399 f"DROP TABLE {if_exists_str} {old_table_name}", 2400 ] 2401 ) 2402 2403 queries = [] 2404 if schema != new_schema: 2405 if flavor in ('postgresql', 'postgis', 'timescaledb', 'timescaledb-ha', 'citus', 'cockroachdb'): 2406 queries.append(f"ALTER TABLE {old_table_name} SET SCHEMA {sql_item_name(new_schema, flavor)}") 2407 schema = new_schema 2408 old_table_name = sql_item_name(old_table, flavor, schema) 2409 2410 if old_table != new_table: 2411 queries.append(f"ALTER TABLE {old_table_name} RENAME TO {new_table_name}") 2412 2413 return queries 2414 2415 2416 2417def get_create_table_query( 2418 query_or_dtypes: Union[str, Dict[str, str]], 2419 new_table: str, 2420 flavor: str, 2421 schema: Optional[str] = None, 2422) -> str: 2423 """ 2424 NOTE: This function is deprecated. Use `get_create_table_queries()` instead. 2425 2426 Return a query to create a new table from a `SELECT` query. 2427 2428 Parameters 2429 ---------- 2430 query: Union[str, Dict[str, str]] 2431 The select query to use for the creation of the table. 2432 If a dictionary is provided, return a `CREATE TABLE` query from the given `dtypes` columns. 2433 2434 new_table: str 2435 The unquoted name of the new table. 2436 2437 flavor: str 2438 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`). 2439 2440 schema: Optional[str], default None 2441 The schema on which the table will reside. 2442 2443 Returns 2444 ------- 2445 A `CREATE TABLE` (or `SELECT INTO`) query for the database flavor. 2446 """ 2447 return get_create_table_queries( 2448 query_or_dtypes, 2449 new_table, 2450 flavor, 2451 schema=schema, 2452 primary_key=None, 2453 )[0] 2454 2455 2456 2457def get_create_table_queries( 2458 query_or_dtypes: Union[str, Dict[str, str]], 2459 new_table: str, 2460 flavor: str, 2461 schema: Optional[str] = None, 2462 primary_key: Optional[str] = None, 2463 primary_key_db_type: Optional[str] = None, 2464 autoincrement: bool = False, 2465 datetime_column: Optional[str] = None, 2466 hypertable_chunk_interval: Optional[str] = None, 2467 hypertable_segmentby: Optional[List[str]] = None, 2468 hypertable_orderby: Optional[List[str]] = None, 2469 partition_by_column: Optional[str] = None, 2470 partition_bounds: Optional[List] = None, 2471 partition_scheme_name: Optional[str] = None, 2472 _parse_dtypes: bool = True, 2473) -> List[str]: 2474 """ 2475 Return a query to create a new table from a `SELECT` query or a `dtypes` dictionary. 2476 2477 Parameters 2478 ---------- 2479 query_or_dtypes: Union[str, Dict[str, str]] 2480 The select query to use for the creation of the table. 2481 If a dictionary is provided, return a `CREATE TABLE` query from the given `dtypes` columns. 2482 2483 new_table: str 2484 The unquoted name of the new table. 2485 2486 flavor: str 2487 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`). 2488 2489 schema: Optional[str], default None 2490 The schema on which the table will reside. 2491 2492 primary_key: Optional[str], default None 2493 If provided, designate this column as the primary key in the new table. 2494 2495 primary_key_db_type: Optional[str], default None 2496 If provided, alter the primary key to this type (to set NOT NULL constraint). 2497 2498 autoincrement: bool, default False 2499 If `True` and `primary_key` is provided, create the `primary_key` column 2500 as an auto-incrementing integer column. 2501 2502 datetime_column: Optional[str], default None 2503 If provided, include this column in the primary key. 2504 Applicable to TimescaleDB only. 2505 2506 hypertable_chunk_interval: Optional[str], default None 2507 If provided (and `flavor` is TimescaleDB), create the table as a hypertable using the 2508 declarative `CREATE TABLE ... WITH (tsdb.hypertable, ...)` syntax, partitioned on 2509 `datetime_column` with this chunk interval (e.g. `'10080 minutes'` or `'100000'`). 2510 Requires TimescaleDB 2.21+; callers should fall back to `create_hypertable()` on failure. 2511 2512 hypertable_segmentby: Optional[List[str]], default None 2513 If provided (alongside `hypertable_chunk_interval`), enable the Hypercore columnstore on 2514 the new hypertable and segment compressed chunks by these columns (`tsdb.segmentby`). 2515 Should be low-cardinality columns; high-cardinality / unique columns belong in 2516 `hypertable_orderby` instead. 2517 2518 hypertable_orderby: Optional[List[str]], default None 2519 If provided (alongside `hypertable_chunk_interval`), order the columnstore by these 2520 columns (`tsdb.orderby`, e.g. `['"timestamp" DESC']`). Declaring `segmentby`/`orderby` 2521 in `CREATE TABLE` causes TimescaleDB to auto-create a columnstore policy. 2522 2523 _parse_dtypes: bool, default True 2524 If `True`, cast Pandas dtypes to SQL dtypes. 2525 Otherwise pass through the given value directly. 2526 2527 Returns 2528 ------- 2529 A `CREATE TABLE` (or `SELECT INTO`) query for the database flavor. 2530 """ 2531 if not isinstance(query_or_dtypes, (str, dict)): 2532 raise TypeError("`query_or_dtypes` must be a query or a dtypes dictionary.") 2533 2534 method = ( 2535 _get_create_table_query_from_cte 2536 if isinstance(query_or_dtypes, str) 2537 else _get_create_table_query_from_dtypes 2538 ) 2539 return method( 2540 query_or_dtypes, 2541 new_table, 2542 flavor, 2543 schema=schema, 2544 primary_key=primary_key, 2545 primary_key_db_type=primary_key_db_type, 2546 autoincrement=(autoincrement and flavor not in SKIP_AUTO_INCREMENT_FLAVORS), 2547 datetime_column=datetime_column, 2548 hypertable_chunk_interval=hypertable_chunk_interval, 2549 hypertable_segmentby=hypertable_segmentby, 2550 hypertable_orderby=hypertable_orderby, 2551 partition_by_column=partition_by_column, 2552 partition_bounds=partition_bounds, 2553 partition_scheme_name=partition_scheme_name, 2554 _parse_dtypes=_parse_dtypes, 2555 ) 2556 2557 2558 2559def _get_create_table_query_from_dtypes( 2560 dtypes: Dict[str, str], 2561 new_table: str, 2562 flavor: str, 2563 schema: Optional[str] = None, 2564 primary_key: Optional[str] = None, 2565 primary_key_db_type: Optional[str] = None, 2566 autoincrement: bool = False, 2567 datetime_column: Optional[str] = None, 2568 hypertable_chunk_interval: Optional[str] = None, 2569 hypertable_segmentby: Optional[List[str]] = None, 2570 hypertable_orderby: Optional[List[str]] = None, 2571 partition_by_column: Optional[str] = None, 2572 partition_bounds: Optional[List] = None, 2573 partition_scheme_name: Optional[str] = None, 2574 _parse_dtypes: bool = True, 2575) -> List[str]: 2576 """ 2577 Create a new table from a `dtypes` dictionary. 2578 """ 2579 from meerschaum.utils.dtypes.sql import get_db_type_from_pd_type, AUTO_INCREMENT_COLUMN_FLAVORS 2580 if not dtypes and not primary_key: 2581 raise ValueError(f"Expecting columns for table '{new_table}'.") 2582 2583 if flavor in SKIP_AUTO_INCREMENT_FLAVORS: 2584 autoincrement = False 2585 2586 cols_types = ( 2587 [ 2588 ( 2589 primary_key, 2590 get_db_type_from_pd_type(dtypes.get(primary_key, 'int') or 'int', flavor=flavor) 2591 ) if _parse_dtypes else ( 2592 primary_key, 2593 dtypes.get(primary_key, 'INT') or 'INT' 2594 ) 2595 ] 2596 if primary_key 2597 else [] 2598 ) + [ 2599 ( 2600 col, 2601 ( 2602 get_db_type_from_pd_type(typ, flavor=flavor) 2603 if _parse_dtypes 2604 else typ 2605 ) 2606 ) 2607 for col, typ in dtypes.items() 2608 if col != primary_key 2609 ] 2610 2611 table_name = sql_item_name(new_table, schema=schema, flavor=flavor) 2612 primary_key_name = sql_item_name(primary_key, flavor) if primary_key else None 2613 primary_key_constraint_name = ( 2614 sql_item_name(f'PK_{new_table}', flavor, None) 2615 if primary_key 2616 else None 2617 ) 2618 datetime_column_name = sql_item_name(datetime_column, flavor) if datetime_column else None 2619 2620 ### Native (non-TimescaleDB) declarative range partitioning. The partition column must be 2621 ### part of the primary key, so it is folded into a composite PK just like the TimescaleDB 2622 ### datetime column below. 2623 is_native_partitioned = bool( 2624 partition_by_column is not None 2625 and flavor in NATIVE_PARTITION_FLAVORS 2626 ) 2627 partition_by_column_name = ( 2628 sql_item_name(partition_by_column, flavor) 2629 if (partition_by_column is not None and flavor in NATIVE_PARTITION_FLAVORS | {'mssql'}) 2630 else None 2631 ) 2632 ### MSSQL partitions a table by placing its clustered index on a partition scheme (created 2633 ### beforehand by the connector); the partitioning column must be part of that clustered key. 2634 is_mssql_partitioned = bool( 2635 flavor == 'mssql' 2636 and partition_scheme_name is not None 2637 and partition_by_column is not None 2638 ) 2639 partition_scheme_item = ( 2640 sql_item_name(partition_scheme_name, flavor, None) if is_mssql_partitioned else None 2641 ) 2642 primary_key_clustered = ( 2643 "CLUSTERED" 2644 if not datetime_column or datetime_column == primary_key 2645 else "NONCLUSTERED" 2646 ) 2647 query = f"CREATE TABLE {table_name} (" 2648 if primary_key: 2649 col_db_type = cols_types[0][1] 2650 auto_increment_str = (' ' + AUTO_INCREMENT_COLUMN_FLAVORS.get( 2651 flavor, 2652 AUTO_INCREMENT_COLUMN_FLAVORS['default'] 2653 )) if autoincrement or primary_key not in dtypes else '' 2654 col_name = sql_item_name(primary_key, flavor=flavor, schema=None) 2655 2656 if flavor in ('sqlite', 'geopackage'): 2657 query += ( 2658 f"\n {col_name} " 2659 + (f"{col_db_type}" if not auto_increment_str else 'INTEGER') 2660 + f" PRIMARY KEY{auto_increment_str} NOT NULL," 2661 ) 2662 elif flavor == 'oracle': 2663 query += f"\n {col_name} {col_db_type} {auto_increment_str} PRIMARY KEY," 2664 elif ( 2665 flavor in ('timescaledb', 'timescaledb-ha') 2666 and datetime_column 2667 and datetime_column != primary_key 2668 ): 2669 query += f"\n {col_name} {col_db_type}{auto_increment_str} NOT NULL," 2670 elif is_native_partitioned and partition_by_column != primary_key: 2671 ### Defer the PK to a composite `PRIMARY KEY(partition_col, pk)` below. 2672 query += f"\n {col_name} {col_db_type}{auto_increment_str} NOT NULL," 2673 elif flavor == 'mssql': 2674 query += f"\n {col_name} {col_db_type}{auto_increment_str} NOT NULL," 2675 else: 2676 query += f"\n {col_name} {col_db_type} PRIMARY KEY{auto_increment_str} NOT NULL," 2677 2678 for col, db_type in cols_types: 2679 if col == primary_key: 2680 continue 2681 col_name = sql_item_name(col, schema=None, flavor=flavor) 2682 query += f"\n {col_name} {db_type}," 2683 if ( 2684 flavor in ('timescaledb', 'timescaledb-ha') 2685 and datetime_column 2686 and primary_key 2687 and datetime_column != primary_key 2688 ): 2689 query += f"\n PRIMARY KEY({datetime_column_name}, {primary_key_name})," 2690 2691 if is_native_partitioned and primary_key and partition_by_column != primary_key: 2692 ### MySQL requires an `AUTO_INCREMENT` column to be the first column of a key, so list the 2693 ### primary key first there; PostgreSQL is order-insensitive. 2694 if flavor in MYSQL_PARTITION_FLAVORS: 2695 query += f"\n PRIMARY KEY({primary_key_name}, {partition_by_column_name})," 2696 else: 2697 query += f"\n PRIMARY KEY({partition_by_column_name}, {primary_key_name})," 2698 2699 if flavor == 'mssql' and primary_key: 2700 if is_mssql_partitioned and partition_by_column != primary_key: 2701 ### Make the primary key CLUSTERED and place it on the partition scheme, including the 2702 ### partition column as the leading key. This is what partitions the table's storage. 2703 query += ( 2704 f"\n CONSTRAINT {primary_key_constraint_name} PRIMARY KEY CLUSTERED " 2705 f"({partition_by_column_name}, {primary_key_name}) " 2706 f"ON {partition_scheme_item}({partition_by_column_name})," 2707 ) 2708 else: 2709 query += f"\n CONSTRAINT {primary_key_constraint_name} PRIMARY KEY {primary_key_clustered} ({primary_key_name})," 2710 2711 query = query[:-1] 2712 query += "\n)" 2713 2714 if is_mssql_partitioned and not primary_key: 2715 ### No primary key to anchor storage on the scheme, so place the table itself on it. 2716 query += f"\nON {partition_scheme_item}({partition_by_column_name})" 2717 2718 if is_native_partitioned and flavor in PG_PARTITION_FLAVORS: 2719 ### The parent table holds no rows directly; child partitions are created on demand 2720 ### (see `SQLConnector._create_missing_partitions`). 2721 query += f"\nPARTITION BY RANGE ({partition_by_column_name})" 2722 2723 elif is_native_partitioned and flavor in MYSQL_PARTITION_FLAVORS: 2724 ### MySQL/MariaDB cannot create a RANGE-partitioned table with zero partitions, so the 2725 ### initial partitions (computed by the connector from the creation dataframe) are declared 2726 ### inline. Further partitions are appended later via `ALTER TABLE ... ADD PARTITION`. 2727 ### `RANGE COLUMNS` partitions on the DATETIME column directly (no integer conversion). 2728 bounds = partition_bounds or [] 2729 partition_defs = ',\n'.join( 2730 f" PARTITION {sql_item_name(part_name, flavor)} VALUES LESS THAN ({hi_literal})" 2731 for part_name, hi_literal in bounds 2732 ) 2733 query += ( 2734 f"\nPARTITION BY RANGE COLUMNS ({partition_by_column_name}) (\n" 2735 f"{partition_defs}\n)" 2736 ) 2737 2738 if ( 2739 hypertable_chunk_interval is not None 2740 and flavor in ('timescaledb', 'timescaledb-ha') 2741 and datetime_column 2742 ): 2743 partition_column = datetime_column.replace("'", "''") 2744 chunk_interval = hypertable_chunk_interval.replace("'", "''") 2745 with_options = [ 2746 "tsdb.hypertable", 2747 f"tsdb.partition_column='{partition_column}'", 2748 f"tsdb.chunk_interval='{chunk_interval}'", 2749 ] 2750 ### Declaring `segmentby`/`orderby` enables the Hypercore columnstore and causes 2751 ### TimescaleDB to auto-create a columnstore (compression) policy. 2752 if hypertable_segmentby: 2753 segmentby = ', '.join(hypertable_segmentby).replace("'", "''") 2754 with_options.append(f"tsdb.segmentby='{segmentby}'") 2755 if hypertable_orderby: 2756 orderby = ', '.join(hypertable_orderby).replace("'", "''") 2757 with_options.append(f"tsdb.orderby='{orderby}'") 2758 query += "\nWITH (\n " + ",\n ".join(with_options) + "\n)" 2759 2760 queries = [query] 2761 return queries 2762 2763 2764 2765def _get_create_table_query_from_cte( 2766 query: str, 2767 new_table: str, 2768 flavor: str, 2769 schema: Optional[str] = None, 2770 primary_key: Optional[str] = None, 2771 primary_key_db_type: Optional[str] = None, 2772 autoincrement: bool = False, 2773 datetime_column: Optional[str] = None, 2774 hypertable_chunk_interval: Optional[str] = None, 2775 hypertable_segmentby: Optional[List[str]] = None, 2776 hypertable_orderby: Optional[List[str]] = None, 2777 partition_by_column: Optional[str] = None, 2778 partition_bounds: Optional[List] = None, 2779 partition_scheme_name: Optional[str] = None, 2780 _parse_dtypes=None, 2781) -> List[str]: 2782 """ 2783 Create a new table from a CTE query. 2784 NOTE: `hypertable_chunk_interval`/`hypertable_segmentby`/`hypertable_orderby` are ignored here; 2785 `CREATE TABLE ... AS` cannot declare a hypertable. Convert via `create_hypertable()` afterward. 2786 """ 2787 import textwrap 2788 create_cte = 'create_query' 2789 create_cte_name = sql_item_name(create_cte, flavor, None) 2790 new_table_name = sql_item_name(new_table, flavor, schema) 2791 primary_key_constraint_name = ( 2792 sql_item_name(f'PK_{new_table}', flavor, None) 2793 if primary_key 2794 else None 2795 ) 2796 primary_key_name = ( 2797 sql_item_name(primary_key, flavor, None) 2798 if primary_key 2799 else None 2800 ) 2801 primary_key_clustered = ( 2802 "CLUSTERED" 2803 if not datetime_column or datetime_column == primary_key 2804 else "NONCLUSTERED" 2805 ) 2806 datetime_column_name = ( 2807 sql_item_name(datetime_column, flavor) 2808 if datetime_column 2809 else None 2810 ) 2811 query = query.lstrip() 2812 is_with_query = get_leading_with_index(query) != -1 2813 if flavor in ('mssql',): 2814 if is_with_query: 2815 final_select_ix = find_top_level_select_index(query, last=True) 2816 if final_select_ix == -1: 2817 from meerschaum.utils.warnings import error 2818 error( 2819 "Could not locate the top-level SELECT of the following query:\n" 2820 + query 2821 ) 2822 create_table_queries = [ 2823 ( 2824 query[:final_select_ix].rstrip() + ',\n' 2825 + f"{create_cte_name} AS (\n" 2826 + textwrap.indent(query[final_select_ix:], ' ') 2827 + "\n)\n" 2828 + f"SELECT *\nINTO {new_table_name}\nFROM {create_cte_name}" 2829 ), 2830 ] 2831 else: 2832 create_table_queries = [ 2833 ( 2834 "SELECT *\n" 2835 f"INTO {new_table_name}\n" 2836 f"FROM (\n{textwrap.indent(query, ' ')}\n) AS {create_cte_name}" 2837 ), 2838 ] 2839 2840 alter_type_queries = [] 2841 if primary_key_db_type: 2842 alter_type_queries.extend([ 2843 ( 2844 f"ALTER TABLE {new_table_name}\n" 2845 f"ALTER COLUMN {primary_key_name} {primary_key_db_type} NOT NULL" 2846 ), 2847 ]) 2848 alter_type_queries.extend([ 2849 ( 2850 f"ALTER TABLE {new_table_name}\n" 2851 f"ADD CONSTRAINT {primary_key_constraint_name} " 2852 f"PRIMARY KEY {primary_key_clustered} ({primary_key_name})" 2853 ), 2854 ]) 2855 elif flavor in (None,): 2856 create_table_queries = [ 2857 ( 2858 f"WITH {create_cte_name} AS (\n{textwrap.indent(query, ' ')}\n)\n" 2859 f"CREATE TABLE {new_table_name} AS\n" 2860 "SELECT *\n" 2861 f"FROM {create_cte_name}" 2862 ), 2863 ] 2864 2865 alter_type_queries = [ 2866 ( 2867 f"ALTER TABLE {new_table_name}\n" 2868 f"ADD PRIMARY KEY ({primary_key_name})" 2869 ), 2870 ] 2871 elif flavor in ('sqlite', 'geopackage'): 2872 if is_with_query: 2873 create_table_queries = [ 2874 f"CREATE TABLE {new_table_name} AS\n{query}" 2875 ] 2876 else: 2877 create_table_queries = [ 2878 ( 2879 f"CREATE TABLE {new_table_name} AS\n" 2880 "SELECT *\n" 2881 f"FROM (\n{textwrap.indent(query, ' ')}\n)" 2882 ), 2883 ] 2884 ### SQLite does not support ALTER TABLE ... ADD PRIMARY KEY. 2885 alter_type_queries = [] 2886 elif flavor in ( 2887 'mysql', 'mariadb', 'duckdb', 'oracle', 2888 'postgresql', 'postgis', 'timescaledb', 'timescaledb-ha', 'citus', 'cockroachdb' 2889 ): 2890 if is_with_query: 2891 create_table_queries = [ 2892 f"CREATE TABLE {new_table_name} AS\n{query}" 2893 ] 2894 else: 2895 create_table_queries = [ 2896 ( 2897 f"CREATE TABLE {new_table_name} AS\n" 2898 "SELECT *\n" 2899 f"FROM (\n{textwrap.indent(query, ' ')}\n)" 2900 + (f" AS {create_cte_name}" if flavor != 'oracle' else '') 2901 ), 2902 ] 2903 2904 pk_cols_str = ( 2905 f"{datetime_column_name}, {primary_key_name}" 2906 if ( 2907 flavor in ('timescaledb', 'timescaledb-ha') 2908 and datetime_column 2909 and datetime_column != primary_key 2910 ) 2911 else f"{primary_key_name}" 2912 ) 2913 alter_type_queries = [ 2914 ( 2915 f"ALTER TABLE {new_table_name}\n" 2916 f"ADD PRIMARY KEY ({pk_cols_str})" 2917 ), 2918 ] 2919 else: 2920 if is_with_query: 2921 create_table_queries = [ 2922 ( 2923 "SELECT *\n" 2924 f"INTO {new_table_name}\n" 2925 f"FROM (\n{textwrap.indent(query, ' ')}\n) AS {create_cte_name}" 2926 ), 2927 ] 2928 else: 2929 create_table_queries = [ 2930 ( 2931 "SELECT *\n" 2932 f"INTO {new_table_name}\n" 2933 f"FROM (\n{textwrap.indent(query, ' ')}\n) AS {create_cte_name}" 2934 ), 2935 ] 2936 2937 alter_type_queries = [ 2938 ( 2939 f"ALTER TABLE {new_table_name}\n" 2940 f"ADD PRIMARY KEY ({primary_key_name})" 2941 ), 2942 ] 2943 2944 if not primary_key: 2945 return create_table_queries 2946 2947 return create_table_queries + alter_type_queries 2948 2949 2950 2951def get_reset_autoincrement_queries( 2952 table: str, 2953 column: str, 2954 connector: mrsm.connectors.SQLConnector, 2955 schema: Optional[str] = None, 2956 debug: bool = False, 2957) -> List[str]: 2958 """ 2959 Return a list of queries to reset a table's auto-increment counter to the next largest value. 2960 2961 Parameters 2962 ---------- 2963 table: str 2964 The name of the table on which the auto-incrementing column exists. 2965 2966 column: str 2967 The name of the auto-incrementing column. 2968 2969 connector: mrsm.connectors.SQLConnector 2970 The SQLConnector to the database on which the table exists. 2971 2972 schema: Optional[str], default None 2973 The schema of the table. Defaults to `connector.schema`. 2974 2975 Returns 2976 ------- 2977 A list of queries to be executed to reset the auto-incrementing column. 2978 """ 2979 if not table_exists(table, connector, schema=schema, debug=debug): 2980 return [] 2981 2982 schema = schema or connector.schema 2983 max_id_name = sql_item_name('max_id', connector.flavor) 2984 table_name = sql_item_name(table, connector.flavor, schema) 2985 table_seq_name = sql_item_name(table + '_' + column + '_seq', connector.flavor, schema) 2986 column_name = sql_item_name(column, connector.flavor) 2987 max_id = connector.value( 2988 f""" 2989 SELECT COALESCE(MAX({column_name}), 0) AS {max_id_name} 2990 FROM {table_name} 2991 """, 2992 debug=debug, 2993 silent=True, 2994 ) 2995 if max_id is None: 2996 return [] 2997 2998 reset_queries = reset_autoincrement_queries.get( 2999 connector.flavor, 3000 reset_autoincrement_queries['default'] 3001 ) 3002 if not isinstance(reset_queries, list): 3003 reset_queries = [reset_queries] 3004 3005 return [ 3006 query.format( 3007 column=column, 3008 column_name=column_name, 3009 table=table, 3010 table_name=table_name, 3011 table_seq_name=table_seq_name, 3012 val=max_id, 3013 val_plus_1=(max_id + 1), 3014 ) 3015 for query in reset_queries 3016 ] 3017 3018 3019 3020def get_create_schema_if_not_exists_queries( 3021 schema: str, 3022 flavor: str, 3023) -> List[str]: 3024 """ 3025 Return the queries to create a schema if it does not yet exist. 3026 For databases which do not support schemas, an empty list will be returned. 3027 """ 3028 if not schema: 3029 return [] 3030 3031 if flavor in NO_SCHEMA_FLAVORS: 3032 return [] 3033 3034 if schema == DEFAULT_SCHEMA_FLAVORS.get(flavor, None): 3035 return [] 3036 3037 clean(schema) 3038 3039 if flavor == 'mssql': 3040 return [ 3041 ( 3042 f"IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{schema}')\n" 3043 "BEGIN\n" 3044 f" EXEC('CREATE SCHEMA {sql_item_name(schema, flavor)}');\n" 3045 "END;" 3046 ) 3047 ] 3048 3049 if flavor == 'oracle': 3050 return [] 3051 3052 return [ 3053 f"CREATE SCHEMA IF NOT EXISTS {sql_item_name(schema, flavor)};" 3054 ] 3055 3056 3057NATIVE_PARTITION_FLAVORS = PG_PARTITION_FLAVORS | MYSQL_PARTITION_FLAVORS 3058 3059 3060### `table_name` is the escaped name of the table. 3061### `table` is the unescaped name of the table. 3062exists_queries = { 3063 'default': "SELECT COUNT(*) FROM {table_name} WHERE 1 = 0", 3064} 3065 3066 3067 3068DROP_INDEX_IF_EXISTS_FLAVORS = { 3069 'mssql', 3070 'timescaledb', 3071 'timescaledb-ha', 3072 'postgresql', 3073 'postgis', 3074 'sqlite', 3075 'geopackage', 3076 'citus', 3077} 3078 3079 3080 3081COALESCE_UNIQUE_INDEX_FLAVORS = { 3082 'timescaledb', 3083 'timescaledb-ha', 3084 'postgresql', 3085 'postgis', 3086 'citus', 3087} 3088 3089 3090 3091hypertable_queries = { 3092 'timescaledb': 'SELECT hypertable_size(\'{table_name}\')', 3093 'timescaledb-ha': 'SELECT hypertable_size(\'{table_name}\')', 3094 'citus': 'SELECT citus_table_size(\'{table_name}\')', 3095} 3096 3097 3098 3099table_wrappers = { 3100 'default' : ('"', '"'), 3101 'timescaledb' : ('"', '"'), 3102 'timescaledb-ha': ('"', '"'), 3103 'citus' : ('"', '"'), 3104 'duckdb' : ('"', '"'), 3105 'postgresql' : ('"', '"'), 3106 'postgis' : ('"', '"'), 3107 'sqlite' : ('"', '"'), 3108 'geopackage' : ('"', '"'), 3109 'mysql' : ('`', '`'), 3110 'mariadb' : ('`', '`'), 3111 'mssql' : ('[', ']'), 3112 'cockroachdb' : ('"', '"'), 3113 'oracle' : ('"', '"'), 3114} 3115 3116 3117 3118OMIT_NULLSFIRST_FLAVORS = { 3119 'mariadb', 3120 'mysql', 3121 'mssql', 3122} 3123 3124 3125 3126SINGLE_ALTER_TABLE_FLAVORS = { 3127 'duckdb', 3128 'sqlite', 3129 'geopackage', 3130 'mssql', 3131 'oracle', 3132} 3133 3134 3135 3136NO_SELECT_INTO_FLAVORS = { 3137 'sqlite', 3138 'geopackage', 3139 'oracle', 3140 'mysql', 3141 'mariadb', 3142 'duckdb', 3143}
test_queries =
{'default': 'SELECT 1', 'oracle': 'SELECT 1 FROM DUAL', 'informix': 'SELECT COUNT(*) FROM systables', 'hsqldb': 'SELECT 1 FROM INFORMATION_SCHEMA.SYSTEM_USERS'}
version_queries =
{'default': 'SELECT VERSION() AS {version_name}', 'sqlite': 'SELECT SQLITE_VERSION() AS {version_name}', 'geopackage': 'SELECT SQLITE_VERSION() AS {version_name}', 'mssql': 'SELECT @@version', 'oracle': 'SELECT version from PRODUCT_COMPONENT_VERSION WHERE rownum = 1'}
SKIP_IF_EXISTS_FLAVORS =
{'oracle', 'mssql'}
DROP_IF_EXISTS_FLAVORS =
{'timescaledb', 'postgis', 'mssql', 'sqlite', 'timescaledb-ha', 'mysql', 'geopackage', 'postgresql', 'citus', 'mariadb'}
SKIP_AUTO_INCREMENT_FLAVORS =
{'citus', 'duckdb'}
columns_types_queries =
{'default': "\n SELECT\n table_catalog AS database,\n table_schema AS schema,\n table_name AS table,\n column_name AS column,\n data_type AS type,\n numeric_precision,\n numeric_scale\n FROM information_schema.columns\n WHERE table_name IN ('{table}', '{table_trunc}')\n ", 'sqlite': '\n SELECT\n \'\' "database",\n \'\' "schema",\n m.name "table",\n p.name "column",\n p.type "type"\n FROM sqlite_master m\n LEFT OUTER JOIN pragma_table_info(m.name) p\n ON m.name <> p.name\n WHERE m.type = \'table\'\n AND m.name IN (\'{table}\', \'{table_trunc}\')\n ', 'geopackage': '\n SELECT\n \'\' "database",\n \'\' "schema",\n m.name "table",\n p.name "column",\n p.type "type"\n FROM sqlite_master m\n LEFT OUTER JOIN pragma_table_info(m.name) p\n ON m.name <> p.name\n WHERE m.type = \'table\'\n AND m.name IN (\'{table}\', \'{table_trunc}\')\n ', 'mssql': "\n SELECT\n TABLE_CATALOG AS [database],\n TABLE_SCHEMA AS [schema],\n TABLE_NAME AS [table],\n COLUMN_NAME AS [column],\n CASE WHEN CHARACTER_MAXIMUM_LENGTH = -1 THEN DATA_TYPE + '(max)' ELSE DATA_TYPE END AS [type],\n NUMERIC_PRECISION AS [numeric_precision],\n NUMERIC_SCALE AS [numeric_scale]\n FROM {db_prefix}INFORMATION_SCHEMA.COLUMNS WITH (NOLOCK)\n WHERE TABLE_NAME IN (\n '{table}',\n '{table_trunc}'\n )\n\n ", 'mysql': "\n SELECT\n TABLE_SCHEMA `database`,\n TABLE_SCHEMA `schema`,\n TABLE_NAME `table`,\n COLUMN_NAME `column`,\n DATA_TYPE `type`,\n NUMERIC_PRECISION `numeric_precision`,\n NUMERIC_SCALE `numeric_scale`\n FROM INFORMATION_SCHEMA.COLUMNS\n WHERE TABLE_NAME IN ('{table}', '{table_trunc}')\n ", 'mariadb': "\n SELECT\n TABLE_SCHEMA `database`,\n TABLE_SCHEMA `schema`,\n TABLE_NAME `table`,\n COLUMN_NAME `column`,\n DATA_TYPE `type`,\n NUMERIC_PRECISION `numeric_precision`,\n NUMERIC_SCALE `numeric_scale`\n FROM INFORMATION_SCHEMA.COLUMNS\n WHERE TABLE_NAME IN ('{table}', '{table_trunc}')\n ", 'oracle': '\n SELECT\n NULL AS "database",\n NULL AS "schema",\n TABLE_NAME AS "table",\n COLUMN_NAME AS "column",\n DATA_TYPE AS "type",\n DATA_PRECISION AS "numeric_precision",\n DATA_SCALE AS "numeric_scale"\n FROM all_tab_columns\n WHERE TABLE_NAME IN (\n \'{table}\',\n \'{table_trunc}\',\n \'{table_lower}\',\n \'{table_lower_trunc}\',\n \'{table_upper}\',\n \'{table_upper_trunc}\'\n )\n '}
columns_indices_queries =
{'default': '\n SELECT\n current_database() AS "database",\n n.nspname AS "schema",\n t.relname AS "table",\n c.column_name AS "column",\n i.relname AS "index",\n CASE WHEN con.contype = \'p\' THEN \'PRIMARY KEY\' ELSE \'INDEX\' END AS "index_type"\n FROM pg_class t\n INNER JOIN pg_index AS ix\n ON t.oid = ix.indrelid\n INNER JOIN pg_class AS i\n ON i.oid = ix.indexrelid\n INNER JOIN pg_namespace AS n\n ON n.oid = t.relnamespace\n INNER JOIN pg_attribute AS a\n ON a.attnum = ANY(ix.indkey)\n AND a.attrelid = t.oid\n INNER JOIN information_schema.columns AS c\n ON c.column_name = a.attname\n AND c.table_name = t.relname\n AND c.table_schema = n.nspname\n LEFT JOIN pg_constraint AS con\n ON con.conindid = i.oid\n AND con.contype = \'p\'\n WHERE\n t.relname IN (\'{table}\', \'{table_trunc}\')\n AND n.nspname = \'{schema}\'\n ', 'sqlite': '\n WITH indexed_columns AS (\n SELECT\n \'{table}\' AS table_name,\n pi.name AS column_name,\n i.name AS index_name,\n \'INDEX\' AS index_type\n FROM\n sqlite_master AS i,\n pragma_index_info(i.name) AS pi\n WHERE\n i.type = \'index\'\n AND i.tbl_name = \'{table}\'\n ),\n primary_key_columns AS (\n SELECT\n \'{table}\' AS table_name,\n ti.name AS column_name,\n \'PRIMARY_KEY\' AS index_name,\n \'PRIMARY KEY\' AS index_type\n FROM\n pragma_table_info(\'{table}\') AS ti\n WHERE\n ti.pk > 0\n )\n SELECT\n NULL AS "database",\n NULL AS "schema",\n "table_name" AS "table",\n "column_name" AS "column",\n "index_name" AS "index",\n "index_type"\n FROM indexed_columns\n UNION ALL\n SELECT\n NULL AS "database",\n NULL AS "schema",\n table_name AS "table",\n column_name AS "column",\n index_name AS "index",\n index_type\n FROM primary_key_columns\n ', 'geopackage': '\n WITH indexed_columns AS (\n SELECT\n \'{table}\' AS table_name,\n pi.name AS column_name,\n i.name AS index_name,\n \'INDEX\' AS index_type\n FROM\n sqlite_master AS i,\n pragma_index_info(i.name) AS pi\n WHERE\n i.type = \'index\'\n AND i.tbl_name = \'{table}\'\n ),\n primary_key_columns AS (\n SELECT\n \'{table}\' AS table_name,\n ti.name AS column_name,\n \'PRIMARY_KEY\' AS index_name,\n \'PRIMARY KEY\' AS index_type\n FROM\n pragma_table_info(\'{table}\') AS ti\n WHERE\n ti.pk > 0\n )\n SELECT\n NULL AS "database",\n NULL AS "schema",\n "table_name" AS "table",\n "column_name" AS "column",\n "index_name" AS "index",\n "index_type"\n FROM indexed_columns\n UNION ALL\n SELECT\n NULL AS "database",\n NULL AS "schema",\n table_name AS "table",\n column_name AS "column",\n index_name AS "index",\n index_type\n FROM primary_key_columns\n ', 'mssql': "\n SELECT\n NULL AS [database],\n s.name AS [schema],\n t.name AS [table],\n c.name AS [column],\n i.name AS [index],\n CASE\n WHEN kc.type = 'PK' THEN 'PRIMARY KEY'\n ELSE 'INDEX'\n END AS [index_type],\n CASE\n WHEN i.type = 1 THEN CAST(1 AS BIT)\n ELSE CAST(0 AS BIT)\n END AS [clustered]\n FROM\n sys.schemas s WITH (NOLOCK)\n INNER JOIN sys.tables t WITH (NOLOCK)\n ON s.schema_id = t.schema_id\n INNER JOIN sys.indexes i WITH (NOLOCK)\n ON t.object_id = i.object_id\n INNER JOIN sys.index_columns ic WITH (NOLOCK)\n ON i.object_id = ic.object_id\n AND i.index_id = ic.index_id\n INNER JOIN sys.columns c WITH (NOLOCK)\n ON ic.object_id = c.object_id\n AND ic.column_id = c.column_id\n LEFT JOIN sys.key_constraints kc WITH (NOLOCK)\n ON kc.parent_object_id = i.object_id\n AND kc.type = 'PK'\n AND kc.name = i.name\n WHERE\n t.name IN ('{table}', '{table_trunc}')\n AND s.name = '{schema}'\n AND i.type IN (1, 2)\n ", 'oracle': '\n SELECT\n NULL AS "database",\n ic.table_owner AS "schema",\n ic.table_name AS "table",\n ic.column_name AS "column",\n i.index_name AS "index",\n CASE\n WHEN c.constraint_type = \'P\' THEN \'PRIMARY KEY\'\n WHEN i.uniqueness = \'UNIQUE\' THEN \'UNIQUE INDEX\'\n ELSE \'INDEX\'\n END AS index_type\n FROM\n all_ind_columns ic\n INNER JOIN all_indexes i\n ON ic.index_name = i.index_name\n AND ic.table_owner = i.owner\n LEFT JOIN all_constraints c\n ON i.index_name = c.constraint_name\n AND i.table_owner = c.owner\n AND c.constraint_type = \'P\'\n WHERE ic.table_name IN (\n \'{table}\',\n \'{table_trunc}\',\n \'{table_upper}\',\n \'{table_upper_trunc}\'\n )\n ', 'mysql': "\n SELECT\n TABLE_SCHEMA AS `database`,\n TABLE_SCHEMA AS `schema`,\n TABLE_NAME AS `table`,\n COLUMN_NAME AS `column`,\n INDEX_NAME AS `index`,\n CASE\n WHEN NON_UNIQUE = 0 THEN 'PRIMARY KEY'\n ELSE 'INDEX'\n END AS `index_type`\n FROM\n information_schema.STATISTICS\n WHERE\n TABLE_NAME IN ('{table}', '{table_trunc}')\n ", 'mariadb': "\n SELECT\n TABLE_SCHEMA AS `database`,\n TABLE_SCHEMA AS `schema`,\n TABLE_NAME AS `table`,\n COLUMN_NAME AS `column`,\n INDEX_NAME AS `index`,\n CASE\n WHEN NON_UNIQUE = 0 THEN 'PRIMARY KEY'\n ELSE 'INDEX'\n END AS `index_type`\n FROM\n information_schema.STATISTICS\n WHERE\n TABLE_NAME IN ('{table}', '{table_trunc}')\n "}
reset_autoincrement_queries: Dict[str, str | List[str]] =
{'default': "\n SELECT SETVAL(pg_get_serial_sequence('{table_name}', '{column}'), {val})\n FROM {table_name}\n ", 'mssql': "\n DBCC CHECKIDENT ('{table_name}', RESEED, {val})\n ", 'mysql': '\n ALTER TABLE {table_name} AUTO_INCREMENT = {val}\n ', 'mariadb': '\n ALTER TABLE {table_name} AUTO_INCREMENT = {val}\n ', 'sqlite': "\n UPDATE sqlite_sequence\n SET seq = {val}\n WHERE name = '{table}'\n ", 'geopackage': "\n UPDATE sqlite_sequence\n SET seq = {val}\n WHERE name = '{table}'\n ", 'oracle': 'ALTER TABLE {table_name} MODIFY {column_name} GENERATED BY DEFAULT ON NULL AS IDENTITY (START WITH {val_plus_1})'}
max_name_lens =
{'default': 64, 'mssql': 128, 'oracle': 30, 'postgresql': 64, 'postgis': 64, 'timescaledb': 64, 'timescaledb-ha': 64, 'citus': 64, 'cockroachdb': 64, 'sqlite': 1024, 'geopackage': 1024, 'mysql': 64, 'mariadb': 64}
json_flavors =
{'timescaledb', 'postgis', 'cockroachdb', 'postgresql', 'citus', 'timescaledb-ha'}
NO_SCHEMA_FLAVORS =
{'sqlite', 'mysql', 'oracle', 'duckdb', 'mariadb', 'geopackage'}
DEFAULT_SCHEMA_FLAVORS =
{'postgresql': 'public', 'postgis': 'public', 'timescaledb': 'public', 'timescaledb-ha': 'public', 'citus': 'public', 'cockroachdb': 'public', 'mysql': 'mysql', 'mariadb': 'mysql', 'mssql': 'dbo'}
NO_CTE_FLAVORS =
{'mariadb', 'mysql'}
def
clean(substring: str) -> None:
448def clean(substring: str) -> None: 449 """ 450 Ensure a substring is clean enough to be inserted into a SQL query. 451 Raises an exception when banned words are used. 452 """ 453 from meerschaum.utils.warnings import error 454 banned_symbols = [ 455 ';', '--', 'drop ', '/*', '*/', 456 'union ', 'exec ', 'execute ', 457 'insert ', 'update ', 'delete ', 458 'create ', 'alter ', 'truncate ', 459 ] 460 for symbol in banned_symbols: 461 if symbol in str(substring).lower(): 462 error(f"Invalid string: '{substring}'")
Ensure a substring is clean enough to be inserted into a SQL query. Raises an exception when banned words are used.
def
dateadd_str( flavor: str = 'postgresql', datepart: str = 'day', number: int | float = 0, begin: str | datetime.datetime | int = 'now', db_type: str | None = None) -> str:
468def dateadd_str( 469 flavor: str = 'postgresql', 470 datepart: str = 'day', 471 number: Union[int, float] = 0, 472 begin: Union[str, datetime, int] = 'now', 473 db_type: Optional[str] = None, 474) -> str: 475 """ 476 Generate a `DATEADD` clause depending on database flavor. 477 478 Parameters 479 ---------- 480 flavor: str, default `'postgresql'` 481 SQL database flavor, e.g. `'postgresql'`, `'sqlite'`. 482 483 Currently supported flavors: 484 485 - `'postgresql'` 486 - `'postgis'` 487 - `'timescaledb'` 488 - `'timescaledb-ha'` 489 - `'citus'` 490 - `'cockroachdb'` 491 - `'duckdb'` 492 - `'mssql'` 493 - `'mysql'` 494 - `'mariadb'` 495 - `'sqlite'` 496 - `'geopackage'` 497 - `'oracle'` 498 499 datepart: str, default `'day'` 500 Which part of the date to modify. Supported values: 501 502 - `'year'` 503 - `'month'` 504 - `'day'` 505 - `'hour'` 506 - `'minute'` 507 - `'second'` 508 509 number: Union[int, float], default `0` 510 How many units to add to the date part. 511 512 begin: Union[str, datetime], default `'now'` 513 Base datetime to which to add dateparts. 514 515 db_type: Optional[str], default None 516 If provided, cast the datetime string as the type. 517 Otherwise, infer this from the input datetime value. 518 519 Returns 520 ------- 521 The appropriate `DATEADD` string for the corresponding database flavor. 522 523 Examples 524 -------- 525 >>> dateadd_str( 526 ... flavor='mssql', 527 ... begin=datetime(2022, 1, 1, 0, 0), 528 ... number=1, 529 ... ) 530 "DATEADD(day, 1, CAST('2022-01-01 00:00:00' AS DATETIME2))" 531 >>> dateadd_str( 532 ... flavor='postgresql', 533 ... begin=datetime(2022, 1, 1, 0, 0), 534 ... number=1, 535 ... ) 536 "CAST('2022-01-01 00:00:00' AS TIMESTAMP) + INTERVAL '1 day'" 537 538 """ 539 from meerschaum.utils.packages import attempt_import 540 from meerschaum.utils.dtypes.sql import get_db_type_from_pd_type, get_pd_type_from_db_type 541 dateutil_parser = attempt_import('dateutil.parser') 542 if 'int' in str(type(begin)).lower(): 543 num_str = str(begin) 544 if number is not None and number != 0: 545 num_str += ( 546 f' + {number}' 547 if number > 0 548 else f" - {number * -1}" 549 ) 550 return num_str 551 if datepart not in _VALID_DATEPARTS: 552 raise ValueError( 553 f"Invalid datepart '{datepart}'. Must be one of: {sorted(_VALID_DATEPARTS)}" 554 ) 555 if not begin: 556 return '' 557 558 _original_begin = begin 559 begin_time = None 560 ### Sanity check: make sure `begin` is a valid datetime before we inject anything. 561 if not isinstance(begin, datetime): 562 try: 563 begin_time = dateutil_parser.parse(begin) 564 except Exception: 565 begin_time = None 566 else: 567 begin_time = begin 568 569 ### Unable to parse into a datetime. 570 if begin_time is None: 571 ### Throw an error if banned symbols are included in the `begin` string. 572 clean(str(begin)) 573 ### If begin is a valid datetime, wrap it in quotes. 574 else: 575 if isinstance(begin, datetime) and begin.tzinfo is not None: 576 begin = begin.astimezone(timezone.utc) 577 begin = ( 578 f"'{begin.replace(tzinfo=None)}'" 579 if isinstance(begin, datetime) and flavor in TIMEZONE_NAIVE_FLAVORS 580 else f"'{begin}'" 581 ) 582 583 dt_is_utc = ( 584 begin_time.tzinfo is not None 585 if begin_time is not None 586 else ('+' in str(begin) or '-' in str(begin).split(':', maxsplit=1)[-1]) 587 ) 588 if db_type: 589 db_type_is_utc = 'utc' in get_pd_type_from_db_type(db_type).lower() 590 dt_is_utc = dt_is_utc or db_type_is_utc 591 db_type = db_type or get_db_type_from_pd_type( 592 ('datetime64[ns, UTC]' if dt_is_utc else 'datetime64[ns]'), 593 flavor=flavor, 594 ) 595 596 da = "" 597 if flavor in ( 598 'postgresql', 599 'postgis', 600 'timescaledb', 601 'timescaledb-ha', 602 'cockroachdb', 603 'citus', 604 ): 605 begin = ( 606 f"CAST({begin} AS {db_type})" if begin != 'now' 607 else f"CAST(NOW() AT TIME ZONE 'utc' AS {db_type})" 608 ) 609 da = begin + (f" + INTERVAL '{number} {datepart}'" if number != 0 else '') 610 611 elif flavor == 'duckdb': 612 begin = f"CAST({begin} AS {db_type})" if begin != 'now' else 'NOW()' 613 da = begin + (f" + INTERVAL '{number} {datepart}'" if number != 0 else '') 614 615 elif flavor in ('mssql',): 616 if begin_time and begin_time.microsecond != 0 and not dt_is_utc: 617 begin = begin[:-4] + "'" 618 begin = f"CAST({begin} AS {db_type})" if begin != 'now' else 'GETUTCDATE()' 619 if dt_is_utc: 620 begin += " AT TIME ZONE 'UTC'" 621 da = f"DATEADD({datepart}, {number}, {begin})" if number != 0 else begin 622 623 elif flavor in ('mysql', 'mariadb'): 624 begin = ( 625 f"CAST({begin} AS DATETIME(6))" 626 if begin != 'now' 627 else 'UTC_TIMESTAMP(6)' 628 ) 629 da = (f"DATE_ADD({begin}, INTERVAL {number} {datepart})" if number != 0 else begin) 630 631 elif flavor in ('sqlite', 'geopackage'): 632 da = f"datetime({begin}, '{number} {datepart}')" 633 634 elif flavor == 'oracle': 635 if begin == 'now': 636 begin = str( 637 datetime.now(timezone.utc).replace(tzinfo=None).strftime(r'%Y:%m:%d %M:%S.%f') 638 ) 639 elif begin_time: 640 begin = str(begin_time.strftime(r'%Y-%m-%d %H:%M:%S.%f')) 641 dt_format = 'YYYY-MM-DD HH24:MI:SS.FF' 642 _begin = f"'{begin}'" if begin_time else begin 643 da = ( 644 (f"TO_TIMESTAMP({_begin}, '{dt_format}')" if begin_time else _begin) 645 + (f" + INTERVAL '{number}' {datepart}" if number != 0 else "") 646 ) 647 return da
Generate a DATEADD clause depending on database flavor.
Parameters
flavor (str, default
'postgresql'): SQL database flavor, e.g.'postgresql','sqlite'.Currently supported flavors:
'postgresql''postgis''timescaledb''timescaledb-ha''citus''cockroachdb''duckdb''mssql''mysql''mariadb''sqlite''geopackage''oracle'
datepart (str, default
'day'): Which part of the date to modify. Supported values:'year''month''day''hour''minute''second'
- number (Union[int, float], default
0): How many units to add to the date part. - begin (Union[str, datetime], default
'now'): Base datetime to which to add dateparts. - db_type (Optional[str], default None): If provided, cast the datetime string as the type. Otherwise, infer this from the input datetime value.
Returns
- The appropriate
DATEADDstring for the corresponding database flavor.
Examples
>>> dateadd_str(
... flavor='mssql',
... begin=datetime(2022, 1, 1, 0, 0),
... number=1,
... )
"DATEADD(day, 1, CAST('2022-01-01 00:00:00' AS DATETIME2))"
>>> dateadd_str(
... flavor='postgresql',
... begin=datetime(2022, 1, 1, 0, 0),
... number=1,
... )
"CAST('2022-01-01 00:00:00' AS TIMESTAMP) + INTERVAL '1 day'"
def
test_connection(self, **kw: Any) -> bool | None:
650def test_connection( 651 self, 652 **kw: Any 653) -> Union[bool, None]: 654 """ 655 Test if a successful connection to the database may be made. 656 657 Parameters 658 ---------- 659 **kw: 660 `max_retries` (default `1`), `retry_wait` seconds between attempts (default `0`) 661 and `warn` (print the error of a failed attempt, default `False`). 662 663 Returns 664 ------- 665 `True` if a connection is made, otherwise `False`. 666 667 """ 668 import time 669 max_retries = max(int(kw.get('max_retries', 1)), 1) 670 retry_wait = kw.get('retry_wait', 0) 671 warn = kw.get('warn', False) 672 for attempt in range(max_retries): 673 try: 674 ### The Rust connector's `SELECT 1`: no SQLAlchemy engine, so the base wheel 675 ### can check a reachable database. 676 if self._core_connector().test_connection(): 677 return True 678 except Exception as e: 679 if warn: 680 print(e) 681 if attempt + 1 < max_retries and retry_wait: 682 time.sleep(retry_wait) 683 return False
Test if a successful connection to the database may be made.
Parameters
- **kw::
max_retries(default1),retry_waitseconds between attempts (default0) andwarn(print the error of a failed attempt, defaultFalse).
Returns
Trueif a connection is made, otherwiseFalse.
def
get_distinct_col_count( col: str, query: str, connector: meerschaum.connectors.SQLConnector | None = None, debug: bool = False) -> int | None:
686def get_distinct_col_count( 687 col: str, 688 query: str, 689 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 690 debug: bool = False 691) -> Optional[int]: 692 """ 693 Returns the number of distinct items in a column of a SQL query. 694 695 Parameters 696 ---------- 697 col: str: 698 The column in the query to count. 699 700 query: str: 701 The SQL query to count from. 702 703 connector: Optional[mrsm.connectors.sql.SQLConnector], default None: 704 The SQLConnector to execute the query. 705 706 debug: bool, default False: 707 Verbosity toggle. 708 709 Returns 710 ------- 711 An `int` of the number of columns in the query or `None` if the query fails. 712 713 """ 714 if connector is None: 715 connector = mrsm.get_connector('sql') 716 717 _col_name = sql_item_name(col, connector.flavor, None) 718 719 _meta_query = ( 720 f""" 721 WITH src AS ( {query} ), 722 dist AS ( SELECT DISTINCT {_col_name} FROM src ) 723 SELECT COUNT(*) FROM dist""" 724 ) if connector.flavor not in ('mysql', 'mariadb') else ( 725 f""" 726 SELECT COUNT(*) 727 FROM ( 728 SELECT DISTINCT {_col_name} 729 FROM ({query}) AS src 730 ) AS dist""" 731 ) 732 733 result = connector.value(_meta_query, debug=debug, silent=True) 734 try: 735 return int(result) 736 except Exception: 737 return None
Returns the number of distinct items in a column of a SQL query.
Parameters
- col (str:): The column in the query to count.
- query (str:): The SQL query to count from.
- connector (Optional[mrsm.connectors.sql.SQLConnector], default None:): The SQLConnector to execute the query.
- debug (bool, default False:): Verbosity toggle.
Returns
- An
intof the number of columns in the query orNoneif the query fails.
def
sql_item_name(item: str, flavor: str, schema: str | None = None) -> str:
740def sql_item_name(item: str, flavor: str, schema: Optional[str] = None) -> str: 741 """ 742 Parse SQL items depending on the flavor. 743 744 Parameters 745 ---------- 746 item: str 747 The database item (table, view, etc.) in need of quotes. 748 749 flavor: str 750 The database flavor (`'postgresql'`, `'mssql'`, `'sqllite'`, etc.). 751 752 schema: Optional[str], default None 753 If provided, prefix the table name with the schema. 754 755 Returns 756 ------- 757 A `str` which contains the input `item` wrapped in the corresponding escape characters. 758 759 Examples 760 -------- 761 >>> sql_item_name('table', 'sqlite') 762 '"table"' 763 >>> sql_item_name('table', 'mssql') 764 "[table]" 765 >>> sql_item_name('table', 'postgresql', schema='abc') 766 '"abc"."table"' 767 768 """ 769 _rust = _core.sql_item_name 770 return _rust(item, flavor, schema)
Parse SQL items depending on the flavor.
Parameters
- item (str): The database item (table, view, etc.) in need of quotes.
- flavor (str):
The database flavor (
'postgresql','mssql','sqllite', etc.). - schema (Optional[str], default None): If provided, prefix the table name with the schema.
Returns
- A
strwhich contains the inputitemwrapped in the corresponding escape characters.
Examples
>>> sql_item_name('table', 'sqlite')
'"table"'
>>> sql_item_name('table', 'mssql')
"[table]"
>>> sql_item_name('table', 'postgresql', schema='abc')
'"abc"."table"'
def
pg_capital(s: str, quote_capitals: bool = True) -> str:
773def pg_capital(s: str, quote_capitals: bool = True) -> str: 774 """ 775 If string contains a capital letter, wrap it in double quotes. 776 777 Parameters 778 ---------- 779 s: str 780 The string to be escaped. 781 782 quote_capitals: bool, default True 783 If `False`, do not quote strings with contain only a mix of capital and lower-case letters. 784 785 Returns 786 ------- 787 The input string wrapped in quotes only if it needs them. 788 789 Examples 790 -------- 791 >>> pg_capital("My Table") 792 '"My Table"' 793 >>> pg_capital('my_table') 794 'my_table' 795 796 """ 797 if s.startswith('"') and s.endswith('"'): 798 return s 799 800 s = s.replace('"', '') 801 802 needs_quotes = s.startswith('_') 803 if not needs_quotes: 804 for c in s: 805 if c == '_': 806 continue 807 808 if not c.isalnum() or (quote_capitals and c.isupper()): 809 needs_quotes = True 810 break 811 812 if needs_quotes: 813 return '"' + s + '"' 814 815 return s
If string contains a capital letter, wrap it in double quotes.
Parameters
- s (str): The string to be escaped.
- quote_capitals (bool, default True):
If
False, do not quote strings with contain only a mix of capital and lower-case letters.
Returns
- The input string wrapped in quotes only if it needs them.
Examples
>>> pg_capital("My Table")
'"My Table"'
>>> pg_capital('my_table')
'my_table'
def
oracle_capital(s: str) -> str:
818def oracle_capital(s: str) -> str: 819 """ 820 Capitalize the string of an item on an Oracle database. 821 """ 822 return s
Capitalize the string of an item on an Oracle database.
def
truncate_item_name(item: str, flavor: str) -> str:
825def truncate_item_name(item: str, flavor: str) -> str: 826 """ 827 Truncate item names to stay within the database flavor's character limit. 828 829 Parameters 830 ---------- 831 item: str 832 The database item being referenced. This string is the "canonical" name internally. 833 834 flavor: str 835 The flavor of the database on which `item` resides. 836 837 Returns 838 ------- 839 The truncated string. 840 """ 841 from meerschaum.utils.misc import truncate_string_sections 842 return truncate_string_sections( 843 item, max_len=max_name_lens.get(flavor, max_name_lens['default']) 844 )
Truncate item names to stay within the database flavor's character limit.
Parameters
- item (str): The database item being referenced. This string is the "canonical" name internally.
- flavor (str):
The flavor of the database on which
itemresides.
Returns
- The truncated string.
def
build_where( params: Dict[str, Any], connector: meerschaum.connectors.SQLConnector | None = None, with_where: bool = True, flavor: str = 'postgresql') -> str:
847def build_where( 848 params: Dict[str, Any], 849 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 850 with_where: bool = True, 851 flavor: str = 'postgresql', 852) -> str: 853 """ 854 Build the `WHERE` clause based on the input criteria. 855 856 Parameters 857 ---------- 858 params: Dict[str, Any]: 859 The keywords dictionary to convert into a WHERE clause. 860 If a value is a string which begins with an underscore, negate that value 861 (e.g. `!=` instead of `=` or `NOT IN` instead of `IN`). 862 A value of `_None` will be interpreted as `IS NOT NULL`. 863 864 connector: Optional[meerschaum.connectors.sql.SQLConnector], default None: 865 The Meerschaum SQLConnector that will be executing the query. 866 The connector is used to extract the SQL dialect. 867 868 with_where: bool, default True: 869 If `True`, include the leading `'WHERE'` string. 870 871 flavor: str, default 'postgresql' 872 If `connector` is `None`, fall back to this flavor. 873 874 Returns 875 ------- 876 A `str` of the `WHERE` clause from the input `params` dictionary for the connector's flavor. 877 878 Examples 879 -------- 880 ``` 881 >>> print(build_where({'foo': [1, 2, 3]})) 882 883 WHERE 884 "foo" IN ('1', '2', '3') 885 ``` 886 """ 887 _rust = _core.build_where 888 _query_flavor = ( 889 getattr(connector, 'flavor', flavor) if connector is not None else flavor 890 ) 891 return _rust(params, _query_flavor, with_where)
Build the WHERE clause based on the input criteria.
Parameters
- params (Dict[str, Any]:):
The keywords dictionary to convert into a WHERE clause.
If a value is a string which begins with an underscore, negate that value
(e.g.
!=instead of=orNOT INinstead ofIN). A value of_Nonewill be interpreted asIS NOT NULL. - connector (Optional[meerschaum.connectors.sql.SQLConnector], default None:): The Meerschaum SQLConnector that will be executing the query. The connector is used to extract the SQL dialect.
- with_where (bool, default True:):
If
True, include the leading'WHERE'string. - flavor (str, default 'postgresql'):
If
connectorisNone, fall back to this flavor.
Returns
- A
strof theWHEREclause from the inputparamsdictionary for the connector's flavor.
Examples
>>> print(build_where({'foo': [1, 2, 3]}))
WHERE
"foo" IN ('1', '2', '3')
def
table_exists( table: str, connector: meerschaum.connectors.SQLConnector, schema: str | None = None, debug: bool = False) -> bool:
894def table_exists( 895 table: str, 896 connector: mrsm.connectors.sql.SQLConnector, 897 schema: Optional[str] = None, 898 debug: bool = False, 899) -> bool: 900 """Check if a table exists. 901 902 Parameters 903 ---------- 904 table: str: 905 The name of the table in question. 906 907 connector: mrsm.connectors.sql.SQLConnector 908 The connector to the database which holds the table. 909 910 schema: Optional[str], default None 911 Optionally specify the table schema. 912 Defaults to `connector.schema`. 913 914 debug: bool, default False : 915 Verbosity toggle. 916 917 Returns 918 ------- 919 A `bool` indicating whether or not the table exists on the database. 920 """ 921 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 922 from meerschaum.connectors.sql._native import table_exists as native_table_exists 923 schema = schema or connector.schema 924 truncated_table_name = truncate_item_name(str(table), connector.flavor) 925 ### Raw SQL through the Rust connector: no SQLAlchemy (rc15 Q53). 926 exists = native_table_exists(connector, truncated_table_name, schema) 927 if exists is not None: 928 return exists 929 sqlalchemy = import_sqlalchemy() 930 insp = sqlalchemy.inspect(connector.engine) 931 return insp.has_table(truncated_table_name, schema=schema)
Check if a table exists.
Parameters
- table (str:): The name of the table in question.
- connector (mrsm.connectors.sql.SQLConnector): The connector to the database which holds the table.
- schema (Optional[str], default None):
Optionally specify the table schema.
Defaults to
connector.schema. - debug (bool, default False :): Verbosity toggle.
Returns
- A
boolindicating whether or not the table exists on the database.
def
get_sqlalchemy_table( table: str, connector: meerschaum.connectors.SQLConnector | None = None, schema: str | None = None, refresh: bool = False, debug: bool = False) -> "Union['sqlalchemy.Table', None]":
934def get_sqlalchemy_table( 935 table: str, 936 connector: Optional[mrsm.connectors.sql.SQLConnector] = None, 937 schema: Optional[str] = None, 938 refresh: bool = False, 939 debug: bool = False, 940) -> Union['sqlalchemy.Table', None]: 941 """ 942 Construct a SQLAlchemy table from its name. 943 944 Parameters 945 ---------- 946 table: str 947 The name of the table on the database. Does not need to be escaped. 948 949 connector: Optional[meerschaum.connectors.sql.SQLConnector], default None: 950 The connector to the database which holds the table. 951 952 schema: Optional[str], default None 953 Specify on which schema the table resides. 954 Defaults to the schema set in `connector`. 955 956 refresh: bool, default False 957 If `True`, rebuild the cached table object. 958 959 debug: bool, default False: 960 Verbosity toggle. 961 962 Returns 963 ------- 964 A `sqlalchemy.Table` object for the table. 965 966 """ 967 if connector is None: 968 from meerschaum import get_connector 969 connector = get_connector('sql') 970 971 if connector.flavor == 'duckdb': 972 return None 973 974 from meerschaum.connectors.sql.tables import get_tables 975 from meerschaum.utils.packages import attempt_import 976 from meerschaum.utils.warnings import warn 977 if refresh: 978 connector.metadata.clear() 979 tables = get_tables(mrsm_instance=connector, debug=debug, create=False) 980 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 981 sqlalchemy = import_sqlalchemy() 982 truncated_table_name = truncate_item_name(str(table), connector.flavor) 983 table_kwargs = { 984 'autoload_with': connector.engine, 985 } 986 if schema: 987 table_kwargs['schema'] = schema 988 989 if refresh or truncated_table_name not in tables: 990 try: 991 tables[truncated_table_name] = sqlalchemy.Table( 992 truncated_table_name, 993 connector.metadata, 994 **table_kwargs 995 ) 996 except sqlalchemy.exc.NoSuchTableError: 997 warn(f"Table '{truncated_table_name}' does not exist in '{connector}'.") 998 return None 999 return tables[truncated_table_name]
Construct a SQLAlchemy table from its name.
Parameters
- table (str): The name of the table on the database. Does not need to be escaped.
- connector (Optional[meerschaum.connectors.sql.SQLConnector], default None:): The connector to the database which holds the table.
- schema (Optional[str], default None):
Specify on which schema the table resides.
Defaults to the schema set in
connector. - refresh (bool, default False):
If
True, rebuild the cached table object. - debug (bool, default False:): Verbosity toggle.
Returns
- A
sqlalchemy.Tableobject for the table.
def
get_table_cols_types( table: str, connectable: "Union['mrsm.connectors.sql.SQLConnector', 'sqlalchemy.orm.session.Session', 'sqlalchemy.engine.base.Engine']", flavor: str | None = None, schema: str | None = None, database: str | None = None, debug: bool = False) -> Dict[str, str]:
1012def get_table_cols_types( 1013 table: str, 1014 connectable: Union[ 1015 'mrsm.connectors.sql.SQLConnector', 1016 'sqlalchemy.orm.session.Session', 1017 'sqlalchemy.engine.base.Engine' 1018 ], 1019 flavor: Optional[str] = None, 1020 schema: Optional[str] = None, 1021 database: Optional[str] = None, 1022 debug: bool = False, 1023) -> Dict[str, str]: 1024 """ 1025 Return a dictionary mapping a table's columns to data types. 1026 This is useful for inspecting tables creating during a not-yet-committed session. 1027 1028 NOTE: This may return incorrect columns if the schema is not explicitly stated. 1029 Use this function if you are confident the table name is unique or if you have 1030 and explicit schema. 1031 To use the configured schema, get the columns from `get_sqlalchemy_table()` instead. 1032 1033 Parameters 1034 ---------- 1035 table: str 1036 The name of the table (unquoted). 1037 1038 connectable: Union[ 1039 'mrsm.connectors.sql.SQLConnector', 1040 'sqlalchemy.orm.session.Session', 1041 'sqlalchemy.engine.base.Engine' 1042 ] 1043 The connection object used to fetch the columns and types. 1044 1045 flavor: Optional[str], default None 1046 The database dialect flavor to use for the query. 1047 If omitted, default to `connectable.flavor`. 1048 1049 schema: Optional[str], default None 1050 If provided, restrict the query to this schema. 1051 1052 database: Optional[str]. default None 1053 If provided, restrict the query to this database. 1054 1055 Returns 1056 ------- 1057 A dictionary mapping column names to data types. 1058 """ 1059 import textwrap 1060 from meerschaum.connectors import SQLConnector 1061 from meerschaum.connectors.sql._native import sql_text, uses_native 1062 flavor = flavor or getattr(connectable, 'flavor', None) 1063 if not flavor: 1064 raise ValueError("Please provide a database flavor.") 1065 if flavor == 'duckdb' and not isinstance(connectable, SQLConnector): 1066 raise ValueError("You must provide a SQLConnector when using DuckDB.") 1067 native = isinstance(connectable, SQLConnector) and uses_native(connectable) 1068 if flavor in NO_SCHEMA_FLAVORS: 1069 schema = None 1070 if schema is None: 1071 schema = DEFAULT_SCHEMA_FLAVORS.get(flavor, None) 1072 if flavor in ('sqlite', 'duckdb', 'oracle', 'geopackage'): 1073 database = None 1074 table_trunc = truncate_item_name(table, flavor=flavor) 1075 table_lower = table.lower() 1076 table_upper = table.upper() 1077 table_lower_trunc = truncate_item_name(table_lower, flavor=flavor) 1078 table_upper_trunc = truncate_item_name(table_upper, flavor=flavor) 1079 db_prefix = ( 1080 "tempdb." 1081 if flavor == 'mssql' and table.startswith('#') 1082 else "" 1083 ) 1084 1085 def _esc(s: str) -> str: 1086 return s.replace("'", "''") 1087 1088 cols_types_query = sql_text( 1089 connectable, 1090 textwrap.dedent(columns_types_queries.get( 1091 flavor, 1092 columns_types_queries['default'] 1093 ).format( 1094 table=_esc(table), 1095 table_trunc=_esc(table_trunc), 1096 table_lower=_esc(table_lower), 1097 table_lower_trunc=_esc(table_lower_trunc), 1098 table_upper=_esc(table_upper), 1099 table_upper_trunc=_esc(table_upper_trunc), 1100 db_prefix=db_prefix, 1101 )).lstrip().rstrip() 1102 ) 1103 1104 cols = ['database', 'schema', 'table', 'column', 'type', 'numeric_precision', 'numeric_scale'] 1105 result_cols_ix = dict(enumerate(cols)) 1106 1107 debug_kwargs = {'debug': debug} if isinstance(connectable, SQLConnector) else {} 1108 if not debug_kwargs and debug: 1109 dprint(cols_types_query) 1110 1111 try: 1112 result_rows = ( 1113 [ 1114 row 1115 for row in connectable.execute(cols_types_query, **debug_kwargs).fetchall() 1116 ] 1117 if flavor != 'duckdb' or native 1118 else [ 1119 tuple([doc[col] for col in cols]) 1120 for doc in _read_records(connectable, cols_types_query, debug=debug) 1121 ] 1122 ) 1123 cols_types_docs = [ 1124 { 1125 result_cols_ix[i]: val 1126 for i, val in enumerate(row) 1127 } 1128 for row in result_rows 1129 ] 1130 cols_types_docs_filtered = [ 1131 doc 1132 for doc in cols_types_docs 1133 if ( 1134 ( 1135 not schema 1136 or doc['schema'] == schema 1137 ) 1138 and 1139 ( 1140 not database 1141 or doc['database'] == database 1142 ) 1143 ) 1144 ] 1145 1146 ### NOTE: This may return incorrect columns if the schema is not explicitly stated. 1147 if cols_types_docs and not cols_types_docs_filtered: 1148 cols_types_docs_filtered = cols_types_docs 1149 1150 ### NOTE: Check for PostGIS GEOMETRY columns. 1151 geometry_cols_types = {} 1152 user_defined_cols = [ 1153 doc 1154 for doc in cols_types_docs_filtered 1155 if str(doc.get('type', None)).upper() == 'USER-DEFINED' 1156 ] 1157 if user_defined_cols: 1158 geometry_cols_types.update( 1159 get_postgis_geo_columns_types( 1160 connectable, 1161 table, 1162 schema=schema, 1163 debug=debug, 1164 ) 1165 ) 1166 1167 cols_types = { 1168 ( 1169 doc['column'] 1170 if flavor != 'oracle' else ( 1171 ( 1172 doc['column'].lower() 1173 if (doc['column'].isupper() and doc['column'].replace('_', '').isalpha()) 1174 else doc['column'] 1175 ) 1176 ) 1177 ): doc['type'].upper() + ( 1178 f'({precision},{scale})' 1179 if ( 1180 (precision := doc.get('numeric_precision', None)) 1181 and 1182 (scale := doc.get('numeric_scale', None)) 1183 ) 1184 else '' 1185 ) 1186 for doc in cols_types_docs_filtered 1187 } 1188 cols_types.update(geometry_cols_types) 1189 return cols_types 1190 except Exception as e: 1191 warn(f"Failed to fetch columns for table '{table}':\n{e}") 1192 return {}
Return a dictionary mapping a table's columns to data types. This is useful for inspecting tables creating during a not-yet-committed session.
NOTE: This may return incorrect columns if the schema is not explicitly stated.
Use this function if you are confident the table name is unique or if you have
and explicit schema.
To use the configured schema, get the columns from get_sqlalchemy_table() instead.
Parameters
- table (str): The name of the table (unquoted).
- connectable (Union[): 'mrsm.connectors.sql.SQLConnector', 'sqlalchemy.orm.session.Session', 'sqlalchemy.engine.base.Engine'
- ]: The connection object used to fetch the columns and types.
- flavor (Optional[str], default None):
The database dialect flavor to use for the query.
If omitted, default to
connectable.flavor. - schema (Optional[str], default None): If provided, restrict the query to this schema.
- database (Optional[str]. default None): If provided, restrict the query to this database.
Returns
- A dictionary mapping column names to data types.
def
get_table_cols_indices( table: str, connectable: "Union['mrsm.connectors.sql.SQLConnector', 'sqlalchemy.orm.session.Session', 'sqlalchemy.engine.base.Engine']", flavor: str | None = None, schema: str | None = None, database: str | None = None, debug: bool = False) -> Dict[str, List[str]]:
1195def get_table_cols_indices( 1196 table: str, 1197 connectable: Union[ 1198 'mrsm.connectors.sql.SQLConnector', 1199 'sqlalchemy.orm.session.Session', 1200 'sqlalchemy.engine.base.Engine' 1201 ], 1202 flavor: Optional[str] = None, 1203 schema: Optional[str] = None, 1204 database: Optional[str] = None, 1205 debug: bool = False, 1206) -> Dict[str, List[str]]: 1207 """ 1208 Return a dictionary mapping a table's columns to lists of indices. 1209 This is useful for inspecting tables creating during a not-yet-committed session. 1210 1211 NOTE: This may return incorrect columns if the schema is not explicitly stated. 1212 Use this function if you are confident the table name is unique or if you have 1213 and explicit schema. 1214 To use the configured schema, get the columns from `get_sqlalchemy_table()` instead. 1215 1216 Parameters 1217 ---------- 1218 table: str 1219 The name of the table (unquoted). 1220 1221 connectable: Union[ 1222 'mrsm.connectors.sql.SQLConnector', 1223 'sqlalchemy.orm.session.Session', 1224 'sqlalchemy.engine.base.Engine' 1225 ] 1226 The connection object used to fetch the columns and types. 1227 1228 flavor: Optional[str], default None 1229 The database dialect flavor to use for the query. 1230 If omitted, default to `connectable.flavor`. 1231 1232 schema: Optional[str], default None 1233 If provided, restrict the query to this schema. 1234 1235 database: Optional[str]. default None 1236 If provided, restrict the query to this database. 1237 1238 Returns 1239 ------- 1240 A dictionary mapping column names to a list of indices. 1241 """ 1242 import textwrap 1243 from collections import defaultdict 1244 from meerschaum.connectors import SQLConnector 1245 from meerschaum.connectors.sql._native import sql_text, uses_native 1246 flavor = flavor or getattr(connectable, 'flavor', None) 1247 if not flavor: 1248 raise ValueError("Please provide a database flavor.") 1249 if flavor == 'duckdb' and not isinstance(connectable, SQLConnector): 1250 raise ValueError("You must provide a SQLConnector when using DuckDB.") 1251 native = isinstance(connectable, SQLConnector) and uses_native(connectable) 1252 if flavor in NO_SCHEMA_FLAVORS: 1253 schema = None 1254 if schema is None: 1255 schema = DEFAULT_SCHEMA_FLAVORS.get(flavor, None) 1256 if flavor in ('sqlite', 'duckdb', 'oracle', 'geopackage'): 1257 database = None 1258 table_trunc = truncate_item_name(table, flavor=flavor) 1259 table_lower = table.lower() 1260 table_upper = table.upper() 1261 table_lower_trunc = truncate_item_name(table_lower, flavor=flavor) 1262 table_upper_trunc = truncate_item_name(table_upper, flavor=flavor) 1263 db_prefix = ( 1264 "tempdb." 1265 if flavor == 'mssql' and table.startswith('#') 1266 else "" 1267 ) 1268 1269 def _esc(s: str) -> str: 1270 return s.replace("'", "''") 1271 1272 cols_indices_query = sql_text( 1273 connectable, 1274 textwrap.dedent(columns_indices_queries.get( 1275 flavor, 1276 columns_indices_queries['default'] 1277 ).format( 1278 table=_esc(table), 1279 table_trunc=_esc(table_trunc), 1280 table_lower=_esc(table_lower), 1281 table_lower_trunc=_esc(table_lower_trunc), 1282 table_upper=_esc(table_upper), 1283 table_upper_trunc=_esc(table_upper_trunc), 1284 db_prefix=db_prefix, 1285 schema=_esc(schema) if schema else schema, 1286 )).lstrip().rstrip() 1287 ) 1288 1289 cols = ['database', 'schema', 'table', 'column', 'index', 'index_type'] 1290 if flavor == 'mssql': 1291 cols.append('clustered') 1292 result_cols_ix = dict(enumerate(cols)) 1293 1294 debug_kwargs = {'debug': debug} if isinstance(connectable, SQLConnector) else {} 1295 if not debug_kwargs and debug: 1296 dprint(cols_indices_query) 1297 1298 try: 1299 result_rows = ( 1300 [ 1301 row 1302 for row in connectable.execute(cols_indices_query, **debug_kwargs).fetchall() 1303 ] 1304 if flavor != 'duckdb' or native 1305 else [ 1306 tuple([doc[col] for col in cols]) 1307 for doc in _read_records(connectable, cols_indices_query, debug=debug) 1308 ] 1309 ) 1310 cols_types_docs = [ 1311 { 1312 result_cols_ix[i]: val 1313 for i, val in enumerate(row) 1314 } 1315 for row in result_rows 1316 ] 1317 cols_types_docs_filtered = [ 1318 doc 1319 for doc in cols_types_docs 1320 if ( 1321 ( 1322 not schema 1323 or doc['schema'] == schema 1324 ) 1325 and 1326 ( 1327 not database 1328 or doc['database'] == database 1329 ) 1330 ) 1331 ] 1332 ### NOTE: This may return incorrect columns if the schema is not explicitly stated. 1333 if cols_types_docs and not cols_types_docs_filtered: 1334 cols_types_docs_filtered = cols_types_docs 1335 1336 cols_indices = defaultdict(lambda: []) 1337 for doc in cols_types_docs_filtered: 1338 col = ( 1339 doc['column'] 1340 if flavor != 'oracle' 1341 else ( 1342 doc['column'].lower() 1343 if (doc['column'].isupper() and doc['column'].replace('_', '').isalpha()) 1344 else doc['column'] 1345 ) 1346 ) 1347 index_doc = { 1348 'name': doc.get('index', None), 1349 'type': doc.get('index_type', None) 1350 } 1351 if flavor == 'mssql': 1352 index_doc['clustered'] = doc.get('clustered', None) 1353 cols_indices[col].append(index_doc) 1354 1355 return dict(cols_indices) 1356 except Exception as e: 1357 warn(f"Failed to fetch columns for table '{table}':\n{e}") 1358 return {}
Return a dictionary mapping a table's columns to lists of indices. This is useful for inspecting tables creating during a not-yet-committed session.
NOTE: This may return incorrect columns if the schema is not explicitly stated.
Use this function if you are confident the table name is unique or if you have
and explicit schema.
To use the configured schema, get the columns from get_sqlalchemy_table() instead.
Parameters
- table (str): The name of the table (unquoted).
- connectable (Union[): 'mrsm.connectors.sql.SQLConnector', 'sqlalchemy.orm.session.Session', 'sqlalchemy.engine.base.Engine'
- ]: The connection object used to fetch the columns and types.
- flavor (Optional[str], default None):
The database dialect flavor to use for the query.
If omitted, default to
connectable.flavor. - schema (Optional[str], default None): If provided, restrict the query to this schema.
- database (Optional[str]. default None): If provided, restrict the query to this database.
Returns
- A dictionary mapping column names to a list of indices.
def
get_db_version(conn: "'SQLConnector'", debug: bool = False) -> str | None:
1361def get_db_version(conn: 'SQLConnector', debug: bool = False) -> Union[str, None]: 1362 """ 1363 Fetch the database version if possible. 1364 """ 1365 version_name = sql_item_name('version', conn.flavor, None) 1366 version_query = version_queries.get( 1367 conn.flavor, 1368 version_queries['default'] 1369 ).format(version_name=version_name) 1370 return conn.value(version_query, debug=debug, silent=True)
Fetch the database version if possible.
PG_PARTITION_FLAVORS =
{'postgis', 'postgresql'}
MYSQL_PARTITION_FLAVORS =
{'mariadb', 'mysql'}
def
find_top_level_select_index(query: str, last: bool = False) -> int:
1378def find_top_level_select_index(query: str, last: bool = False) -> int: 1379 """ 1380 Return the index of the first (or last) top-level `SELECT` keyword in a query. 1381 1382 The scan skips `--` line comments, `/* */` block comments, string literals, 1383 quoted identifiers, and parenthesized subqueries, so a definition may safely 1384 contain comments (including comments mentioning "select"). 1385 1386 Parameters 1387 ---------- 1388 query: str 1389 The SQL query to scan. 1390 1391 last: bool, default False 1392 If `True`, return the index of the last top-level `SELECT` 1393 (e.g. the final statement of a CTE query) instead of the first. 1394 1395 Returns 1396 ------- 1397 The character index of the `SELECT` keyword, or `-1` if none exists 1398 at the top level. 1399 """ 1400 q_lower = query.lower() 1401 n = len(q_lower) 1402 depth = 0 1403 found_ix = -1 1404 i = 0 1405 while i < n: 1406 ch = q_lower[i] 1407 nxt = q_lower[i + 1] if (i + 1) < n else '' 1408 if ch == '-' and nxt == '-': 1409 eol = q_lower.find('\n', i) 1410 i = n if eol == -1 else (eol + 1) 1411 continue 1412 if ch == '/' and nxt == '*': 1413 end = q_lower.find('*/', i + 2) 1414 i = n if end == -1 else (end + 2) 1415 continue 1416 if ch == "'": 1417 i += 1 1418 while i < n: 1419 if q_lower[i] == "'": 1420 if (i + 1) < n and q_lower[i + 1] == "'": 1421 i += 2 1422 continue 1423 break 1424 i += 1 1425 i += 1 1426 continue 1427 if ch in ('"', '`'): 1428 end = q_lower.find(ch, i + 1) 1429 i = n if end == -1 else (end + 1) 1430 continue 1431 if ch == '[': 1432 end = q_lower.find(']', i + 1) 1433 i = n if end == -1 else (end + 1) 1434 continue 1435 if ch == '(': 1436 depth += 1 1437 i += 1 1438 continue 1439 if ch == ')': 1440 depth -= 1 1441 i += 1 1442 continue 1443 if depth == 0 and ch == 's' and q_lower[i:i + 6] == 'select': 1444 before_ok = i == 0 or not (q_lower[i - 1].isalnum() or q_lower[i - 1] == '_') 1445 after_ok = (i + 6) == n or not (q_lower[i + 6].isalnum() or q_lower[i + 6] == '_') 1446 if before_ok and after_ok: 1447 if not last: 1448 return i 1449 found_ix = i 1450 i += 6 1451 continue 1452 i += 1 1453 return found_ix
Return the index of the first (or last) top-level SELECT keyword in a query.
The scan skips -- line comments, /* */ block comments, string literals,
quoted identifiers, and parenthesized subqueries, so a definition may safely
contain comments (including comments mentioning "select").
Parameters
- query (str): The SQL query to scan.
- last (bool, default False):
If
True, return the index of the last top-levelSELECT(e.g. the final statement of a CTE query) instead of the first.
Returns
- The character index of the
SELECTkeyword, or-1if none exists - at the top level.
def
get_leading_with_index(query: str) -> int:
1456def get_leading_with_index(query: str) -> int: 1457 """ 1458 Return the index of a query's leading `WITH` keyword, 1459 skipping whitespace and comments, or `-1` if the query 1460 does not begin with `WITH`. 1461 """ 1462 q_lower = query.lower() 1463 n = len(q_lower) 1464 i = 0 1465 while i < n: 1466 ch = q_lower[i] 1467 nxt = q_lower[i + 1] if (i + 1) < n else '' 1468 if ch in (' ', '\t', '\r', '\n'): 1469 i += 1 1470 continue 1471 if ch == '-' and nxt == '-': 1472 eol = q_lower.find('\n', i) 1473 i = n if eol == -1 else (eol + 1) 1474 continue 1475 if ch == '/' and nxt == '*': 1476 end = q_lower.find('*/', i + 2) 1477 i = n if end == -1 else (end + 2) 1478 continue 1479 break 1480 if q_lower[i:i + 4] == 'with' and ( 1481 (i + 4) == n or not (q_lower[i + 4].isalnum() or q_lower[i + 4] == '_') 1482 ): 1483 return i 1484 return -1
Return the index of a query's leading WITH keyword,
skipping whitespace and comments, or -1 if the query
does not begin with WITH.
def
wrap_query_with_cte( sub_query: str, parent_query: str, flavor: str, cte_name: str = 'src') -> str:
1487def wrap_query_with_cte( 1488 sub_query: str, 1489 parent_query: str, 1490 flavor: str, 1491 cte_name: str = "src", 1492) -> str: 1493 """ 1494 Wrap a subquery in a CTE and append an encapsulating query. 1495 1496 Parameters 1497 ---------- 1498 sub_query: str 1499 The query to be referenced. This may itself contain CTEs and comments. 1500 Unless `cte_name` is provided, this will be aliased as `src`. 1501 1502 parent_query: str 1503 The larger query to append which references the subquery. 1504 If it contains CTEs of its own, they are hoisted into the 1505 enclosing `WITH` list. 1506 1507 flavor: str 1508 The database flavor, e.g. `'mssql'`. 1509 1510 cte_name: str, default 'src' 1511 The CTE alias, defaults to `src`. 1512 1513 Returns 1514 ------- 1515 An encapsulating query which allows you to treat `sub_query` as a temporary table. 1516 Raises an exception when a query's top-level `SELECT` cannot be located 1517 (e.g. an incomplete definition), rather than emitting invalid SQL. 1518 1519 Examples 1520 -------- 1521 1522 >>> from meerschaum.utils.sql import wrap_query_with_cte 1523 >>> sub_query = "WITH foo AS (SELECT 1 AS val) SELECT (val * 2) AS newval FROM foo" 1524 >>> parent_query = "SELECT newval * 3 FROM src" 1525 >>> query = wrap_query_with_cte(sub_query, parent_query, 'mssql') 1526 >>> print(query) 1527 >>> # WITH foo AS (SELECT 1 AS val) 1528 >>> # , [src] AS ( 1529 >>> # SELECT (val * 2) AS newval FROM foo 1530 >>> # ) 1531 >>> # SELECT newval * 3 FROM src 1532 1533 """ 1534 import textwrap 1535 from meerschaum.utils.warnings import error 1536 sub_query = sub_query.lstrip() 1537 parent_query = parent_query.lstrip() 1538 cte_name_quoted = sql_item_name(cte_name, flavor, None) 1539 1540 if flavor in NO_CTE_FLAVORS: 1541 return ( 1542 parent_query 1543 .replace(cte_name_quoted, '--MRSM_SUBQUERY--') 1544 .replace(cte_name, '--MRSM_SUBQUERY--') 1545 .replace('--MRSM_SUBQUERY--', f"(\n{sub_query}\n) AS {cte_name_quoted}") 1546 ) 1547 1548 if get_leading_with_index(sub_query) != -1: 1549 select_index = find_top_level_select_index(sub_query) 1550 if select_index == -1: 1551 error( 1552 "Could not locate the top-level SELECT of the following query:\n" 1553 + sub_query 1554 ) 1555 1556 ### Flatten: the definitions part keeps the sub-query's own CTEs 1557 ### (`WITH a AS (...) -- trailing comments included`), and the body 1558 ### (`SELECT ...`) becomes the new CTE. The joining comma goes on its 1559 ### own line so a trailing `--` comment cannot swallow it. 1560 definitions = sub_query[:select_index].rstrip() 1561 if definitions.endswith(','): 1562 definitions = definitions[:-1].rstrip() 1563 body = sub_query[select_index:].strip() 1564 1565 lead = ( 1566 f"{definitions}\n" 1567 f", {cte_name_quoted} AS (\n" 1568 f"{textwrap.indent(body, ' ')}\n" 1569 f")" 1570 ) 1571 else: 1572 lead = ( 1573 f"WITH {cte_name_quoted} AS (\n" 1574 f"{textwrap.indent(sub_query, ' ')}\n" 1575 f")" 1576 ) 1577 1578 parent_with_index = get_leading_with_index(parent_query) 1579 if parent_with_index == -1: 1580 return f"{lead}\n{parent_query}" 1581 1582 ### The parent query declares CTEs of its own: hoist them after the new 1583 ### definition instead of emitting an invalid second `WITH`. 1584 parent_prefix = parent_query[:parent_with_index].rstrip() 1585 parent_rest = parent_query[parent_with_index + len('with'):].lstrip() 1586 return ( 1587 (f"{parent_prefix}\n" if parent_prefix else "") 1588 + f"{lead}\n, {parent_rest}" 1589 )
Wrap a subquery in a CTE and append an encapsulating query.
Parameters
- sub_query (str):
The query to be referenced. This may itself contain CTEs and comments.
Unless
cte_nameis provided, this will be aliased assrc. - parent_query (str):
The larger query to append which references the subquery.
If it contains CTEs of its own, they are hoisted into the
enclosing
WITHlist. - flavor (str):
The database flavor, e.g.
'mssql'. - cte_name (str, default 'src'):
The CTE alias, defaults to
src.
Returns
- An encapsulating query which allows you to treat
sub_queryas a temporary table. - Raises an exception when a query's top-level
SELECTcannot be located - (e.g. an incomplete definition), rather than emitting invalid SQL.
Examples
>>> from meerschaum.utils.sql import wrap_query_with_cte
>>> sub_query = "WITH foo AS (SELECT 1 AS val) SELECT (val * 2) AS newval FROM foo"
>>> parent_query = "SELECT newval * 3 FROM src"
>>> query = wrap_query_with_cte(sub_query, parent_query, 'mssql')
>>> print(query)
>>> # WITH foo AS (SELECT 1 AS val)
>>> # , [src] AS (
>>> # SELECT (val * 2) AS newval FROM foo
>>> # )
>>> # SELECT newval * 3 FROM src
def
format_cte_subquery( sub_query: str, flavor: str, sub_name: str = 'src', cols_to_select: List[str] | str = '*') -> str:
1592def format_cte_subquery( 1593 sub_query: str, 1594 flavor: str, 1595 sub_name: str = 'src', 1596 cols_to_select: Union[List[str], str] = '*', 1597) -> str: 1598 """ 1599 Given a subquery, build a wrapper query that selects from the CTE subquery. 1600 1601 Parameters 1602 ---------- 1603 sub_query: str 1604 The subquery to wrap. 1605 1606 flavor: str 1607 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`. 1608 1609 sub_name: str, default 'src' 1610 If possible, give this name to the CTE (must be unquoted). 1611 1612 cols_to_select: Union[List[str], str], default '' 1613 If specified, choose which columns to select from the CTE. 1614 If a list of strings is provided, each item will be quoted and joined with commas. 1615 If a string is given, assume it is quoted and insert it into the query. 1616 1617 Returns 1618 ------- 1619 A wrapper query that selects from the CTE. 1620 """ 1621 quoted_sub_name = sql_item_name(sub_name, flavor, None) 1622 cols_str = ( 1623 cols_to_select 1624 if isinstance(cols_to_select, str) 1625 else ', '.join([sql_item_name(col, flavor, None) for col in cols_to_select]) 1626 ) 1627 parent_query = ( 1628 f"SELECT {cols_str}\n" 1629 f"FROM {quoted_sub_name}" 1630 ) 1631 return wrap_query_with_cte(sub_query, parent_query, flavor, cte_name=sub_name)
Given a subquery, build a wrapper query that selects from the CTE subquery.
Parameters
- sub_query (str): The subquery to wrap.
- flavor (str):
The database flavor to use for the query (e.g.
'mssql','postgresql'. - sub_name (str, default 'src'): If possible, give this name to the CTE (must be unquoted).
- cols_to_select (Union[List[str], str], default ''): If specified, choose which columns to select from the CTE. If a list of strings is provided, each item will be quoted and joined with commas. If a string is given, assume it is quoted and insert it into the query.
Returns
- A wrapper query that selects from the CTE.
def
session_execute( session: "'sqlalchemy.orm.session.Session'", queries: List[str] | str, with_results: bool = False, debug: bool = False) -> "Union[mrsm.SuccessTuple, Tuple[mrsm.SuccessTuple, List['sqlalchemy.sql.ResultProxy']]]":
1634def session_execute( 1635 session: 'sqlalchemy.orm.session.Session', 1636 queries: Union[List[str], str], 1637 with_results: bool = False, 1638 debug: bool = False, 1639) -> Union[mrsm.SuccessTuple, Tuple[mrsm.SuccessTuple, List['sqlalchemy.sql.ResultProxy']]]: 1640 """ 1641 Similar to `SQLConnector.exec_queries()`, execute a list of queries 1642 and roll back when one fails. 1643 1644 Parameters 1645 ---------- 1646 session: sqlalchemy.orm.session.Session 1647 A SQLAlchemy session representing a transaction. 1648 1649 queries: Union[List[str], str] 1650 A query or list of queries to be executed. 1651 If a query fails, roll back the session. 1652 1653 with_results: bool, default False 1654 If `True`, return a list of result objects. 1655 1656 Returns 1657 ------- 1658 A `SuccessTuple` indicating the queries were successfully executed. 1659 If `with_results`, return the `SuccessTuple` and a list of results. 1660 """ 1661 from meerschaum.connectors.sql._create_engine import import_sqlalchemy 1662 sqlalchemy = import_sqlalchemy() 1663 if not isinstance(queries, list): 1664 queries = [queries] 1665 successes, msgs, results = [], [], [] 1666 for query in queries: 1667 if debug: 1668 dprint(query) 1669 query_text = sqlalchemy.text(query) 1670 fail_msg = "Failed to execute queries." 1671 try: 1672 result = session.execute(query_text) 1673 query_success = result is not None 1674 query_msg = "Success" if query_success else fail_msg 1675 except Exception as e: 1676 query_success = False 1677 query_msg = f"{fail_msg}\n{e}" 1678 result = None 1679 successes.append(query_success) 1680 msgs.append(query_msg) 1681 results.append(result) 1682 if not query_success: 1683 if debug: 1684 dprint("Rolling back session.") 1685 session.rollback() 1686 break 1687 success, msg = all(successes), '\n'.join(msgs) 1688 if with_results: 1689 return (success, msg), results 1690 return success, msg
Similar to SQLConnector.exec_queries(), execute a list of queries
and roll back when one fails.
Parameters
- session (sqlalchemy.orm.session.Session): A SQLAlchemy session representing a transaction.
- queries (Union[List[str], str]): A query or list of queries to be executed. If a query fails, roll back the session.
- with_results (bool, default False):
If
True, return a list of result objects.
Returns
- A
SuccessTupleindicating the queries were successfully executed. - If
with_results, return theSuccessTupleand a list of results.
def
get_postgis_geo_columns_types( connectable: "Union['mrsm.connectors.sql.SQLConnector', 'sqlalchemy.orm.session.Session', 'sqlalchemy.engine.base.Engine']", table: str, schema: str | None = 'public', debug: bool = False) -> Dict[str, str]:
1693def get_postgis_geo_columns_types( 1694 connectable: Union[ 1695 'mrsm.connectors.sql.SQLConnector', 1696 'sqlalchemy.orm.session.Session', 1697 'sqlalchemy.engine.base.Engine' 1698 ], 1699 table: str, 1700 schema: Optional[str] = 'public', 1701 debug: bool = False, 1702) -> Dict[str, str]: 1703 """ 1704 Return a dictionary mapping PostGIS geometry column names to geometry types. 1705 1706 Parameters 1707 ---------- 1708 connectable: Union[SQLConnector, sqlalchemy.orm.session.Session, sqlalchemy.engine.base.Engine] 1709 The SQLConnector, Session, or Engine. 1710 1711 table: str 1712 The table's name. 1713 1714 schema: Optional[str], default 'public' 1715 The schema to use for the fully qualified table name. 1716 1717 Returns 1718 ------- 1719 A dictionary mapping column names to geometry types. 1720 """ 1721 from meerschaum.utils.dtypes import get_geometry_type_srid 1722 from meerschaum.connectors.sql._native import sql_text 1723 default_type, default_srid = get_geometry_type_srid() 1724 default_type = default_type.upper() 1725 1726 clean(table) 1727 clean(str(schema)) 1728 schema = schema or 'public' 1729 truncated_schema_name = truncate_item_name(schema, flavor='postgis') 1730 truncated_table_name = truncate_item_name(table, flavor='postgis') 1731 query = sql_text( 1732 connectable, 1733 "SELECT \"f_geometry_column\" AS \"column\", 'GEOMETRY' AS \"func\", \"type\", geom.\"srid\"\n" 1734 "FROM \"geometry_columns\" AS geom\n" 1735 f"WHERE \"f_table_schema\" = '{truncated_schema_name}'\n" 1736 f" AND \"f_table_name\" = '{truncated_table_name}'\n" 1737 "UNION ALL\n" 1738 "SELECT \"f_geography_column\" AS \"column\", 'GEOGRAPHY' AS \"func\", \"type\", geog.\"srid\"\n" 1739 "FROM \"geography_columns\" AS geog\n" 1740 f"WHERE \"f_table_schema\" = '{truncated_schema_name}'\n" 1741 f" AND \"f_table_name\" = '{truncated_table_name}'\n" 1742 ) 1743 debug_kwargs = {'debug': debug} if isinstance(connectable, mrsm.connectors.SQLConnector) else {} 1744 result_rows = [ 1745 row 1746 for row in connectable.execute(query, **debug_kwargs).fetchall() 1747 ] 1748 cols_type_tuples = { 1749 row[0]: ( 1750 row[1], # func 1751 row[2], # type 1752 ( 1753 row[3] 1754 if row[3] 1755 else 0 1756 ) # srid 1757 ) 1758 for row in result_rows 1759 } 1760 1761 geometry_cols_types = { 1762 col: ( 1763 f"{func}({typ.upper()}, {srid})" 1764 if srid 1765 else ( 1766 func 1767 + ( 1768 f'({typ.upper()})' 1769 if typ.upper() not in ('GEOMETRY', 'GEOGRAPHY') 1770 else '' 1771 ) 1772 ) 1773 ) 1774 for col, (func, typ, srid) in cols_type_tuples.items() 1775 } 1776 return geometry_cols_types
Return a dictionary mapping PostGIS geometry column names to geometry types.
Parameters
- connectable (Union[SQLConnector, sqlalchemy.orm.session.Session, sqlalchemy.engine.base.Engine]): The SQLConnector, Session, or Engine.
- table (str): The table's name.
- schema (Optional[str], default 'public'): The schema to use for the fully qualified table name.
Returns
- A dictionary mapping column names to geometry types.
UPDATE_QUERIES =
{'default': '\n UPDATE {target_table_name} AS f\n {sets_subquery_none}\n FROM {target_table_name} AS t\n INNER JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p\n ON\n {and_subquery_t}\n WHERE\n {and_subquery_f}\n AND\n {date_bounds_subquery}\n ', 'timescaledb-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'timescaledb-ha-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'postgresql-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'postgis-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'citus-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'cockroachdb-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'mysql': '\n UPDATE {target_table_name} AS f\n JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p\n ON\n {and_subquery_f}\n {sets_subquery_f}\n WHERE\n {date_bounds_subquery}\n ', 'mysql-upsert': '\n INSERT {ignore}INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n {on_duplicate_key_update}\n {cols_equal_values}\n ', 'mariadb': '\n UPDATE {target_table_name} AS f\n JOIN (SELECT {patch_cols_str} FROM {patch_table_name}) AS p\n ON\n {and_subquery_f}\n {sets_subquery_f}\n WHERE\n {date_bounds_subquery}\n ', 'mariadb-upsert': '\n INSERT {ignore}INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n {on_duplicate_key_update}\n {cols_equal_values}\n ', 'mssql': '\n {with_temp_date_bounds}\n MERGE {target_table_name} f\n USING (SELECT {patch_cols_str} FROM {patch_table_name}) p\n ON\n {and_subquery_f}\n AND\n {date_bounds_subquery}\n WHEN MATCHED THEN\n UPDATE\n {sets_subquery_none};\n ', 'mssql-upsert': ['{identity_insert_on}', '\n {with_temp_date_bounds}\n MERGE {target_table_name} f\n USING (SELECT {patch_cols_str} FROM {patch_table_name}) p\n ON\n {and_subquery_f}\n AND\n {date_bounds_subquery}{when_matched_update_sets_subquery_none}\n WHEN NOT MATCHED THEN\n INSERT ({patch_cols_str})\n VALUES ({patch_cols_prefixed_str});\n ', '{identity_insert_off}'], 'oracle': '\n MERGE INTO {target_table_name} f\n USING (SELECT {patch_cols_str} FROM {patch_table_name}) p\n ON (\n {and_subquery_f}\n AND\n {date_bounds_subquery}\n )\n WHEN MATCHED THEN\n UPDATE\n {sets_subquery_none}\n ', 'oracle-upsert': '\n MERGE INTO {target_table_name} f\n USING (SELECT {patch_cols_str} FROM {patch_table_name}) p\n ON (\n {and_subquery_f}\n AND\n {date_bounds_subquery}\n ){when_matched_update_sets_subquery_none}\n WHEN NOT MATCHED THEN\n INSERT ({patch_cols_str})\n VALUES ({patch_cols_prefixed_str})\n ', 'sqlite-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n WHERE true\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n ', 'sqlite_delete_insert': ['\n DELETE FROM {target_table_name} AS f\n WHERE ROWID IN (\n SELECT t.ROWID\n FROM {target_table_name} AS t\n INNER JOIN (SELECT * FROM {patch_table_name}) AS p\n ON {and_subquery_t}\n );\n ', '\n INSERT INTO {target_table_name} AS f\n SELECT {patch_cols_str} FROM {patch_table_name} AS p\n '], 'geopackage-upsert': '\n INSERT INTO {target_table_name} ({patch_cols_str})\n SELECT {patch_cols_str}\n FROM {patch_table_name}\n WHERE true\n ON CONFLICT ({join_cols_str}) DO {update_or_nothing} {sets_subquery_none_excluded}\n '}
def
get_update_queries( target: str, patch: str, connectable: "Union[mrsm.connectors.sql.SQLConnector, 'sqlalchemy.orm.session.Session']", join_cols: Iterable[str], flavor: str | None = None, upsert: bool = False, datetime_col: str | None = None, schema: str | None = None, patch_schema: str | None = None, target_cols_types: Dict[str, str] | None = None, patch_cols_types: Dict[str, str] | None = None, identity_insert: bool = False, null_indices: bool = True, cast_columns: bool = True, debug: bool = False) -> List[str]:
1945def get_update_queries( 1946 target: str, 1947 patch: str, 1948 connectable: Union[ 1949 mrsm.connectors.sql.SQLConnector, 1950 'sqlalchemy.orm.session.Session' 1951 ], 1952 join_cols: Iterable[str], 1953 flavor: Optional[str] = None, 1954 upsert: bool = False, 1955 datetime_col: Optional[str] = None, 1956 schema: Optional[str] = None, 1957 patch_schema: Optional[str] = None, 1958 target_cols_types: Optional[Dict[str, str]] = None, 1959 patch_cols_types: Optional[Dict[str, str]] = None, 1960 identity_insert: bool = False, 1961 null_indices: bool = True, 1962 cast_columns: bool = True, 1963 debug: bool = False, 1964) -> List[str]: 1965 """ 1966 Build a list of `MERGE`, `UPDATE`, `DELETE`/`INSERT` queries to apply a patch to target table. 1967 1968 Parameters 1969 ---------- 1970 target: str 1971 The name of the target table. 1972 1973 patch: str 1974 The name of the patch table. This should have the same shape as the target. 1975 1976 connectable: Union[meerschaum.connectors.sql.SQLConnector, sqlalchemy.orm.session.Session] 1977 The `SQLConnector` or SQLAlchemy session which will later execute the queries. 1978 1979 join_cols: List[str] 1980 The columns to use to join the patch to the target. 1981 1982 flavor: Optional[str], default None 1983 If using a SQLAlchemy session, provide the expected database flavor. 1984 1985 upsert: bool, default False 1986 If `True`, return an upsert query rather than an update. 1987 1988 datetime_col: Optional[str], default None 1989 If provided, bound the join query using this column as the datetime index. 1990 This must be present on both tables. 1991 1992 schema: Optional[str], default None 1993 If provided, use this schema when quoting the target table. 1994 Defaults to `connector.schema`. 1995 1996 patch_schema: Optional[str], default None 1997 If provided, use this schema when quoting the patch table. 1998 Defaults to `schema`. 1999 2000 target_cols_types: Optional[Dict[str, Any]], default None 2001 If provided, use these as the columns-types dictionary for the target table. 2002 Default will infer from the database context. 2003 2004 patch_cols_types: Optional[Dict[str, Any]], default None 2005 If provided, use these as the columns-types dictionary for the target table. 2006 Default will infer from the database context. 2007 2008 identity_insert: bool, default False 2009 If `True`, include `SET IDENTITY_INSERT` queries before and after the update queries. 2010 Only applies for MSSQL upserts. 2011 2012 null_indices: bool, default True 2013 If `False`, do not coalesce index columns before joining. 2014 2015 cast_columns: bool, default True 2016 If `False`, do not cast update columns to the target table types. 2017 2018 debug: bool, default False 2019 Verbosity toggle. 2020 2021 Returns 2022 ------- 2023 A list of query strings to perform the update operation. 2024 """ 2025 import textwrap 2026 from meerschaum.connectors import SQLConnector 2027 from meerschaum.utils.debug import dprint 2028 from meerschaum.utils.dtypes import are_dtypes_equal 2029 from meerschaum.utils.dtypes.sql import DB_FLAVORS_CAST_DTYPES, get_pd_type_from_db_type 2030 flavor = flavor or getattr(connectable, 'flavor', None) 2031 if not flavor: 2032 raise ValueError("Provide a flavor if using a SQLAlchemy session.") 2033 if ( 2034 flavor in ('sqlite', 'geopackage') 2035 and isinstance(connectable, SQLConnector) 2036 and connectable.db_version < '3.33.0' 2037 ): 2038 flavor = 'sqlite_delete_insert' 2039 flavor_key = (f'{flavor}-upsert' if upsert else flavor) 2040 base_queries = UPDATE_QUERIES.get( 2041 flavor_key, 2042 UPDATE_QUERIES['default'] 2043 ) 2044 if not isinstance(base_queries, list): 2045 base_queries = [base_queries] 2046 schema = schema or (connectable.schema if isinstance(connectable, SQLConnector) else None) 2047 patch_schema = patch_schema or schema 2048 target_table_columns = get_table_cols_types( 2049 target, 2050 connectable, 2051 flavor=flavor, 2052 schema=schema, 2053 debug=debug, 2054 ) if not target_cols_types else target_cols_types 2055 patch_table_columns = get_table_cols_types( 2056 patch, 2057 connectable, 2058 flavor=flavor, 2059 schema=patch_schema, 2060 debug=debug, 2061 ) if not patch_cols_types else patch_cols_types 2062 2063 patch_cols_str = ', '.join( 2064 [ 2065 sql_item_name(col, flavor) 2066 for col in patch_table_columns 2067 ] 2068 ) 2069 patch_cols_prefixed_str = ', '.join( 2070 [ 2071 'p.' + sql_item_name(col, flavor) 2072 for col in patch_table_columns 2073 ] 2074 ) 2075 2076 join_cols_str = ', '.join( 2077 [ 2078 sql_item_name(col, flavor) 2079 for col in join_cols 2080 ] 2081 ) 2082 2083 value_cols = [] 2084 join_cols_types = [] 2085 if debug: 2086 dprint("target_table_columns:") 2087 mrsm.pprint(target_table_columns) 2088 for c_name, c_type in target_table_columns.items(): 2089 if c_name not in patch_table_columns: 2090 continue 2091 if flavor in DB_FLAVORS_CAST_DTYPES: 2092 c_type = DB_FLAVORS_CAST_DTYPES[flavor].get(c_type.upper(), c_type) 2093 ( 2094 join_cols_types 2095 if c_name in join_cols 2096 else value_cols 2097 ).append((c_name, c_type)) 2098 if debug: 2099 dprint(f"value_cols: {value_cols}") 2100 2101 if not join_cols_types: 2102 return [] 2103 if not value_cols and not upsert: 2104 return [] 2105 2106 coalesce_join_cols_str = ', '.join( 2107 [ 2108 ( 2109 ( 2110 'COALESCE(' 2111 + sql_item_name(c_name, flavor) 2112 + ', ' 2113 + get_null_replacement(c_type, flavor) 2114 + ')' 2115 ) 2116 if null_indices 2117 else sql_item_name(c_name, flavor) 2118 ) 2119 for c_name, c_type in join_cols_types 2120 ] 2121 ) 2122 2123 update_or_nothing = ('UPDATE' if value_cols else 'NOTHING') 2124 2125 def sets_subquery(l_prefix: str, r_prefix: str): 2126 if not value_cols: 2127 return '' 2128 2129 utc_value_cols = { 2130 c_name 2131 for c_name, c_type in value_cols 2132 if ('utc' in get_pd_type_from_db_type(c_type).lower()) 2133 } if flavor not in TIMEZONE_NAIVE_FLAVORS else set() 2134 2135 cast_func_cols = { 2136 c_name: ( 2137 ('', '', '') 2138 if not cast_columns or ( 2139 flavor == 'oracle' 2140 and are_dtypes_equal(get_pd_type_from_db_type(c_type), 'bytes') 2141 ) 2142 else ( 2143 ('CAST(', f" AS {c_type.replace('_', ' ')}", ')' + ( 2144 " AT TIME ZONE 'UTC'" 2145 if c_name in utc_value_cols 2146 else '' 2147 )) 2148 if flavor not in ('sqlite', 'geopackage') 2149 else ('', '', '') 2150 ) 2151 ) 2152 for c_name, c_type in value_cols 2153 } 2154 return 'SET ' + ',\n'.join([ 2155 ( 2156 l_prefix + sql_item_name(c_name, flavor, None) 2157 + ' = ' 2158 + cast_func_cols[c_name][0] 2159 + r_prefix + sql_item_name(c_name, flavor, None) 2160 + cast_func_cols[c_name][1] 2161 + cast_func_cols[c_name][2] 2162 ) for c_name, c_type in value_cols 2163 ]) 2164 2165 def and_subquery(l_prefix: str, r_prefix: str): 2166 return '\n AND\n '.join([ 2167 ( 2168 ( 2169 "COALESCE(" 2170 + l_prefix 2171 + sql_item_name(c_name, flavor, None) 2172 + ", " 2173 + get_null_replacement(c_type, flavor) 2174 + ")" 2175 + '\n =\n ' 2176 + "COALESCE(" 2177 + r_prefix 2178 + sql_item_name(c_name, flavor, None) 2179 + ", " 2180 + get_null_replacement(c_type, flavor) 2181 + ")" 2182 ) 2183 if null_indices 2184 else ( 2185 l_prefix 2186 + sql_item_name(c_name, flavor, None) 2187 + ' = ' 2188 + r_prefix 2189 + sql_item_name(c_name, flavor, None) 2190 ) 2191 ) for c_name, c_type in join_cols_types 2192 ]) 2193 2194 skip_query_val = "" 2195 target_table_name = sql_item_name(target, flavor, schema) 2196 patch_table_name = sql_item_name(patch, flavor, patch_schema) 2197 dt_col_name = sql_item_name(datetime_col, flavor, None) if datetime_col else None 2198 date_bounds_table = patch_table_name if flavor != 'mssql' else '[date_bounds]' 2199 min_dt_col_name = f"MIN({dt_col_name})" if flavor != 'mssql' else '[Min_dt]' 2200 max_dt_col_name = f"MAX({dt_col_name})" if flavor != 'mssql' else '[Max_dt]' 2201 date_bounds_subquery = ( 2202 f"""f.{dt_col_name} >= (SELECT {min_dt_col_name} FROM {date_bounds_table}) 2203 AND 2204 f.{dt_col_name} <= (SELECT {max_dt_col_name} FROM {date_bounds_table})""" 2205 if datetime_col 2206 else "1 = 1" 2207 ) 2208 with_temp_date_bounds = f"""WITH [date_bounds] AS ( 2209 SELECT MIN({dt_col_name}) AS {min_dt_col_name}, MAX({dt_col_name}) AS {max_dt_col_name} 2210 FROM {patch_table_name} 2211 )""" if datetime_col else "" 2212 identity_insert_on = ( 2213 f"SET IDENTITY_INSERT {target_table_name} ON" 2214 if identity_insert 2215 else skip_query_val 2216 ) 2217 identity_insert_off = ( 2218 f"SET IDENTITY_INSERT {target_table_name} OFF" 2219 if identity_insert 2220 else skip_query_val 2221 ) 2222 2223 ### NOTE: MSSQL upserts must exclude the update portion if only upserting indices. 2224 when_matched_update_sets_subquery_none = "" if not value_cols else ( 2225 "\n WHEN MATCHED THEN\n" 2226 f" UPDATE {sets_subquery('', 'p.')}" 2227 ) 2228 2229 cols_equal_values = '\n,'.join( 2230 [ 2231 f"{sql_item_name(c_name, flavor)} = VALUES({sql_item_name(c_name, flavor)})" 2232 for c_name, c_type in value_cols 2233 ] 2234 ) 2235 on_duplicate_key_update = ( 2236 "ON DUPLICATE KEY UPDATE" 2237 if value_cols 2238 else "" 2239 ) 2240 ignore = "IGNORE " if not value_cols else "" 2241 2242 formatted_queries = [ 2243 textwrap.dedent(base_query.format( 2244 sets_subquery_none=sets_subquery('', 'p.'), 2245 sets_subquery_none_excluded=sets_subquery('', 'EXCLUDED.'), 2246 sets_subquery_f=sets_subquery('f.', 'p.'), 2247 and_subquery_f=and_subquery('p.', 'f.'), 2248 and_subquery_t=and_subquery('p.', 't.'), 2249 target_table_name=target_table_name, 2250 patch_table_name=patch_table_name, 2251 patch_cols_str=patch_cols_str, 2252 patch_cols_prefixed_str=patch_cols_prefixed_str, 2253 date_bounds_subquery=date_bounds_subquery, 2254 join_cols_str=join_cols_str, 2255 coalesce_join_cols_str=coalesce_join_cols_str, 2256 update_or_nothing=update_or_nothing, 2257 when_matched_update_sets_subquery_none=when_matched_update_sets_subquery_none, 2258 cols_equal_values=cols_equal_values, 2259 on_duplicate_key_update=on_duplicate_key_update, 2260 ignore=ignore, 2261 with_temp_date_bounds=with_temp_date_bounds, 2262 identity_insert_on=identity_insert_on, 2263 identity_insert_off=identity_insert_off, 2264 )).lstrip().rstrip() 2265 for base_query in base_queries 2266 ] 2267 2268 ### NOTE: Allow for skipping some queries. 2269 return [query for query in formatted_queries if query]
Build a list of MERGE, UPDATE, DELETE/INSERT queries to apply a patch to target table.
Parameters
- target (str): The name of the target table.
- patch (str): The name of the patch table. This should have the same shape as the target.
- connectable (Union[meerschaum.connectors.sql.SQLConnector, sqlalchemy.orm.session.Session]):
The
SQLConnectoror SQLAlchemy session which will later execute the queries. - join_cols (List[str]): The columns to use to join the patch to the target.
- flavor (Optional[str], default None): If using a SQLAlchemy session, provide the expected database flavor.
- upsert (bool, default False):
If
True, return an upsert query rather than an update. - datetime_col (Optional[str], default None): If provided, bound the join query using this column as the datetime index. This must be present on both tables.
- schema (Optional[str], default None):
If provided, use this schema when quoting the target table.
Defaults to
connector.schema. - patch_schema (Optional[str], default None):
If provided, use this schema when quoting the patch table.
Defaults to
schema. - target_cols_types (Optional[Dict[str, Any]], default None): If provided, use these as the columns-types dictionary for the target table. Default will infer from the database context.
- patch_cols_types (Optional[Dict[str, Any]], default None): If provided, use these as the columns-types dictionary for the target table. Default will infer from the database context.
- identity_insert (bool, default False):
If
True, includeSET IDENTITY_INSERTqueries before and after the update queries. Only applies for MSSQL upserts. - null_indices (bool, default True):
If
False, do not coalesce index columns before joining. - cast_columns (bool, default True):
If
False, do not cast update columns to the target table types. - debug (bool, default False): Verbosity toggle.
Returns
- A list of query strings to perform the update operation.
def
get_null_replacement(typ: str, flavor: str) -> str:
2273def get_null_replacement(typ: str, flavor: str) -> str: 2274 """ 2275 Return a value that may temporarily be used in place of NULL for this type. 2276 2277 Parameters 2278 ---------- 2279 typ: str 2280 The typ to be converted to NULL. 2281 2282 flavor: str 2283 The database flavor for which this value will be used. 2284 2285 Returns 2286 ------- 2287 A value which may stand in place of NULL for this type. 2288 `'None'` is returned if a value cannot be determined. 2289 """ 2290 from meerschaum.utils.dtypes.sql import DB_FLAVORS_CAST_DTYPES 2291 if 'geometry' in typ.lower(): 2292 return "'010100000000008058346FCDC100008058346FCDC1'" 2293 if 'int' in typ.lower() or typ.lower() in ('numeric', 'number'): 2294 return '-987654321' 2295 if 'bool' in typ.lower() or typ.lower() == 'bit': 2296 bool_typ = ( 2297 PD_TO_DB_DTYPES_FLAVORS 2298 .get('bool', {}) 2299 .get(flavor, PD_TO_DB_DTYPES_FLAVORS['bool']['default']) 2300 ) 2301 if flavor in DB_FLAVORS_CAST_DTYPES: 2302 bool_typ = DB_FLAVORS_CAST_DTYPES[flavor].get(bool_typ, bool_typ) 2303 val_to_cast = ( 2304 -987654321 2305 if flavor in ('mysql', 'mariadb') 2306 else 0 2307 ) 2308 return f'CAST({val_to_cast} AS {bool_typ})' 2309 if 'time' in typ.lower() or 'date' in typ.lower(): 2310 db_type = typ if typ.isupper() else None 2311 return dateadd_str(flavor=flavor, begin='1900-01-01', db_type=db_type) 2312 if 'float' in typ.lower() or 'double' in typ.lower() or typ.lower() in ('decimal',): 2313 return '-987654321.0' 2314 if flavor == 'oracle' and typ.lower().split('(', maxsplit=1)[0] == 'char': 2315 return "'-987654321'" 2316 if flavor == 'oracle' and typ.lower() in ('blob', 'bytes'): 2317 return '00' 2318 if typ.lower() in ('uniqueidentifier', 'guid', 'uuid'): 2319 magic_val = 'DEADBEEF-ABBA-BABE-CAFE-DECAFC0FFEE5' 2320 if flavor == 'mssql': 2321 return f"CAST('{magic_val}' AS UNIQUEIDENTIFIER)" 2322 return f"'{magic_val}'" 2323 return ('n' if flavor == 'oracle' else '') + "'-987654321'"
Return a value that may temporarily be used in place of NULL for this type.
Parameters
- typ (str): The typ to be converted to NULL.
- flavor (str): The database flavor for which this value will be used.
Returns
- A value which may stand in place of NULL for this type.
'None'is returned if a value cannot be determined.
def
get_rename_table_queries( old_table: str, new_table: str, flavor: str, schema: str | None = None, new_schema: str | None = None) -> List[str]:
2327def get_rename_table_queries( 2328 old_table: str, 2329 new_table: str, 2330 flavor: str, 2331 schema: Optional[str] = None, 2332 new_schema: Optional[str] = None, 2333) -> List[str]: 2334 """ 2335 Return queries to alter a table's name. 2336 2337 Parameters 2338 ---------- 2339 old_table: str 2340 The unquoted name of the old table. 2341 2342 new_table: str 2343 The unquoted name of the new table. 2344 2345 flavor: str 2346 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`. 2347 2348 schema: Optional[str], default None 2349 The schema on which the table resides. 2350 2351 new_schema: Optional[str], default None 2352 If provided, move the table to this schema. 2353 Defaults to `schema`. 2354 2355 Returns 2356 ------- 2357 A list of `ALTER TABLE` or equivalent queries for the database flavor. 2358 """ 2359 if new_schema is None: 2360 new_schema = schema 2361 2362 if flavor in ('sqlite', 'geopackage'): 2363 schema = None 2364 new_schema = None 2365 2366 old_table_name = sql_item_name(old_table, flavor, schema) 2367 new_table_name = sql_item_name(new_table, flavor, None) 2368 2369 if flavor == 'mssql': 2370 queries = [] 2371 if schema != new_schema: 2372 queries.append(f"ALTER SCHEMA {sql_item_name(new_schema, flavor)} TRANSFER {old_table_name}") 2373 schema = new_schema 2374 old_table_name = sql_item_name(old_table, flavor, schema) 2375 if old_table != new_table: 2376 queries.append(f"EXEC sp_rename '{old_table_name}', '{new_table}'") 2377 return queries 2378 2379 if flavor in ('mysql', 'mariadb'): 2380 new_table_qualified = sql_item_name(new_table, flavor, new_schema) 2381 return [f"RENAME TABLE {old_table_name} TO {new_table_qualified}"] 2382 2383 if_exists_str = "IF EXISTS" if flavor in DROP_IF_EXISTS_FLAVORS else "" 2384 if flavor == 'duckdb': 2385 tmp_table = '_tmp_rename_' + new_table 2386 tmp_table_name = sql_item_name(tmp_table, flavor, schema) 2387 return ( 2388 get_create_table_queries( 2389 f"SELECT * FROM {old_table_name}", 2390 tmp_table, 2391 'duckdb', 2392 schema, 2393 ) + get_create_table_queries( 2394 f"SELECT * FROM {tmp_table_name}", 2395 new_table, 2396 'duckdb', 2397 new_schema, 2398 ) + [ 2399 f"DROP TABLE {if_exists_str} {tmp_table_name}", 2400 f"DROP TABLE {if_exists_str} {old_table_name}", 2401 ] 2402 ) 2403 2404 queries = [] 2405 if schema != new_schema: 2406 if flavor in ('postgresql', 'postgis', 'timescaledb', 'timescaledb-ha', 'citus', 'cockroachdb'): 2407 queries.append(f"ALTER TABLE {old_table_name} SET SCHEMA {sql_item_name(new_schema, flavor)}") 2408 schema = new_schema 2409 old_table_name = sql_item_name(old_table, flavor, schema) 2410 2411 if old_table != new_table: 2412 queries.append(f"ALTER TABLE {old_table_name} RENAME TO {new_table_name}") 2413 2414 return queries
Return queries to alter a table's name.
Parameters
- old_table (str): The unquoted name of the old table.
- new_table (str): The unquoted name of the new table.
- flavor (str):
The database flavor to use for the query (e.g.
'mssql','postgresql'. - schema (Optional[str], default None): The schema on which the table resides.
- new_schema (Optional[str], default None):
If provided, move the table to this schema.
Defaults to
schema.
Returns
- A list of
ALTER TABLEor equivalent queries for the database flavor.
def
get_create_table_query( query_or_dtypes: str | Dict[str, str], new_table: str, flavor: str, schema: str | None = None) -> str:
2418def get_create_table_query( 2419 query_or_dtypes: Union[str, Dict[str, str]], 2420 new_table: str, 2421 flavor: str, 2422 schema: Optional[str] = None, 2423) -> str: 2424 """ 2425 NOTE: This function is deprecated. Use `get_create_table_queries()` instead. 2426 2427 Return a query to create a new table from a `SELECT` query. 2428 2429 Parameters 2430 ---------- 2431 query: Union[str, Dict[str, str]] 2432 The select query to use for the creation of the table. 2433 If a dictionary is provided, return a `CREATE TABLE` query from the given `dtypes` columns. 2434 2435 new_table: str 2436 The unquoted name of the new table. 2437 2438 flavor: str 2439 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`). 2440 2441 schema: Optional[str], default None 2442 The schema on which the table will reside. 2443 2444 Returns 2445 ------- 2446 A `CREATE TABLE` (or `SELECT INTO`) query for the database flavor. 2447 """ 2448 return get_create_table_queries( 2449 query_or_dtypes, 2450 new_table, 2451 flavor, 2452 schema=schema, 2453 primary_key=None, 2454 )[0]
NOTE: This function is deprecated. Use get_create_table_queries() instead.
Return a query to create a new table from a SELECT query.
Parameters
- query (Union[str, Dict[str, str]]):
The select query to use for the creation of the table.
If a dictionary is provided, return a
CREATE TABLEquery from the givendtypescolumns. - new_table (str): The unquoted name of the new table.
- flavor (str):
The database flavor to use for the query (e.g.
'mssql','postgresql'). - schema (Optional[str], default None): The schema on which the table will reside.
Returns
- A
CREATE TABLE(orSELECT INTO) query for the database flavor.
def
get_create_table_queries( query_or_dtypes: str | Dict[str, str], new_table: str, flavor: str, schema: str | None = None, primary_key: str | None = None, primary_key_db_type: str | None = None, autoincrement: bool = False, datetime_column: str | None = None, hypertable_chunk_interval: str | None = None, hypertable_segmentby: List[str] | None = None, hypertable_orderby: List[str] | None = None, partition_by_column: str | None = None, partition_bounds: List | None = None, partition_scheme_name: str | None = None, _parse_dtypes: bool = True) -> List[str]:
2458def get_create_table_queries( 2459 query_or_dtypes: Union[str, Dict[str, str]], 2460 new_table: str, 2461 flavor: str, 2462 schema: Optional[str] = None, 2463 primary_key: Optional[str] = None, 2464 primary_key_db_type: Optional[str] = None, 2465 autoincrement: bool = False, 2466 datetime_column: Optional[str] = None, 2467 hypertable_chunk_interval: Optional[str] = None, 2468 hypertable_segmentby: Optional[List[str]] = None, 2469 hypertable_orderby: Optional[List[str]] = None, 2470 partition_by_column: Optional[str] = None, 2471 partition_bounds: Optional[List] = None, 2472 partition_scheme_name: Optional[str] = None, 2473 _parse_dtypes: bool = True, 2474) -> List[str]: 2475 """ 2476 Return a query to create a new table from a `SELECT` query or a `dtypes` dictionary. 2477 2478 Parameters 2479 ---------- 2480 query_or_dtypes: Union[str, Dict[str, str]] 2481 The select query to use for the creation of the table. 2482 If a dictionary is provided, return a `CREATE TABLE` query from the given `dtypes` columns. 2483 2484 new_table: str 2485 The unquoted name of the new table. 2486 2487 flavor: str 2488 The database flavor to use for the query (e.g. `'mssql'`, `'postgresql'`). 2489 2490 schema: Optional[str], default None 2491 The schema on which the table will reside. 2492 2493 primary_key: Optional[str], default None 2494 If provided, designate this column as the primary key in the new table. 2495 2496 primary_key_db_type: Optional[str], default None 2497 If provided, alter the primary key to this type (to set NOT NULL constraint). 2498 2499 autoincrement: bool, default False 2500 If `True` and `primary_key` is provided, create the `primary_key` column 2501 as an auto-incrementing integer column. 2502 2503 datetime_column: Optional[str], default None 2504 If provided, include this column in the primary key. 2505 Applicable to TimescaleDB only. 2506 2507 hypertable_chunk_interval: Optional[str], default None 2508 If provided (and `flavor` is TimescaleDB), create the table as a hypertable using the 2509 declarative `CREATE TABLE ... WITH (tsdb.hypertable, ...)` syntax, partitioned on 2510 `datetime_column` with this chunk interval (e.g. `'10080 minutes'` or `'100000'`). 2511 Requires TimescaleDB 2.21+; callers should fall back to `create_hypertable()` on failure. 2512 2513 hypertable_segmentby: Optional[List[str]], default None 2514 If provided (alongside `hypertable_chunk_interval`), enable the Hypercore columnstore on 2515 the new hypertable and segment compressed chunks by these columns (`tsdb.segmentby`). 2516 Should be low-cardinality columns; high-cardinality / unique columns belong in 2517 `hypertable_orderby` instead. 2518 2519 hypertable_orderby: Optional[List[str]], default None 2520 If provided (alongside `hypertable_chunk_interval`), order the columnstore by these 2521 columns (`tsdb.orderby`, e.g. `['"timestamp" DESC']`). Declaring `segmentby`/`orderby` 2522 in `CREATE TABLE` causes TimescaleDB to auto-create a columnstore policy. 2523 2524 _parse_dtypes: bool, default True 2525 If `True`, cast Pandas dtypes to SQL dtypes. 2526 Otherwise pass through the given value directly. 2527 2528 Returns 2529 ------- 2530 A `CREATE TABLE` (or `SELECT INTO`) query for the database flavor. 2531 """ 2532 if not isinstance(query_or_dtypes, (str, dict)): 2533 raise TypeError("`query_or_dtypes` must be a query or a dtypes dictionary.") 2534 2535 method = ( 2536 _get_create_table_query_from_cte 2537 if isinstance(query_or_dtypes, str) 2538 else _get_create_table_query_from_dtypes 2539 ) 2540 return method( 2541 query_or_dtypes, 2542 new_table, 2543 flavor, 2544 schema=schema, 2545 primary_key=primary_key, 2546 primary_key_db_type=primary_key_db_type, 2547 autoincrement=(autoincrement and flavor not in SKIP_AUTO_INCREMENT_FLAVORS), 2548 datetime_column=datetime_column, 2549 hypertable_chunk_interval=hypertable_chunk_interval, 2550 hypertable_segmentby=hypertable_segmentby, 2551 hypertable_orderby=hypertable_orderby, 2552 partition_by_column=partition_by_column, 2553 partition_bounds=partition_bounds, 2554 partition_scheme_name=partition_scheme_name, 2555 _parse_dtypes=_parse_dtypes, 2556 )
Return a query to create a new table from a SELECT query or a dtypes dictionary.
Parameters
- query_or_dtypes (Union[str, Dict[str, str]]):
The select query to use for the creation of the table.
If a dictionary is provided, return a
CREATE TABLEquery from the givendtypescolumns. - new_table (str): The unquoted name of the new table.
- flavor (str):
The database flavor to use for the query (e.g.
'mssql','postgresql'). - schema (Optional[str], default None): The schema on which the table will reside.
- primary_key (Optional[str], default None): If provided, designate this column as the primary key in the new table.
- primary_key_db_type (Optional[str], default None): If provided, alter the primary key to this type (to set NOT NULL constraint).
- autoincrement (bool, default False):
If
Trueandprimary_keyis provided, create theprimary_keycolumn as an auto-incrementing integer column. - datetime_column (Optional[str], default None): If provided, include this column in the primary key. Applicable to TimescaleDB only.
- hypertable_chunk_interval (Optional[str], default None):
If provided (and
flavoris TimescaleDB), create the table as a hypertable using the declarativeCREATE TABLE ... WITH (tsdb.hypertable, ...)syntax, partitioned ondatetime_columnwith this chunk interval (e.g.'10080 minutes'or'100000'). Requires TimescaleDB 2.21+; callers should fall back tocreate_hypertable()on failure. - hypertable_segmentby (Optional[List[str]], default None):
If provided (alongside
hypertable_chunk_interval), enable the Hypercore columnstore on the new hypertable and segment compressed chunks by these columns (tsdb.segmentby). Should be low-cardinality columns; high-cardinality / unique columns belong inhypertable_orderbyinstead. - hypertable_orderby (Optional[List[str]], default None):
If provided (alongside
hypertable_chunk_interval), order the columnstore by these columns (tsdb.orderby, e.g.['"timestamp" DESC']). Declaringsegmentby/orderbyinCREATE TABLEcauses TimescaleDB to auto-create a columnstore policy. - _parse_dtypes (bool, default True):
If
True, cast Pandas dtypes to SQL dtypes. Otherwise pass through the given value directly.
Returns
- A
CREATE TABLE(orSELECT INTO) query for the database flavor.
def
get_reset_autoincrement_queries( table: str, column: str, connector: meerschaum.connectors.SQLConnector, schema: str | None = None, debug: bool = False) -> List[str]:
2952def get_reset_autoincrement_queries( 2953 table: str, 2954 column: str, 2955 connector: mrsm.connectors.SQLConnector, 2956 schema: Optional[str] = None, 2957 debug: bool = False, 2958) -> List[str]: 2959 """ 2960 Return a list of queries to reset a table's auto-increment counter to the next largest value. 2961 2962 Parameters 2963 ---------- 2964 table: str 2965 The name of the table on which the auto-incrementing column exists. 2966 2967 column: str 2968 The name of the auto-incrementing column. 2969 2970 connector: mrsm.connectors.SQLConnector 2971 The SQLConnector to the database on which the table exists. 2972 2973 schema: Optional[str], default None 2974 The schema of the table. Defaults to `connector.schema`. 2975 2976 Returns 2977 ------- 2978 A list of queries to be executed to reset the auto-incrementing column. 2979 """ 2980 if not table_exists(table, connector, schema=schema, debug=debug): 2981 return [] 2982 2983 schema = schema or connector.schema 2984 max_id_name = sql_item_name('max_id', connector.flavor) 2985 table_name = sql_item_name(table, connector.flavor, schema) 2986 table_seq_name = sql_item_name(table + '_' + column + '_seq', connector.flavor, schema) 2987 column_name = sql_item_name(column, connector.flavor) 2988 max_id = connector.value( 2989 f""" 2990 SELECT COALESCE(MAX({column_name}), 0) AS {max_id_name} 2991 FROM {table_name} 2992 """, 2993 debug=debug, 2994 silent=True, 2995 ) 2996 if max_id is None: 2997 return [] 2998 2999 reset_queries = reset_autoincrement_queries.get( 3000 connector.flavor, 3001 reset_autoincrement_queries['default'] 3002 ) 3003 if not isinstance(reset_queries, list): 3004 reset_queries = [reset_queries] 3005 3006 return [ 3007 query.format( 3008 column=column, 3009 column_name=column_name, 3010 table=table, 3011 table_name=table_name, 3012 table_seq_name=table_seq_name, 3013 val=max_id, 3014 val_plus_1=(max_id + 1), 3015 ) 3016 for query in reset_queries 3017 ]
Return a list of queries to reset a table's auto-increment counter to the next largest value.
Parameters
- table (str): The name of the table on which the auto-incrementing column exists.
- column (str): The name of the auto-incrementing column.
- connector (mrsm.connectors.SQLConnector): The SQLConnector to the database on which the table exists.
- schema (Optional[str], default None):
The schema of the table. Defaults to
connector.schema.
Returns
- A list of queries to be executed to reset the auto-incrementing column.
def
get_create_schema_if_not_exists_queries(schema: str, flavor: str) -> List[str]:
3021def get_create_schema_if_not_exists_queries( 3022 schema: str, 3023 flavor: str, 3024) -> List[str]: 3025 """ 3026 Return the queries to create a schema if it does not yet exist. 3027 For databases which do not support schemas, an empty list will be returned. 3028 """ 3029 if not schema: 3030 return [] 3031 3032 if flavor in NO_SCHEMA_FLAVORS: 3033 return [] 3034 3035 if schema == DEFAULT_SCHEMA_FLAVORS.get(flavor, None): 3036 return [] 3037 3038 clean(schema) 3039 3040 if flavor == 'mssql': 3041 return [ 3042 ( 3043 f"IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{schema}')\n" 3044 "BEGIN\n" 3045 f" EXEC('CREATE SCHEMA {sql_item_name(schema, flavor)}');\n" 3046 "END;" 3047 ) 3048 ] 3049 3050 if flavor == 'oracle': 3051 return [] 3052 3053 return [ 3054 f"CREATE SCHEMA IF NOT EXISTS {sql_item_name(schema, flavor)};" 3055 ]
Return the queries to create a schema if it does not yet exist. For databases which do not support schemas, an empty list will be returned.
NATIVE_PARTITION_FLAVORS =
{'postgis', 'mariadb', 'postgresql', 'mysql'}
exists_queries =
{'default': 'SELECT COUNT(*) FROM {table_name} WHERE 1 = 0'}
DROP_INDEX_IF_EXISTS_FLAVORS =
{'timescaledb', 'postgis', 'mssql', 'sqlite', 'timescaledb-ha', 'geopackage', 'postgresql', 'citus'}
COALESCE_UNIQUE_INDEX_FLAVORS =
{'timescaledb', 'postgis', 'postgresql', 'citus', 'timescaledb-ha'}
hypertable_queries =
{'timescaledb': "SELECT hypertable_size('{table_name}')", 'timescaledb-ha': "SELECT hypertable_size('{table_name}')", 'citus': "SELECT citus_table_size('{table_name}')"}
table_wrappers =
{'default': ('"', '"'), 'timescaledb': ('"', '"'), 'timescaledb-ha': ('"', '"'), 'citus': ('"', '"'), 'duckdb': ('"', '"'), 'postgresql': ('"', '"'), 'postgis': ('"', '"'), 'sqlite': ('"', '"'), 'geopackage': ('"', '"'), 'mysql': ('`', '`'), 'mariadb': ('`', '`'), 'mssql': ('[', ']'), 'cockroachdb': ('"', '"'), 'oracle': ('"', '"')}
OMIT_NULLSFIRST_FLAVORS =
{'mssql', 'mariadb', 'mysql'}
SINGLE_ALTER_TABLE_FLAVORS =
{'mssql', 'sqlite', 'oracle', 'duckdb', 'geopackage'}
NO_SELECT_INTO_FLAVORS =
{'sqlite', 'mysql', 'oracle', 'duckdb', 'mariadb', 'geopackage'}