Skip to content

Commit 85a7d5c

Browse files
betodealmeidabetodealmeida
authored
feat: improve SQLite DB engine spec (#24909)
1 parent 3b81538 commit 85a7d5c

File tree

5 files changed

+383
-37
lines changed

5 files changed

+383
-37
lines changed

‎superset/db_engine_specs/gsheets.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
from superset import security_manager
3333
from superset.constants import PASSWORD_MASK
3434
from superset.databases.schemas import encrypted_field_properties, EncryptedString
35-
from superset.db_engine_specs.sqlite import SqliteEngineSpec
35+
from superset.db_engine_specs.shillelagh import ShillelaghEngineSpec
3636
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
3737

3838
if TYPE_CHECKING:
@@ -65,14 +65,13 @@ class GSheetsPropertiesType(TypedDict):
6565
catalog: dict[str, str]
6666

6767

68-
class GSheetsEngineSpec(SqliteEngineSpec):
68+
class GSheetsEngineSpec(ShillelaghEngineSpec):
6969
"""Engine for Google spreadsheets"""
7070

71-
engine = "gsheets"
7271
engine_name = "Google Sheets"
72+
engine = "gsheets"
7373
allows_joins = True
7474
allows_subqueries = True
75-
disable_ssh_tunneling = True
7675

7776
parameters_schema = GSheetsParametersSchema()
7877
default_driver = "apsw"

‎superset/db_engine_specs/shillelagh.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,15 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17+
from __future__ import annotations
18+
19+
from typing import TYPE_CHECKING
20+
1721
from superset.db_engine_specs.sqlite import SqliteEngineSpec
1822

23+
if TYPE_CHECKING:
24+
from superset.models.core import Database
25+
1926

2027
class ShillelaghEngineSpec(SqliteEngineSpec):
2128
"""Engine for shillelagh"""
@@ -28,3 +35,14 @@ class ShillelaghEngineSpec(SqliteEngineSpec):
2835

2936
allows_joins = True
3037
allows_subqueries = True
38+
39+
@classmethod
40+
def get_function_names(
41+
cls,
42+
database: Database,
43+
) -> list[str]:
44+
return super().get_function_names(database) + [
45+
"sleep",
46+
"version",
47+
"get_metadata",
48+
]

‎superset/db_engine_specs/sqlite.py‎

Lines changed: 171 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,13 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17+
18+
from __future__ import annotations
19+
1720
import re
1821
from datetime import datetime
1922
from re import Pattern
20-
from typing import Any, Optional, TYPE_CHECKING
23+
from typing import Any, TYPE_CHECKING
2124

2225
from flask_babel import gettext as __
2326
from sqlalchemy import types
@@ -39,11 +42,41 @@ class SqliteEngineSpec(BaseEngineSpec):
3942
engine = "sqlite"
4043
engine_name = "SQLite"
4144

45+
disable_ssh_tunneling = True
46+
4247
_time_grain_expressions = {
4348
None: "{col}",
4449
TimeGrain.SECOND: "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:%S', {col}))",
50+
TimeGrain.FIVE_SECONDS: (
51+
"DATETIME({col}, printf('-%d seconds', "
52+
"CAST(strftime('%S', {col}) AS INT) % 5))"
53+
),
54+
TimeGrain.THIRTY_SECONDS: (
55+
"DATETIME({col}, printf('-%d seconds', "
56+
"CAST(strftime('%S', {col}) AS INT) % 30))"
57+
),
4558
TimeGrain.MINUTE: "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}))",
59+
TimeGrain.FIVE_MINUTES: (
60+
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
61+
"CAST(strftime('%M', {col}) AS INT) % 5))"
62+
),
63+
TimeGrain.TEN_MINUTES: (
64+
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
65+
"CAST(strftime('%M', {col}) AS INT) % 10))"
66+
),
67+
TimeGrain.FIFTEEN_MINUTES: (
68+
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
69+
"CAST(strftime('%M', {col}) AS INT) % 15))"
70+
),
71+
TimeGrain.THIRTY_MINUTES: (
72+
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
73+
"CAST(strftime('%M', {col}) AS INT) % 30))"
74+
),
4675
TimeGrain.HOUR: "DATETIME(STRFTIME('%Y-%m-%dT%H:00:00', {col}))",
76+
TimeGrain.SIX_HOURS: (
77+
"DATETIME(STRFTIME('%Y-%m-%dT%H:00:00', {col}), printf('-%d hours', "
78+
"CAST(strftime('%H', {col}) AS INT) % 6))"
79+
),
4780
TimeGrain.DAY: "DATETIME({col}, 'start of day')",
4881
TimeGrain.WEEK: "DATETIME({col}, 'start of day', \
4982
-strftime('%w', {col}) || ' days')",
@@ -62,6 +95,13 @@ class SqliteEngineSpec(BaseEngineSpec):
6295
"DATETIME({col}, 'start of day', 'weekday 1', '-7 days')"
6396
),
6497
}
98+
# not sure why these are diffenret
99+
_time_grain_expressions.update(
100+
{
101+
TimeGrain.HALF_HOUR: _time_grain_expressions[TimeGrain.THIRTY_MINUTES],
102+
TimeGrain.QUARTER_YEAR: _time_grain_expressions[TimeGrain.QUARTER],
103+
}
104+
)
65105

