Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: improve SQLite DB engine spec #24909

Merged
merged 4 commits into from
Aug 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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

Check warning on line 24 in superset/db_engine_specs/shillelagh.py

View check run for this annotation

Codecov / codecov/patch

superset/db_engine_specs/shillelagh.py#L24

Added line #L24 was not covered by tests


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

allows_joins = True
allows_subqueries = True

@classmethod
def get_function_names(
cls,
database: Database,
) -> list[str]:
return super().get_function_names(database) + [

Check warning on line 44 in superset/db_engine_specs/shillelagh.py

View check run for this annotation

Codecov / codecov/patch

superset/db_engine_specs/shillelagh.py#L44

Added line #L44 was not covered by tests
"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 [
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@betodealmeida I'm somewhat surprised sqlite3 doesn't provide these.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I wish it exposed more of the internals of SQLite!

"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
Loading