add db_engine_spec

Author: Yibo-Chen13

timeplus_connect/tp_superset/__init__.py

@@ -0,0 +1,138 @@
+import logging
+import re
+from datetime import datetime
+from marshmallow import fields, Schema
+from marshmallow.validate import Range
+from typing import Any
+from sqlalchemy.engine.url import URL
+
+from flask import current_app
+from flask_babel import gettext as __
+from sqlalchemy import types
+from superset.db_engine_specs import BaseEngineSpec
+from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+
+try:
+    from timeplus_connect.common import set_setting
+    from timeplus_connect.datatypes.format import set_default_formats
+
+    # override default formats for compatibility
+    set_default_formats(
+        "fixed_string",
+        "string",
+        "ipv*",
+        "string",
+        "uint64",
+        "signed",
+        "UUID",
+        "string",
+        "*int256",
+        "string",
+        "*int128",
+        "string",
+    )
+    set_setting(
+        "product_name",
+        f"superset/{current_app.config.get('VERSION_STRING', 'dev')}",
+    )
+except ImportError:  # Timeplus Connect not installed, do nothing
+    pass
+
+class TimeplusParametersSchema(Schema):
+    username = fields.String(allow_none=True, metadata={"description": __("Username")})
+    password = fields.String(allow_none=True, metadata={"description": __("Password")})
+    host = fields.String(
+        required=True, metadata={"description": __("Hostname or IP address")}
+    )
+    port = fields.Integer(
+        allow_none=True,
+        metadata={"description": __("Database port")},
+        validate=Range(min=0, max=65535),
+    )
+    database = fields.String(
+        allow_none=True, metadata={"description": __("Database name")}
+    )
+    encryption = fields.Boolean(
+        dump_default=True,
+        metadata={"description": __("Use an encrypted connection to the database")},
+    )
+    query = fields.Dict(
+        keys=fields.Str(),
+        values=fields.Raw(),
+        metadata={"description": __("Additional parameters")},
+    )
+
+
+class TimeplusEngineSpec(BaseEngineSpec):
+    """Engine spec for timeplus-connect connector"""
+
+    engine = "timeplusdb"
+    engine_name = "Timeplus Connect (Superset)"
+
+    _show_functions_column = "name"
+    supports_file_upload = False
+
+    sqlalchemy_uri_placeholder = (
+        "timeplusdb://user:password@host[:port][/dbname][?secure=value&=value...]"
+    )
+
+    parameters_schema = TimeplusParametersSchema()
+    encryption_parameters = {"secure": "true"}
+
+    @classmethod
+    def epoch_to_dttm(cls) -> str:
+        return "{col}"
+
+    @classmethod
+    def convert_dttm(
+        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.Date):
+            return f"to_date('{dttm.date().isoformat()}')"
+        if isinstance(sqla_type, types.DateTime):
+            return f"""to_datetime('{dttm.isoformat(sep=" ", timespec="seconds")}')"""
+        return None
+
+
+    @classmethod
+    def get_function_names(cls, database: Database) -> list[str]:
+        # pylint: disable=import-outside-toplevel, import-error
+        from timeplus_connect.driver.exceptions import TimeplusError
+
+        if cls._function_names:
+            return cls._function_names
+        try:
+            names = database.get_df(
+                "SELECT name FROM system.functions UNION ALL "
+                + "SELECT name FROM system.table_functions"
+            )["name"].tolist()
+            cls._function_names = names
+            return names
+        except TimeplusError:
+            logger.exception("Error retrieving system.functions")
+            return []
+    
+    @classmethod
+    def get_datatype(cls, type_code: str) -> str:
+        # keep it lowercase, as Timeplus types aren't typical SHOUTCASE ANSI SQL
+        return type_code
+
+    @classmethod
+    def build_sqlalchemy_uri(
+        cls,
+        parameters: dict[str, str],
+        encrypted_extra: dict[str, str] | None = None,
+    ) -> str:
+        url_params = parameters.copy()
+        if url_params.get("encryption"):
+            query = parameters.get("query", {}).copy()
+            query.update(cls.encryption_parameters)
+            url_params["query"] = query
+        if not url_params.get("database"):
+            url_params["database"] = "default"
+        url_params.pop("encryption", None)
+        return str(URL.create(f"{cls.engine}+{cls.default_driver}", **url_params))