66106
custom_errors: dict[Pattern[str], tuple[str, SupersetErrorType, dict[str, Any]]] = {
67107
COLUMN_DOES_NOT_EXIST_REGEX: (
@@ -77,16 +117,143 @@ def epoch_to_dttm(cls) -> str:
77117

78118
@classmethod
79119
def convert_dttm(
80-
cls, target_type: str, dttm: datetime, db_extra: Optional[dict[str, Any]] = None
81-
) -> Optional[str]:
120+
cls, target_type: str, dttm: datetime, db_extra: dict[str, Any] | None = None
121+
) -> str | None:
82122
sqla_type = cls.get_sqla_column_type(target_type)
83123
if isinstance(sqla_type, (types.String, types.DateTime)):
84124
return f"""'{dttm.isoformat(sep=" ", timespec="seconds")}'"""
85125
return None
86126

87127
@classmethod
88128
def get_table_names(
89-
cls, database: "Database", inspector: Inspector, schema: Optional[str]
129+
cls,
130+
database: Database,
131+
inspector: Inspector,
132+
schema: str | None,
90133
) -> set[str]:
91134
"""Need to disregard the schema for Sqlite"""
92135
return set(inspector.get_table_names())
136+
137+
@classmethod
138+
def get_function_names(
139+
cls,
140+
database: Database,
141+
) -> list[str]:
142+
"""
143+
Return function names.
144+
"""
145+
return [
146+
"abs",
147+
"acos",
148+
"acosh",
149+
"asin",
150+
"asinh",
151+
"atan",
152+
"atan2",
153+
"atanh",
154+
"avg",
155+
"ceil",
156+
"ceiling",
157+
"changes",
158+
"char",
159+
"coalesce",
160+
"cos",
161+
"cosh",
162+
"count",
163+
"cume_dist",
164+
"date",
165+
"datetime",
166+
"degrees",
167+
"dense_rank",
168+
"exp",
169+
"first_value",
170+
"floor",
171+
"format",
172+
"glob",
173+
"group_concat",
174+
"hex",
175+
"ifnull",
176+
"iif",
177+
"instr",
178+
"json",
179+
"json_array",
180+
"json_array_length",
181+
"json_each",
182+
"json_error_position",
183+
"json_extract",
184+
"json_group_array",
185+
"json_group_object",
186+
"json_insert",
187+
"json_object",
188+
"json_patch",
189+
"json_quote",
190+
"json_remove",
191+
"json_replace",
192+
"json_set",
193+
"json_tree",
194+
"json_type",
195+
"json_valid",
196+
"julianday",
197+
"lag",
198+
"last_insert_rowid",
199+
"last_value",
200+
"lead",
201+
"length",
202+
"like",
203+
"likelihood",
204+
"likely",
205+
"ln",
206+
"load_extension",
207+
"log",
208+
"log10",
209+
"log2",
210+
"lower",
211+
"ltrim",
212+
"max",
213+
"min",
214+
"mod",
215+
"nth_value",
216+
"ntile",
217+
"nullif",
218+
"percent_rank",
219+
"pi",
220+
"pow",
221+
"power",
222+
"printf",
223+
"quote",
224+
"radians",
225+
"random",
226+
"randomblob",
227+
"rank",
228+
"replace",
229+
"round",
230+
"row_number",
231+
"rtrim",
232+
"sign",
233+
"sin",
234+
"sinh",
235+
"soundex",
236+
"sqlite_compileoption_get",
237+
"sqlite_compileoption_used",
238+
"sqlite_offset",
239+
"sqlite_source_id",
240+
"sqlite_version",
241+
"sqrt",
242+
"strftime",
243+
"substr",
244+
"substring",
245+
"sum",
246+
"tan",
247+
"tanh",
248+
"time",
249+
"total_changes",
250+
"trim",
251+
"trunc",
252+
"typeof",
253+
"unhex",
254+
"unicode",
255+
"unixepoch",
256+
"unlikely",
257+
"upper",
258+
"zeroblob",
259+
]

0 commit comments

Comments
 (0)