-
-
Notifications
You must be signed in to change notification settings - Fork 6
/
sql_utils.py
553 lines (423 loc) · 16.1 KB
/
sql_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
import ast
import os
import re
import sqlalchemy as sa
from sqlalchemy.dialects.mysql import dialect as mysql_dialect
from sqlalchemy.dialects.postgresql import dialect as postgresql_dialect
from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
from sqlalchemy.engine import reflection
from sqlalchemy.engine.url import make_url
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import expression as exp
import sqlparse as sp
from zillion.core import *
# from zillion.nlp import build_chain, PromptTemplate
DIGIT_THRESHOLD_FOR_MEAN_AGGR = 1
INTEGER_SA_TYPES = [
sa.BigInteger,
sa.BIGINT,
sa.Integer,
sa.INT,
sa.INTEGER,
sa.SmallInteger,
sa.SMALLINT,
]
FLOAT_SA_TYPES = [
sa.DECIMAL,
sa.Float,
sa.FLOAT,
sa.Numeric,
sa.NUMERIC,
sa.REAL,
sa.dialects.postgresql.DOUBLE_PRECISION,
sa.dialects.postgresql.MONEY,
]
NUMERIC_SA_TYPES = INTEGER_SA_TYPES + FLOAT_SA_TYPES
DATETIME_SA_TYPES = [sa.DateTime, sa.DATETIME, sa.Time, sa.TIME, sa.TIMESTAMP]
DATE_SA_TYPES = [sa.Date, sa.DATE]
AGGREGATION_SQLA_FUNC_MAP = {
AggregationTypes.MEAN: sa.func.AVG,
AggregationTypes.COUNT: sa.func.COUNT,
AggregationTypes.COUNT_DISTINCT: lambda x: sa.func.COUNT(sa.distinct(x)),
AggregationTypes.MIN: sa.func.MIN,
AggregationTypes.MAX: sa.func.MAX,
AggregationTypes.SUM: sa.func.SUM,
}
SQL_AGGREGATION_FUNCS = set(
[
"AVG",
"SUM",
"MIN",
"MAX",
"COUNT",
"COUNT_DISTINCT",
"STD",
"MEDIAN",
"MODE",
"VAR",
]
)
# This establishes a baseline of schemas to ignore during reflection
DIALECT_IGNORE_SCHEMAS = {
"mysql": set(["information_schema", "performance_schema", "mysql", "sys"]),
"postgresql": set(["information_schema", r"pg_(.*)"]),
}
class InvalidSQLAlchemyTypeString(Exception):
pass
def contains_sql_keywords(sql):
"""Determine whether a SQL query contains special SQL keywords (DML, DDL,
etc.)
**Parameters:**
* **sql** - (*str or sqlparse result*) The SQL query to check for keywords
**Returns:**
(*bool*) - True if the SQL string contains keywords
"""
if isinstance(sql, str):
sql = sp.parse(sql)
for token in sql:
if token.ttype in (sp.tokens.DML, sp.tokens.DDL, sp.tokens.CTE):
return True
if isinstance(token, sp.sql.TokenList):
token_result = contains_sql_keywords(token)
if token_result:
return True
return False
def contains_aggregation(sql):
"""Determine whether a SQL query contains aggregation functions.
**Warning:**
This relies on a non-exhaustive list of SQL aggregation functions
to look for. This will likely need updating.
**Parameters:**
* **sql** - (*str or sqlparse result*) The SQL query to check for
aggregation functions
**Returns:**
(*bool*) - True if the SQL string contains aggregation
"""
if isinstance(sql, str):
sql = sp.parse(sql)
for token in sql:
if isinstance(token, sp.sql.Function):
name = token.get_name()
if name.upper() in SQL_AGGREGATION_FUNCS:
return True
if isinstance(token, sp.sql.TokenList):
token_result = contains_aggregation(token)
if token_result:
return True
return False
def type_string_to_sa_type(type_string):
"""Convert a field type string to a SQLAlchemy type. The type string will be
evaluated as a python statement or class name to init from the SQLAlchemy
top level module. Dialect-specific SQLAlchemy types are not currently
supported.
**Parameters:**
* **type_string** - (*str*) A string representing a SQLAlchemy type, such as
"Integer", or "String(32)". This does a case-insensitive search and will
return the first matching SQLAlchemy type.
**Returns:**
(*SQLAlchemy type object*) - An init'd SQLAlchemy type object
"""
try:
tree = ast.parse(type_string)
ast_obj = tree.body[0].value
if isinstance(ast_obj, ast.Name):
type_name = ast_obj.id
type_args = []
type_kwargs = {}
else:
type_name = ast_obj.func.id
type_args = [arg.n for arg in ast_obj.args]
type_kwargs = {k.arg: k.value.n for k in ast_obj.keywords}
type_cls = igetattr(sa.types, type_name, None)
if not type_cls:
raise InvalidSQLAlchemyTypeString(
"Could not find matching type for %s" % type_name
)
return type_cls(*type_args, **type_kwargs)
except Exception as e:
raise InvalidSQLAlchemyTypeString("Unable to parse %s" % type_string) from e
def to_generic_sa_type(type):
"""Return a generic SQLAlchemy type object from a type that may be dialect-
specific. This will attempt to preserve common type settings such as
specified field length, scale, and precision. On error it will fall back to
trying to init the generic type with no params.
"""
params = {}
for param in ["length", "precision", "scale"]:
if hasattr(type, param):
params[param] = getattr(type, param)
try:
return type._type_affinity(**params)
except Exception as e:
if "unexpected keyword" not in str(e):
raise
return type._type_affinity()
def infer_aggregation_and_rounding(column):
"""Infer the aggregation and rounding settings based on the column type.
This is just a rough / best guess based on the column type, precision
and rounding settings.
**Parameters:**
* **column** - (*SQLAlchemy column*) The column to analyze
**Returns:**
(*AggregationType, int*) - A 2-item tuple of the aggregation type and
rounding to use
"""
if isinstance(column.type, tuple(INTEGER_SA_TYPES)):
return AggregationTypes.SUM, 0
if isinstance(column.type, tuple(FLOAT_SA_TYPES)):
rounding = column.type.scale
precision = column.type.precision
if rounding is None and precision is None:
aggregation = AggregationTypes.SUM
else:
if rounding:
whole_digits = precision - rounding
if whole_digits <= DIGIT_THRESHOLD_FOR_MEAN_AGGR:
aggregation = AggregationTypes.MEAN
else:
aggregation = AggregationTypes.SUM
else:
# We really don't know what to do here, so we'll just
# guess SUM.
aggregation = AggregationTypes.SUM
return aggregation, rounding
raise ZillionException("Column %s is not a numeric type" % column)
def aggregation_to_sqla_func(aggregation):
"""Convert an AggregationType string to a SQLAlchemy function"""
return AGGREGATION_SQLA_FUNC_MAP[aggregation]
def is_numeric_type(type):
"""Determine if this is a numeric SQLAlchemy type"""
raiseif(isinstance(type, str), "Expected a SQLAlchemy type, got string")
if isinstance(type, tuple(NUMERIC_SA_TYPES)):
return True
return False
def is_probably_metric(column, formula=None, nlp_column_info=None):
"""Determine if a column is probably a metric. This is used when trying to
automatically init/reflect a datasource and determine the field types for
columns. The logic is very coarse, and should not be relied on for more than
quick/convenient use cases.
**Parameters:**
* **column** - (*SQLAlchemy column*) The column to analyze
* **formula** - (*str, optional*) A formula to calculate the column
* **nlp_column_info** - (*dict, optional*) Column attributes inferred from
natural language processing of the table/column definitions
**Returns:**
(*bool*) - True if the column is probably a metric
"""
if formula and contains_aggregation(formula):
return True
if column.primary_key:
return False
if column.name.endswith("_id") or column.name.endswith("Id") or column.name == "id":
return False
if nlp_column_info:
if nlp_column_info.get("type", None) == FieldTypes.METRIC:
return True
return False
if not isinstance(column.type, tuple(NUMERIC_SA_TYPES)):
return False
return True
def sqla_compile(expr):
"""Compile a SQL expression
**Parameters:**
* **expr** - (*SQLAlchemy expression*) The SQLAlchemy expression to compile
**Returns:**
(*str*) - The compiled expression string
"""
return str(expr.compile(compile_kwargs={"literal_binds": True}))
def printexpr(expr):
"""Print a SQLAlchemy expression"""
print(sqla_compile(expr))
def column_fullname(column, prefix=None):
"""Get a fully qualified name for a column
**Parameters:**
* **column** - (*SQLAlchemy column*) A SQLAlchemy column object to get the
full name for
* **prefix** - (*str, optional*) If specified, a manual prefix to prepend to
the output string. This will automatically be separted with a ".".
**Returns:**
(*str*) - A fully qualified column name. The exact format will vary
depending on your SQLAlchemy metadata, but an example would be:
schema.table.column
"""
name = "%s.%s" % (column.table.fullname, column.name)
if prefix:
name = prefix + "." + name
return name
def get_schema_and_table_name(table):
"""Extract the schema and table name from a full table name. If the table
name is not schema-qualified, return None for the schema name"""
schema = None
table_name = table
if "." in table:
parts = table.split(".")
raiseifnot(len(parts) == 2, "Invalid table name: %s" % table)
schema, table_name = parts
return schema, table_name
def get_sqla_criterion_expr(column, criterion, negate=False):
"""Create a SQLAlchemy criterion expression
**Parameters:**
* **column** - (*SQLAlchemy column*) A SQLAlchemy column object to be used
in the expression
* **criterion** - (*3-item iterable*) A 3-item tuple or list of the format
[field, operation, value(s)]. See `core.CRITERIA_OPERATIONS` for supported
operations. The value item may take on different formats depending on the
operation. In most cases passing an iterable will result in multiple
criteria of that operation being formed. For example, ("my_field", "=",
[1,2,3]) would logically or 3 conditions of equality to the 3 values in the
list. The "between" operations expect each value to be a 2-item iterable
representing the lower and upper bound of the criterion.
* **negate** - (*bool, optional*) Negate the expression
**Returns:**
(*SQLAlchemy expression*) - A SQLALchemy expression representing the
criterion
**Notes:**
Postgresql "like" is case sensitive, but mysql "like" is not. Postgresql
also supports "ilike" to specify case insensitive, so one option is to look
at the dialect to determine the function, but that is not supported yet.
"""
field, op, values = criterion
op = op.lower()
if not isinstance(values, (list, tuple)):
values = [values]
use_or = True
has_null = any([v is None for v in values])
if op == "=":
clauses = [column == v if v is not None else column.is_(None) for v in values]
elif op == "!=":
clauses = [column != v if v is not None else column.isnot(None) for v in values]
elif op == ">":
clauses = [column > v for v in values]
elif op == "<":
clauses = [column < v for v in values]
elif op == ">=":
clauses = [column >= v for v in values]
elif op == "<=":
clauses = [column <= v for v in values]
elif op == "in":
if has_null:
clauses = [
column == v if v is not None else column.is_(None) for v in values
]
else:
clauses = [column.in_(values)]
elif op == "not in":
use_or = False
if has_null:
clauses = [
column != v if v is not None else column.isnot(None) for v in values
]
else:
clauses = [sa.not_(column.in_(values))]
elif op == "between":
raiseifnot(len(values) == 2, "Between clause value must have length of 2")
clauses = [column.between(values[0], values[1])]
elif op == "not between":
raiseifnot(len(values) == 2, "Between clause value must have length of 2")
clauses = [sa.not_(column.between(values[0], values[1]))]
elif op == "like":
clauses = [column.like(v) for v in values]
elif op == "not like":
use_or = False
clauses = [sa.not_(column.like(v)) for v in values]
else:
raise ZillionException("Invalid criterion operand: %s" % op)
if use_or:
clause = sa.or_(*clauses)
else:
clause = sa.and_(*clauses)
if negate:
clause = sa.not_(clause)
return clause
def check_metadata_url(url, confirm_exists=False):
"""Check validity of the metadata URL"""
url = make_url(url)
dialect = url.get_dialect().name
if confirm_exists:
if dialect == "sqlite":
raiseifnot(
os.path.isfile(url.database),
"SQLite DB does not exist: %s" % url.database,
)
else:
raise AssertionError(
"confirm_exists not supported for dialect: %s" % dialect
)
def comment(self, c):
"""See https://github.com/sqlalchemy/sqlalchemy/wiki/CompiledComments"""
self._added_comment = c
return self
exp.ClauseElement.comment = comment
exp.ClauseElement._added_comment = None
def _compile_element(elem, prepend_newline=False):
"""See https://github.com/sqlalchemy/sqlalchemy/wiki/CompiledComments"""
@compiles(elem)
def add_comment(element, compiler, **kw):
meth = getattr(compiler, "visit_%s" % element.__visit_name__)
text = meth(element, **kw)
if element._added_comment:
# Modified this line to not add newline
text = "-- %s\n" % element._added_comment + text
elif prepend_newline:
text = "\n" + text
return text
_compile_element(exp.Case)
_compile_element(exp.Label, True)
_compile_element(exp.ColumnClause)
_compile_element(exp.Join)
_compile_element(exp.Select)
_compile_element(exp.Alias)
_compile_element(exp.Exists)
def get_schemas(engine):
"""Inspect the SQLAlchemy engine to get a list of schemas"""
insp = reflection.Inspector.from_engine(engine)
return insp.get_schema_names()
# -------- Some DB-specific stuff
def to_mysql_type(type):
"""Compile into a MySQL SQLAlchemy type"""
return type.compile(dialect=mysql_dialect())
def to_postgresql_type(type):
"""Compile into a PostgreSQL SQLAlchemy type"""
return type.compile(dialect=postgresql_dialect())
def to_sqlite_type(type):
"""Compile into a SQLite SQLAlchemy type"""
return type.compile(dialect=sqlite_dialect())
def to_duckdb_type(type):
"""Compile into a DuckDB SQLAlchemy type"""
from duckdb_engine import Dialect as duckdb_dialect
return type.compile(dialect=duckdb_dialect())
def filter_dialect_schemas(schemas, dialect):
"""Filter out a set of baseline/system schemas for a dialect
**Parameters:**
* **schemas** - (*list*) A list of schema names
* **dialect** - (*str*) The name of a SQLAlchemy dialect
**Returns:**
(*list*) - A filtered list of schema names
"""
ignores = DIALECT_IGNORE_SCHEMAS.get(dialect, None)
if not ignores:
return schemas
final = []
for schema in schemas:
add = True
for ignore in ignores:
if re.match(ignore, schema):
add = False
break
if add:
final.append(schema)
return final
def get_postgres_schemas(conn):
"""Helper to list PostgreSQL schemas"""
qr = conn.execute(
sa.text(
"SELECT schema_name FROM information_schema.schemata "
"WHERE schema_name not LIKE 'pg_%' and schema_name != 'information_schema'"
)
)
return [x["schema_name"] for x in qr.fetchall()]
def get_postgres_pid(conn):
"""Helper to get the PostgreSQL connection PID"""
qr = conn.execute("select pg_backend_pid()")
pid = qr.fetchone()[0]
return pid