-
Notifications
You must be signed in to change notification settings - Fork 14.4k
/
Copy pathhelpers.py
2134 lines (1882 loc) · 79.7 KB
/
helpers.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
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=too-many-lines
"""a collection of model-related helper classes and functions"""
import builtins
import dataclasses
import logging
import re
import uuid
from collections.abc import Hashable
from datetime import datetime, timedelta
from typing import Any, cast, NamedTuple, Optional, TYPE_CHECKING, Union
import dateutil.parser
import humanize
import numpy as np
import pandas as pd
import pytz
import sqlalchemy as sa
import sqlparse
import yaml
from flask import g
from flask_appbuilder import Model
from flask_appbuilder.models.decorators import renders
from flask_appbuilder.models.mixins import AuditMixin
from flask_appbuilder.security.sqla.models import User
from flask_babel import lazy_gettext as _
from jinja2.exceptions import TemplateError
from markupsafe import escape, Markup
from sqlalchemy import and_, Column, or_, UniqueConstraint
from sqlalchemy.exc import MultipleResultsFound
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import Mapper, validates
from sqlalchemy.sql.elements import ColumnElement, literal_column, TextClause
from sqlalchemy.sql.expression import Label, Select, TextAsFrom
from sqlalchemy.sql.selectable import Alias, TableClause
from sqlalchemy_utils import UUIDType
from superset import app, db, is_feature_enabled
from superset.advanced_data_type.types import AdvancedDataTypeResponse
from superset.common.db_query_status import QueryStatus
from superset.common.utils.time_range_utils import get_since_until_from_time_range
from superset.constants import EMPTY_STRING, NULL_STRING
from superset.db_engine_specs.base import TimestampExpression
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
AdvancedDataTypeResponseError,
ColumnNotFoundException,
QueryClauseValidationException,
QueryObjectValidationError,
SupersetParseError,
SupersetSecurityException,
)
from superset.extensions import feature_flag_manager
from superset.jinja_context import BaseTemplateProcessor
from superset.sql.parse import SQLScript
from superset.sql_parse import (
has_table_query,
insert_rls_in_predicate,
sanitize_clause,
)
from superset.superset_typing import (
AdhocMetric,
Column as ColumnTyping,
FilterValue,
FilterValues,
Metric,
OrderBy,
QueryObjectDict,
)
from superset.utils import core as utils, json
from superset.utils.core import (
GenericDataType,
get_column_name,
get_non_base_axis_columns,
get_user_id,
is_adhoc_column,
MediumText,
remove_duplicates,
)
from superset.utils.dates import datetime_to_epoch
if TYPE_CHECKING:
from superset.connectors.sqla.models import SqlMetric, TableColumn
from superset.db_engine_specs import BaseEngineSpec
from superset.models.core import Database
config = app.config
logger = logging.getLogger(__name__)
VIRTUAL_TABLE_ALIAS = "virtual_table"
SERIES_LIMIT_SUBQ_ALIAS = "series_limit"
ADVANCED_DATA_TYPES = config["ADVANCED_DATA_TYPES"]
def validate_adhoc_subquery(
sql: str,
database_id: int,
engine: str,
default_schema: str,
) -> str:
"""
Check if adhoc SQL contains sub-queries or nested sub-queries with table.
If sub-queries are allowed, the adhoc SQL is modified to insert any applicable RLS
predicates to it.
:param sql: adhoc sql expression
:raise SupersetSecurityException if sql contains sub-queries or
nested sub-queries with table
"""
statements = []
for statement in sqlparse.parse(sql):
try:
has_table = has_table_query(str(statement), engine)
except SupersetParseError:
has_table = True
if has_table:
if not is_feature_enabled("ALLOW_ADHOC_SUBQUERY"):
raise SupersetSecurityException(
SupersetError(
error_type=SupersetErrorType.ADHOC_SUBQUERY_NOT_ALLOWED_ERROR,
message=_("Custom SQL fields cannot contain sub-queries."),
level=ErrorLevel.ERROR,
)
)
# TODO (betodealmeida): reimplement with sqlglot
statement = insert_rls_in_predicate(statement, database_id, default_schema)
statements.append(statement)
return ";\n".join(str(statement) for statement in statements)
def json_to_dict(json_str: str) -> dict[Any, Any]:
if json_str:
val = re.sub(",[ \t\r\n]+}", "}", json_str)
val = re.sub(",[ \t\r\n]+\\]", "]", val)
return json.loads(val)
return {}
def convert_uuids(obj: Any) -> Any:
"""
Convert UUID objects to str so we can use yaml.safe_dump
"""
if isinstance(obj, uuid.UUID):
return str(obj)
if isinstance(obj, list):
return [convert_uuids(el) for el in obj]
if isinstance(obj, dict):
return {k: convert_uuids(v) for k, v in obj.items()}
return obj
class UUIDMixin: # pylint: disable=too-few-public-methods
uuid = sa.Column(
UUIDType(binary=True), primary_key=False, unique=True, default=uuid.uuid4
)
@property
def short_uuid(self) -> str:
return str(self.uuid)[:8]
class ImportExportMixin(UUIDMixin):
export_parent: Optional[str] = None
# The name of the attribute
# with the SQL Alchemy back reference
export_children: list[str] = []
# List of (str) names of attributes
# with the SQL Alchemy forward references
export_fields: list[str] = []
# The names of the attributes
# that are available for import and export
extra_import_fields: list[str] = []
# Additional fields that should be imported,
# even though they were not exported
__mapper__: Mapper
@classmethod
def _unique_constraints(cls) -> list[set[str]]:
"""Get all (single column and multi column) unique constraints"""
unique = [
{c.name for c in u.columns}
for u in cls.__table_args__ # type: ignore
if isinstance(u, UniqueConstraint)
]
unique.extend({c.name} for c in cls.__table__.columns if c.unique) # type: ignore
return unique
@classmethod
def parent_foreign_key_mappings(cls) -> dict[str, str]:
"""Get a mapping of foreign name to the local name of foreign keys"""
parent_rel = cls.__mapper__.relationships.get(cls.export_parent)
if parent_rel:
return {
local.name: remote.name
for (local, remote) in parent_rel.local_remote_pairs
}
return {}
@classmethod
def export_schema(
cls, recursive: bool = True, include_parent_ref: bool = False
) -> dict[str, Any]:
"""Export schema as a dictionary"""
parent_excludes = set()
if not include_parent_ref:
parent_ref = cls.__mapper__.relationships.get(cls.export_parent)
if parent_ref:
parent_excludes = {column.name for column in parent_ref.local_columns}
def formatter(column: sa.Column) -> str:
return (
f"{str(column.type)} Default ({column.default.arg})"
if column.default
else str(column.type)
)
schema: dict[str, Any] = {
column.name: formatter(column)
for column in cls.__table__.columns # type: ignore
if (column.name in cls.export_fields and column.name not in parent_excludes)
}
if recursive:
for column in cls.export_children:
child_class = cls.__mapper__.relationships[column].argument.class_
schema[column] = [
child_class.export_schema(
recursive=recursive, include_parent_ref=include_parent_ref
)
]
return schema
@classmethod
def import_from_dict( # noqa: C901
# pylint: disable=too-many-arguments,too-many-branches,too-many-locals
cls,
dict_rep: dict[Any, Any],
parent: Optional[Any] = None,
recursive: bool = True,
sync: Optional[list[str]] = None,
allow_reparenting: bool = False,
) -> Any:
"""Import obj from a dictionary"""
if sync is None:
sync = []
parent_refs = cls.parent_foreign_key_mappings()
export_fields = (
set(cls.export_fields)
| set(cls.extra_import_fields)
| set(parent_refs.keys())
| {"uuid"}
)
new_children = {c: dict_rep[c] for c in cls.export_children if c in dict_rep}
unique_constraints = cls._unique_constraints()
filters = [] # Using these filters to check if obj already exists
# Remove fields that should not get imported
for k in list(dict_rep):
if k not in export_fields and k not in parent_refs:
del dict_rep[k]
if not parent:
if cls.export_parent:
for prnt in parent_refs.keys():
if prnt not in dict_rep:
raise RuntimeError(f"{cls.__name__}: Missing field {prnt}")
else:
# Set foreign keys to parent obj
for k, v in parent_refs.items():
dict_rep[k] = getattr(parent, v)
if not allow_reparenting:
# Add filter for parent obj
filters.extend(
[getattr(cls, k) == dict_rep.get(k) for k in parent_refs.keys()]
)
# Add filter for unique constraints
ucs = [
and_(
*[
getattr(cls, k) == dict_rep.get(k)
for k in cs
if dict_rep.get(k) is not None
]
)
for cs in unique_constraints
]
filters.append(or_(*ucs))
# Check if object already exists in DB, break if more than one is found
try:
obj_query = db.session.query(cls).filter(and_(*filters))
obj = obj_query.one_or_none()
except MultipleResultsFound:
logger.error(
"Error importing %s \n %s \n %s",
cls.__name__,
str(obj_query),
yaml.safe_dump(dict_rep),
exc_info=True,
)
raise
if not obj:
is_new_obj = True
# Create new DB object
obj = cls(**dict_rep)
logger.debug("Importing new %s %s", obj.__tablename__, str(obj))
if cls.export_parent and parent:
setattr(obj, cls.export_parent, parent)
db.session.add(obj)
else:
is_new_obj = False
logger.debug("Updating %s %s", obj.__tablename__, str(obj))
# Update columns
for k, v in dict_rep.items():
setattr(obj, k, v)
# Recursively create children
if recursive:
for child in cls.export_children:
argument = cls.__mapper__.relationships[child].argument
child_class = (
argument.class_ if hasattr(argument, "class_") else argument
)
added = []
for c_obj in new_children.get(child, []):
added.append(
child_class.import_from_dict(
dict_rep=c_obj, parent=obj, sync=sync
)
)
# If children should get synced, delete the ones that did not
# get updated.
if child in sync and not is_new_obj:
back_refs = child_class.parent_foreign_key_mappings()
delete_filters = [
getattr(child_class, k) == getattr(obj, back_refs.get(k))
for k in back_refs.keys()
]
to_delete = set(
db.session.query(child_class).filter(and_(*delete_filters))
).difference(set(added))
for o in to_delete:
logger.debug("Deleting %s %s", child, str(obj))
db.session.delete(o)
return obj
def export_to_dict(
self,
recursive: bool = True,
include_parent_ref: bool = False,
include_defaults: bool = False,
export_uuids: bool = False,
) -> dict[Any, Any]:
"""Export obj to dictionary"""
export_fields = set(self.export_fields)
if export_uuids:
export_fields.add("uuid")
if "id" in export_fields:
export_fields.remove("id")
cls = self.__class__
parent_excludes = set()
if recursive and not include_parent_ref:
parent_ref = cls.__mapper__.relationships.get(cls.export_parent)
if parent_ref:
parent_excludes = {c.name for c in parent_ref.local_columns}
dict_rep = {
c.name: getattr(self, c.name)
for c in cls.__table__.columns # type: ignore
if (
c.name in export_fields
and c.name not in parent_excludes
and (
include_defaults
or (
getattr(self, c.name) is not None
and (not c.default or getattr(self, c.name) != c.default.arg)
)
)
)
}
# sort according to export_fields using DSU (decorate, sort, undecorate)
order = {field: i for i, field in enumerate(self.export_fields)}
decorated_keys = [(order.get(k, len(order)), k) for k in dict_rep]
decorated_keys.sort()
dict_rep = {k: dict_rep[k] for _, k in decorated_keys}
if recursive:
for cld in self.export_children:
# sorting to make lists of children stable
dict_rep[cld] = sorted(
[
child.export_to_dict(
recursive=recursive,
include_parent_ref=include_parent_ref,
include_defaults=include_defaults,
)
for child in getattr(self, cld)
],
key=lambda k: sorted(str(k.items())),
)
return convert_uuids(dict_rep)
def override(self, obj: Any) -> None:
"""Overrides the plain fields of the dashboard."""
for field in obj.__class__.export_fields:
setattr(self, field, getattr(obj, field))
def copy(self) -> Any:
"""Creates a copy of the dashboard without relationships."""
new_obj = self.__class__()
new_obj.override(self)
return new_obj
def alter_params(self, **kwargs: Any) -> None:
params = self.params_dict
params.update(kwargs)
self.params = json.dumps(params)
def remove_params(self, param_to_remove: str) -> None:
params = self.params_dict
params.pop(param_to_remove, None)
self.params = json.dumps(params)
def reset_ownership(self) -> None:
"""object will belong to the user the current user"""
# make sure the object doesn't have relations to a user
# it will be filled by appbuilder on save
self.created_by = None
self.changed_by = None
# flask global context might not exist (in cli or tests for example)
self.owners = []
if g and hasattr(g, "user"):
self.owners = [g.user]
@property
def params_dict(self) -> dict[Any, Any]:
return json_to_dict(self.params)
@property
def template_params_dict(self) -> dict[Any, Any]:
return json_to_dict(self.template_params) # type: ignore
def _user(user: User) -> str:
if not user:
return ""
return escape(user)
class AuditMixinNullable(AuditMixin):
"""Altering the AuditMixin to use nullable fields
Allows creating objects programmatically outside of CRUD
"""
created_on = sa.Column(sa.DateTime, default=datetime.now, nullable=True)
changed_on = sa.Column(
sa.DateTime, default=datetime.now, onupdate=datetime.now, nullable=True
)
@declared_attr
def created_by_fk(self) -> sa.Column: # pylint: disable=arguments-renamed
return sa.Column(
sa.Integer,
sa.ForeignKey("ab_user.id"),
default=get_user_id,
nullable=True,
)
@declared_attr
def changed_by_fk(self) -> sa.Column: # pylint: disable=arguments-renamed
return sa.Column(
sa.Integer,
sa.ForeignKey("ab_user.id"),
default=get_user_id,
onupdate=get_user_id,
nullable=True,
)
@property
def created_by_name(self) -> str:
if self.created_by:
return escape(f"{self.created_by}")
return ""
@property
def changed_by_name(self) -> str:
if self.changed_by:
return escape(f"{self.changed_by}")
return ""
@renders("created_by")
def creator(self) -> Union[Markup, str]:
return _user(self.created_by)
@property
def changed_by_(self) -> Union[Markup, str]:
return _user(self.changed_by)
@renders("changed_on")
def changed_on_(self) -> Markup:
return Markup(f'<span class="no-wrap">{self.changed_on}</span>')
@renders("changed_on")
def changed_on_delta_humanized(self) -> str:
return self.changed_on_humanized
@renders("changed_on")
def changed_on_dttm(self) -> float:
return datetime_to_epoch(self.changed_on)
@renders("created_on")
def created_on_delta_humanized(self) -> str:
return self.created_on_humanized
@renders("changed_on")
def changed_on_utc(self) -> str:
# Convert naive datetime to UTC
return self.changed_on.astimezone(pytz.utc).strftime("%Y-%m-%dT%H:%M:%S.%f%z")
@property
def changed_on_humanized(self) -> str:
return humanize.naturaltime(datetime.now() - self.changed_on)
@property
def created_on_humanized(self) -> str:
return humanize.naturaltime(datetime.now() - self.created_on)
@renders("changed_on")
def modified(self) -> Markup:
return Markup(f'<span class="no-wrap">{self.changed_on_humanized}</span>')
class QueryResult: # pylint: disable=too-few-public-methods
"""Object returned by the query interface"""
def __init__( # pylint: disable=too-many-arguments
self,
df: pd.DataFrame,
query: str,
duration: timedelta,
applied_template_filters: Optional[list[str]] = None,
applied_filter_columns: Optional[list[ColumnTyping]] = None,
rejected_filter_columns: Optional[list[ColumnTyping]] = None,
status: str = QueryStatus.SUCCESS,
error_message: Optional[str] = None,
errors: Optional[list[dict[str, Any]]] = None,
from_dttm: Optional[datetime] = None,
to_dttm: Optional[datetime] = None,
) -> None:
self.df = df
self.query = query
self.duration = duration
self.applied_template_filters = applied_template_filters or []
self.applied_filter_columns = applied_filter_columns or []
self.rejected_filter_columns = rejected_filter_columns or []
self.status = status
self.error_message = error_message
self.errors = errors or []
self.from_dttm = from_dttm
self.to_dttm = to_dttm
self.sql_rowcount = len(self.df.index) if not self.df.empty else 0
class ExtraJSONMixin:
"""Mixin to add an `extra` column (JSON) and utility methods"""
extra_json = sa.Column(MediumText(), default="{}")
@property
def extra(self) -> dict[str, Any]:
try:
return json.loads(self.extra_json or "{}") or {}
except (TypeError, json.JSONDecodeError) as exc:
logger.error(
"Unable to load an extra json: %r. Leaving empty.", exc, exc_info=True
)
return {}
@extra.setter
def extra(self, extras: dict[str, Any]) -> None:
self.extra_json = json.dumps(extras)
def set_extra_json_key(self, key: str, value: Any) -> None:
extra = self.extra
extra[key] = value
self.extra_json = json.dumps(extra)
@validates("extra_json")
def ensure_extra_json_is_not_none(
self,
_: str,
value: Optional[dict[str, Any]],
) -> Any:
if value is None:
return "{}"
return value
class CertificationMixin:
"""Mixin to add extra certification fields"""
extra = sa.Column(sa.Text, default="{}")
def get_extra_dict(self) -> dict[str, Any]:
try:
return json.loads(self.extra)
except (TypeError, json.JSONDecodeError):
return {}
@property
def is_certified(self) -> bool:
return bool(self.get_extra_dict().get("certification"))
@property
def certified_by(self) -> Optional[str]:
return self.get_extra_dict().get("certification", {}).get("certified_by")
@property
def certification_details(self) -> Optional[str]:
return self.get_extra_dict().get("certification", {}).get("details")
@property
def warning_markdown(self) -> Optional[str]:
return self.get_extra_dict().get("warning_markdown")
def clone_model(
target: Model,
ignore: Optional[list[str]] = None,
keep_relations: Optional[list[str]] = None,
**kwargs: Any,
) -> Model:
"""
Clone a SQLAlchemy model. By default will only clone naive column attributes.
To include relationship attributes, use `keep_relations`.
"""
ignore = ignore or []
table = target.__table__
primary_keys = table.primary_key.columns.keys()
data = {
attr: getattr(target, attr)
for attr in list(table.columns.keys()) + (keep_relations or [])
if attr not in primary_keys and attr not in ignore
}
data.update(kwargs)
return target.__class__(**data)
# todo(hugh): centralize where this code lives
class QueryStringExtended(NamedTuple):
applied_template_filters: Optional[list[str]]
applied_filter_columns: list[ColumnTyping]
rejected_filter_columns: list[ColumnTyping]
labels_expected: list[str]
prequeries: list[str]
sql: str
class SqlaQuery(NamedTuple):
applied_template_filters: list[str]
applied_filter_columns: list[ColumnTyping]
rejected_filter_columns: list[ColumnTyping]
cte: Optional[str]
extra_cache_keys: list[Any]
labels_expected: list[str]
prequeries: list[str]
sqla_query: Select
class ExploreMixin: # pylint: disable=too-many-public-methods
"""
Allows any flask_appbuilder.Model (Query, Table, etc.)
to be used to power a chart inside /explore
"""
sqla_aggregations = {
"COUNT_DISTINCT": lambda column_name: sa.func.COUNT(sa.distinct(column_name)),
"COUNT": sa.func.COUNT,
"SUM": sa.func.SUM,
"AVG": sa.func.AVG,
"MIN": sa.func.MIN,
"MAX": sa.func.MAX,
}
fetch_values_predicate = None
@property
def type(self) -> str:
raise NotImplementedError()
@property
def db_extra(self) -> Optional[dict[str, Any]]:
raise NotImplementedError()
def query(self, query_obj: QueryObjectDict) -> QueryResult:
raise NotImplementedError()
@property
def database_id(self) -> int:
raise NotImplementedError()
@property
def owners_data(self) -> list[Any]:
raise NotImplementedError()
@property
def metrics(self) -> list[Any]:
return []
@property
def uid(self) -> str:
raise NotImplementedError()
@property
def is_rls_supported(self) -> bool:
raise NotImplementedError()
@property
def cache_timeout(self) -> int:
raise NotImplementedError()
@property
def column_names(self) -> list[str]:
raise NotImplementedError()
@property
def offset(self) -> int:
raise NotImplementedError()
@property
def main_dttm_col(self) -> Optional[str]:
raise NotImplementedError()
@property
def always_filter_main_dttm(self) -> Optional[bool]:
return False
@property
def dttm_cols(self) -> list[str]:
raise NotImplementedError()
@property
def db_engine_spec(self) -> builtins.type["BaseEngineSpec"]:
raise NotImplementedError()
@property
def database(self) -> "Database":
raise NotImplementedError()
@property
def catalog(self) -> str:
raise NotImplementedError()
@property
def schema(self) -> str:
raise NotImplementedError()
@property
def sql(self) -> str:
raise NotImplementedError()
@property
def columns(self) -> list[Any]:
raise NotImplementedError()
def get_extra_cache_keys(self, query_obj: dict[str, Any]) -> list[Hashable]:
raise NotImplementedError()
def get_template_processor(self, **kwargs: Any) -> BaseTemplateProcessor:
raise NotImplementedError()
def get_fetch_values_predicate(
self,
template_processor: Optional[ # pylint: disable=unused-argument
BaseTemplateProcessor
] = None,
) -> TextClause:
return self.fetch_values_predicate
def get_sqla_row_level_filters(
self,
template_processor: Optional[BaseTemplateProcessor] = None, # pylint: disable=unused-argument
) -> list[TextClause]:
# TODO: We should refactor this mixin and remove this method
# as it exists in the BaseDatasource and is not applicable
# for datasources of type query
return []
def _process_sql_expression( # pylint: disable=too-many-arguments
self,
expression: Optional[str],
database_id: int,
engine: str,
schema: str,
template_processor: Optional[BaseTemplateProcessor],
) -> Optional[str]:
if template_processor and expression:
expression = template_processor.process_template(expression)
if expression:
expression = validate_adhoc_subquery(
expression,
database_id,
engine,
schema,
)
try:
expression = sanitize_clause(expression)
except QueryClauseValidationException as ex:
raise QueryObjectValidationError(ex.message) from ex
return expression
def make_sqla_column_compatible(
self, sqla_col: ColumnElement, label: Optional[str] = None
) -> ColumnElement:
"""Takes a sqlalchemy column object and adds label info if supported by engine.
:param sqla_col: sqlalchemy column instance
:param label: alias/label that column is expected to have
:return: either a sql alchemy column or label instance if supported by engine
"""
label_expected = label or sqla_col.name
db_engine_spec = self.db_engine_spec
# add quotes to tables
if db_engine_spec.get_allows_alias_in_select(self.database):
label = db_engine_spec.make_label_compatible(label_expected)
sqla_col = sqla_col.label(label)
sqla_col.key = label_expected
return sqla_col
@staticmethod
def _apply_cte(sql: str, cte: Optional[str]) -> str:
"""
Append a CTE before the SELECT statement if defined
:param sql: SELECT statement
:param cte: CTE statement
:return:
"""
if cte:
sql = f"{cte}\n{sql}"
return sql
def get_query_str_extended(
self,
query_obj: QueryObjectDict,
mutate: bool = True,
) -> QueryStringExtended:
sqlaq = self.get_sqla_query(**query_obj)
sql = self.database.compile_sqla_query(
sqlaq.sqla_query,
catalog=self.catalog,
schema=self.schema,
is_virtual=bool(self.sql),
)
sql = self._apply_cte(sql, sqlaq.cte)
if mutate:
sql = self.database.mutate_sql_based_on_config(sql)
return QueryStringExtended(
applied_template_filters=sqlaq.applied_template_filters,
applied_filter_columns=sqlaq.applied_filter_columns,
rejected_filter_columns=sqlaq.rejected_filter_columns,
labels_expected=sqlaq.labels_expected,
prequeries=sqlaq.prequeries,
sql=sql,
)
def _normalize_prequery_result_type(
self,
row: pd.Series,
dimension: str,
columns_by_name: dict[str, "TableColumn"],
) -> Union[str, int, float, bool, str]:
"""
Convert a prequery result type to its equivalent Python type.
Some databases like Druid will return timestamps as strings, but do not perform
automatic casting when comparing these strings to a timestamp. For cases like
this we convert the value via the appropriate SQL transform.
:param row: A prequery record
:param dimension: The dimension name
:param columns_by_name: The mapping of columns by name
:return: equivalent primitive python type
"""
value = row[dimension]
if isinstance(value, np.generic):
value = value.item()
column_ = columns_by_name[dimension]
db_extra: dict[str, Any] = self.database.get_extra()
if isinstance(column_, dict):
if (
column_.get("type")
and column_.get("is_temporal")
and isinstance(value, str)
):
sql = self.db_engine_spec.convert_dttm(
column_.get("type"), dateutil.parser.parse(value), db_extra=None
)
if sql:
value = self.db_engine_spec.get_text_clause(sql)
else:
if column_.type and column_.is_temporal and isinstance(value, str):
sql = self.db_engine_spec.convert_dttm(
column_.type, dateutil.parser.parse(value), db_extra=db_extra
)
if sql:
value = self.text(sql)
return value
def make_orderby_compatible(
self, select_exprs: list[ColumnElement], orderby_exprs: list[ColumnElement]
) -> None:
"""
If needed, make sure aliases for selected columns are not used in
`ORDER BY`.
In some databases (e.g. Presto), `ORDER BY` clause is not able to
automatically pick the source column if a `SELECT` clause alias is named
the same as a source column. In this case, we update the SELECT alias to
another name to avoid the conflict.
"""
if self.db_engine_spec.allows_alias_to_source_column:
return
def is_alias_used_in_orderby(col: ColumnElement) -> bool:
if not isinstance(col, Label):
return False
regexp = re.compile(f"\\(.*\\b{re.escape(col.name)}\\b.*\\)", re.IGNORECASE)
return any(regexp.search(str(x)) for x in orderby_exprs)
# Iterate through selected columns, if column alias appears in orderby
# use another `alias`. The final output columns will still use the
# original names, because they are updated by `labels_expected` after
# querying.
for col in select_exprs:
if is_alias_used_in_orderby(col):
col.name = f"{col.name}__"
def exc_query(self, qry: Any) -> QueryResult:
qry_start_dttm = datetime.now()
query_str_ext = self.get_query_str_extended(qry)
sql = query_str_ext.sql
status = QueryStatus.SUCCESS
errors = None
error_message = None
def assign_column_label(df: pd.DataFrame) -> Optional[pd.DataFrame]:
"""
Some engines change the case or generate bespoke column names, either by
default or due to lack of support for aliasing. This function ensures that
the column names in the DataFrame correspond to what is expected by
the viz components.
Sometimes a query may also contain only order by columns that are not used
as metrics or groupby columns, but need to present in the SQL `select`,
filtering by `labels_expected` make sure we only return columns users want.
:param df: Original DataFrame returned by the engine