diff --git a/superset_iris/engine.py b/superset_iris/engine.py index 4bfac23..c5c282a 100644 --- a/superset_iris/engine.py +++ b/superset_iris/engine.py @@ -4,10 +4,11 @@ from typing import ( Any, Dict, - ContextManager, + List, TYPE_CHECKING, Pattern, Tuple, + Optional, ) import pandas as pd @@ -18,11 +19,16 @@ BaseEngineSpec, BasicParametersType, BasicParametersMixin, + BasicPropertiesType, ) -from superset.errors import SupersetErrorType +from superset.constants import USER_AGENT +from superset.errors import ErrorLevel, SupersetErrorType, SupersetError from superset.sql_parse import Table +from superset.databases.utils import make_url_safe from sqlalchemy.engine.base import Engine +from sqlalchemy.engine.url import URL from superset.utils.core import GenericDataType +from typing_extensions import TypedDict from sqlalchemy_iris import BIT @@ -43,15 +49,23 @@ class IRISParametersSchema(Schema): - username = fields.String(allow_none=True, description=__("Username")) - password = fields.String(allow_none=True, description=__("Password")) - host = fields.String(required=True, description=__("Hostname or IP address")) - port = fields.Integer( - allow_none=True, - description=__("Database port"), - validate=Range(min=0, max=65535), - ) - database = fields.String(allow_none=True, description=__("Database name")) + username = fields.Str(required=True) + password = fields.Str(required=True) + host = fields.Str(required=True) + port = fields.Integer(required=True) + database = fields.Str(required=True) + + +class IRISParametersType(TypedDict): + username: str + password: str + host: str + port: int + database: str + + +class IRISPropertiesType(TypedDict): + parameters: IRISParametersType class IRISEngineSpec(BaseEngineSpec, BasicParametersMixin): @@ -62,11 +76,14 @@ class IRISEngineSpec(BaseEngineSpec, BasicParametersMixin): engine = "iris" engine_name = "InterSystems IRIS" + parameters_schema = IRISParametersSchema() + default_driver = "iris" + allow_limit_clause = False max_column_name_length = 50 - sqlalchemy_uri_placeholder = "iris://_SYSTEM:SYS@iris:1972/USER" + sqlalchemy_uri_placeholder = "iris://{username}:{password}@{host}:{port}/{database}" parameters_schema = IRISParametersSchema() column_type_mappings = ( @@ -82,6 +99,80 @@ class IRISEngineSpec(BaseEngineSpec, BasicParametersMixin): ), ) + @staticmethod + def get_extra_params(database: "Database") -> Dict[str, Any]: + """ + Add a user agent to be used in the connection. + """ + extra: Dict[str, Any] = BaseEngineSpec.get_extra_params(database) + engine_params: Dict[str, Any] = extra.setdefault("engine_params", {}) + connect_args: Dict[str, Any] = engine_params.setdefault("connect_args", {}) + + connect_args.setdefault("application_name", USER_AGENT) + + return extra + + def build_sqlalchemy_uri( + cls, + parameters: IRISParametersType, + encrypted_extra: Optional[ # pylint: disable=unused-argument + Dict[str, Any] + ] = None, + ) -> str: + url = str( + URL( + "iris", + username=parameters.get("username"), + password=parameters.get("password"), + host=parameters.get("host"), + port=parameters.get("port"), + database=parameters.get("database"), + ) + ) + return url + + @classmethod + def get_parameters_from_uri( + cls, + uri: str, + encrypted_extra: Optional[ # pylint: disable=unused-argument + Dict[str, str] + ] = None, + ) -> Any: + url = make_url_safe(uri) + return { + "username": url.username, + "password": url.password, + "host": url.host, + "port": url.port, + "database": url.database, + } + + @classmethod + def validate_parameters(cls, properties: IRISPropertiesType) -> List[SupersetError]: + errors: List[SupersetError] = [] + required = { + "username", + "password", + "host", + "port", + "database", + } + parameters = properties.get("parameters", {}) + present = {key for key in parameters if parameters.get(key, ())} + missing = sorted(required - present) + + if missing: + errors.append( + SupersetError( + message=f'One or more parameters are missing: {", ".join(missing)}', + error_type=SupersetErrorType.CONNECTION_MISSING_PARAMETERS_ERROR, + level=ErrorLevel.WARNING, + extra={"missing": missing}, + ), + ) + return errors + _time_grain_expressions = { None: "{col}", "PT1S": "CAST(TO_CHAR(CAST({col} AS TIMESTAMP), 'YYYY-MM-DD HH24:MM:SS') AS DATETIME)",