updated engine

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)",