Skip to content

Commit

Permalink
feat: improve SQLite DB engine spec (#24909)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored Aug 7, 2023
1 parent 3b81538 commit 85a7d5c
Show file tree
Hide file tree
Showing 5 changed files with 383 additions and 37 deletions.
7 changes: 3 additions & 4 deletions superset/db_engine_specs/gsheets.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from superset import security_manager
from superset.constants import PASSWORD_MASK
from superset.databases.schemas import encrypted_field_properties, EncryptedString
from superset.db_engine_specs.sqlite import SqliteEngineSpec
from superset.db_engine_specs.shillelagh import ShillelaghEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType

if TYPE_CHECKING:
Expand Down Expand Up @@ -65,14 +65,13 @@ class GSheetsPropertiesType(TypedDict):
catalog: dict[str, str]


class GSheetsEngineSpec(SqliteEngineSpec):
class GSheetsEngineSpec(ShillelaghEngineSpec):
"""Engine for Google spreadsheets"""

engine = "gsheets"
engine_name = "Google Sheets"
engine = "gsheets"
allows_joins = True
allows_subqueries = True
disable_ssh_tunneling = True

parameters_schema = GSheetsParametersSchema()
default_driver = "apsw"
Expand Down
18 changes: 18 additions & 0 deletions superset/db_engine_specs/shillelagh.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,15 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from typing import TYPE_CHECKING

from superset.db_engine_specs.sqlite import SqliteEngineSpec

if TYPE_CHECKING:
from superset.models.core import Database


class ShillelaghEngineSpec(SqliteEngineSpec):
"""Engine for shillelagh"""
Expand All @@ -28,3 +35,14 @@ class ShillelaghEngineSpec(SqliteEngineSpec):

allows_joins = True
allows_subqueries = True

@classmethod
def get_function_names(
cls,
database: Database,
) -> list[str]:
return super().get_function_names(database) + [
"sleep",
"version",
"get_metadata",
]
175 changes: 171 additions & 4 deletions superset/db_engine_specs/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,13 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

import re
from datetime import datetime
from re import Pattern
from typing import Any, Optional, TYPE_CHECKING
from typing import Any, TYPE_CHECKING

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

disable_ssh_tunneling = True

_time_grain_expressions = {
None: "{col}",
TimeGrain.SECOND: "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:%S', {col}))",
TimeGrain.FIVE_SECONDS: (
"DATETIME({col}, printf('-%d seconds', "
"CAST(strftime('%S', {col}) AS INT) % 5))"
),
TimeGrain.THIRTY_SECONDS: (
"DATETIME({col}, printf('-%d seconds', "
"CAST(strftime('%S', {col}) AS INT) % 30))"
),
TimeGrain.MINUTE: "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}))",
TimeGrain.FIVE_MINUTES: (
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
"CAST(strftime('%M', {col}) AS INT) % 5))"
),
TimeGrain.TEN_MINUTES: (
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
"CAST(strftime('%M', {col}) AS INT) % 10))"
),
TimeGrain.FIFTEEN_MINUTES: (
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
"CAST(strftime('%M', {col}) AS INT) % 15))"
),
TimeGrain.THIRTY_MINUTES: (
"DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}), printf('-%d minutes', "
"CAST(strftime('%M', {col}) AS INT) % 30))"
),
TimeGrain.HOUR: "DATETIME(STRFTIME('%Y-%m-%dT%H:00:00', {col}))",
TimeGrain.SIX_HOURS: (
"DATETIME(STRFTIME('%Y-%m-%dT%H:00:00', {col}), printf('-%d hours', "
"CAST(strftime('%H', {col}) AS INT) % 6))"
),
TimeGrain.DAY: "DATETIME({col}, 'start of day')",
TimeGrain.WEEK: "DATETIME({col}, 'start of day', \
-strftime('%w', {col}) || ' days')",
Expand All @@ -62,6 +95,13 @@ class SqliteEngineSpec(BaseEngineSpec):
"DATETIME({col}, 'start of day', 'weekday 1', '-7 days')"
),
}
# not sure why these are diffenret
_time_grain_expressions.update(
{
TimeGrain.HALF_HOUR: _time_grain_expressions[TimeGrain.THIRTY_MINUTES],
TimeGrain.QUARTER_YEAR: _time_grain_expressions[TimeGrain.QUARTER],
}
)

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

@classmethod
def convert_dttm(
cls, target_type: str, dttm: datetime, db_extra: Optional[dict[str, Any]] = None
) -> Optional[str]:
cls, target_type: str, dttm: datetime, db_extra: dict[str, Any] | None = None
) -> str | None:
sqla_type = cls.get_sqla_column_type(target_type)
if isinstance(sqla_type, (types.String, types.DateTime)):
return f"""'{dttm.isoformat(sep=" ", timespec="seconds")}'"""
return None

@classmethod
def get_table_names(
cls, database: "Database", inspector: Inspector, schema: Optional[str]
cls,
database: Database,
inspector: Inspector,
schema: str | None,
) -> set[str]:
"""Need to disregard the schema for Sqlite"""
return set(inspector.get_table_names())

@classmethod
def get_function_names(
cls,
database: Database,
) -> list[str]:
"""
Return function names.
"""
return [
"abs",
"acos",
"acosh",
"asin",
"asinh",
"atan",
"atan2",
"atanh",
"avg",
"ceil",
"ceiling",
"changes",
"char",
"coalesce",
"cos",
"cosh",
"count",
"cume_dist",
"date",
"datetime",
"degrees",
"dense_rank",
"exp",
"first_value",
"floor",
"format",
"glob",
"group_concat",
"hex",
"ifnull",
"iif",
"instr",
"json",
"json_array",
"json_array_length",
"json_each",
"json_error_position",
"json_extract",
"json_group_array",
"json_group_object",
"json_insert",
"json_object",
"json_patch",
"json_quote",
"json_remove",
"json_replace",
"json_set",
"json_tree",
"json_type",
"json_valid",
"julianday",
"lag",
"last_insert_rowid",
"last_value",
"lead",
"length",
"like",
"likelihood",
"likely",
"ln",
"load_extension",
"log",
"log10",
"log2",
"lower",
"ltrim",
"max",
"min",
"mod",
"nth_value",
"ntile",
"nullif",
"percent_rank",
"pi",
"pow",
"power",
"printf",
"quote",
"radians",
"random",
"randomblob",
"rank",
"replace",
"round",
"row_number",
"rtrim",
"sign",
"sin",
"sinh",
"soundex",
"sqlite_compileoption_get",
"sqlite_compileoption_used",
"sqlite_offset",
"sqlite_source_id",
"sqlite_version",
"sqrt",
"strftime",
"substr",
"substring",
"sum",
"tan",
"tanh",
"time",
"total_changes",
"trim",
"trunc",
"typeof",
"unhex",
"unicode",
"unixepoch",
"unlikely",
"upper",
"zeroblob",
]
Loading

0 comments on commit 85a7d5c

Please sign in to comment.