superset/db_engine_specs/ocient.py
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import contextlib
import re
import threading
from re import Pattern
from typing import Any, Callable, NamedTuple, Optional
from flask_babel import gettext as __
from sqlalchemy.engine.reflection import Inspector
with contextlib.suppress(ImportError, RuntimeError): # pyocient may not be installed
# Ensure pyocient inherits Superset's logging level
import geojson
import pyocient
from shapely import wkt
from superset import app
superset_log_level = app.config["LOG_LEVEL"]
pyocient.logger.setLevel(superset_log_level)
from superset.constants import TimeGrain
from superset.db_engine_specs.base import BaseEngineSpec
from superset.errors import SupersetErrorType
from superset.models.core import Database
from superset.models.sql_lab import Query
# Regular expressions to catch custom errors
CONNECTION_INVALID_USERNAME_REGEX = re.compile(
r"The referenced user does not exist \(User '(?P<username>.*?)' not found\)"
)
CONNECTION_INVALID_PASSWORD_REGEX = re.compile(
r"The userid/password combination was not valid \(Incorrect password for user\)"
)
CONNECTION_INVALID_HOSTNAME_REGEX = re.compile(
r"Unable to connect to (?P<host>.*?):(?P<port>.*?)"
)
CONNECTION_UNKNOWN_DATABASE_REGEX = re.compile(
r"No database named '(?P<database>.*?)' exists"
)
CONNECTION_INVALID_PORT_ERROR = re.compile("Port out of range 0-65535")
INVALID_CONNECTION_STRING_REGEX = re.compile(
r"An invalid connection string attribute was specified"
r" \(failed to decrypt cipher text\)"
)
SYNTAX_ERROR_REGEX = re.compile(
r"There is a syntax error in your statement \((?P<qualifier>.*?)"
r" input '(?P<input>.*?)' expecting (?P<expected>.*?)\)"
)
TABLE_DOES_NOT_EXIST_REGEX = re.compile(
r"The referenced table or view '(?P<table>.*?)' does not exist"
)
COLUMN_DOES_NOT_EXIST_REGEX = re.compile(
r"The reference to column '(?P<column>.*?)' is not valid"
)
# Custom datatype conversion functions
def _to_hex(data: bytes) -> str:
"""
Converts the bytes object into a string of hexadecimal digits.
:param data: the bytes object
:returns: string of hexadecimal digits representing the bytes
"""
return data.hex()
def _wkt_to_geo_json(geo_as_wkt: str) -> Any:
"""
Converts pyocient geometry objects to their geoJSON representation.
:param geo_as_wkt: the GIS object in WKT format
:returns: the geoJSON encoding of `geo`
"""
# Need to try-catch here because these deps may not be installed
geo = wkt.loads(geo_as_wkt)
return geojson.Feature(geometry=geo, properties={})
def _point_list_to_wkt(
points, # type: list[pyocient._STPoint]
) -> str:
"""
Converts the list of pyocient._STPoint elements to a WKT LineString.
:param points: the list of pyocient._STPoint objects
:returns: WKT LineString
"""
coords = [f"{p.long} {p.lat}" for p in points]
return f"LINESTRING({', '.join(coords)})"
def _point_to_geo_json(
point, # type: pyocient._STPoint
) -> Any:
"""
Converts the pyocient._STPolygon object to the geoJSON format
:param point: the pyocient._STPoint instance
:returns: the geoJSON encoding of this point
"""
wkt_point = str(point)
return _wkt_to_geo_json(wkt_point)
def _linestring_to_geo_json(
linestring, # type: pyocient._STLinestring
) -> Any:
"""
Converts the pyocient._STLinestring object to a GIS format
compatible with the Superset visualization toolkit (powered
by Deck.gl).
:param linestring: the pyocient._STLinestring instance
:returns: the geoJSON of this linestring
"""
if len(linestring.points) == 1:
# While technically an invalid linestring object, Ocient
# permits ST_LINESTRING containers to contain a single
# point. The flexibility allows the database to encode
# geometry collections as an array of the highest dimensional
# element in the collection (i.e. ST_LINESTRING[] or
# ST_POLYGON[]).
point = linestring.points[0]
return _point_to_geo_json(point)
wkt_linestring = str(linestring)
return _wkt_to_geo_json(wkt_linestring)
def _polygon_to_geo_json(
polygon, # type: pyocient._STPolygon
) -> Any:
"""
Converts the pyocient._STPolygon object to a GIS format
compatible with the Superset visualization toolkit (powered
by Deck.gl).
:param polygon: the pyocient._STPolygon instance
:returns: the geoJSON encoding of this polygon
"""
if len(polygon.exterior) > 0 and len(polygon.holes) == 0:
if len(polygon.exterior) == 1:
# The exterior ring contains a single ST_POINT
point = polygon.exterior[0]
return _point_to_geo_json(point)
if polygon.exterior[0] != polygon.exterior[-1]:
# The exterior ring contains an open ST_LINESTRING
wkt_linestring = _point_list_to_wkt(polygon.exterior)
return _wkt_to_geo_json(wkt_linestring)
# else
# This is a valid ST_POLYGON
wkt_polygon = str(polygon)
return _wkt_to_geo_json(wkt_polygon)
# Sanitization function for column values
SanitizeFunc = Callable[[Any], Any]
# Represents a pair of a column index and the sanitization function
# to apply to its values.
class PlacedSanitizeFunc(NamedTuple):
column_index: int
sanitize_func: SanitizeFunc
# This map contains functions used to sanitize values for column types
# that cannot be processed natively by Superset.
#
# Superset serializes temporal objects using a custom serializer
# defined in superset/utils/core.py (#json_int_dttm_ser(...)). Other
# are serialized by the default JSON encoder.
#
# Need to try-catch here because pyocient may not be installed
try:
from pyocient import TypeCodes
_sanitized_ocient_type_codes: dict[int, SanitizeFunc] = {
TypeCodes.BINARY: _to_hex,
TypeCodes.ST_POINT: _point_to_geo_json,
TypeCodes.IP: str,
TypeCodes.IPV4: str,
TypeCodes.ST_LINESTRING: _linestring_to_geo_json,
TypeCodes.ST_POLYGON: _polygon_to_geo_json,
}
except ImportError:
_sanitized_ocient_type_codes = {}
def _find_columns_to_sanitize(cursor: Any) -> list[PlacedSanitizeFunc]:
"""
Cleans the column value for consumption by Superset.
:param cursor: the result set cursor
:returns: the list of tuples consisting of the column index and sanitization function
"""
return [
PlacedSanitizeFunc(i, _sanitized_ocient_type_codes[cursor.description[i][1]])
for i in range(len(cursor.description))
if cursor.description[i][1] in _sanitized_ocient_type_codes
]
class OcientEngineSpec(BaseEngineSpec):
engine = "ocient"
engine_name = "Ocient"
# limit_method = LimitMethod.WRAP_SQL
force_column_alias_quotes = True
max_column_name_length = 30
allows_cte_in_subquery = False
# Ocient does not support cte names starting with underscores
cte_alias = "cte__"
# Store mapping of superset Query id -> Ocient ID
# These are inserted into the cache when executing the query
# They are then removed, either upon cancellation or query completion
query_id_mapping: dict[str, str] = {}
query_id_mapping_lock = threading.Lock()
custom_errors: dict[Pattern[str], tuple[str, SupersetErrorType, dict[str, Any]]] = {
CONNECTION_INVALID_USERNAME_REGEX: (
__('The username "%(username)s" does not exist.'),
SupersetErrorType.CONNECTION_INVALID_USERNAME_ERROR,
{},
),
CONNECTION_INVALID_PASSWORD_REGEX: (
__(
"The user/password combination is not valid"
" (Incorrect password for user)."
),
SupersetErrorType.CONNECTION_INVALID_PASSWORD_ERROR,
{},
),
CONNECTION_UNKNOWN_DATABASE_REGEX: (
__('Could not connect to database: "%(database)s"'),
SupersetErrorType.CONNECTION_UNKNOWN_DATABASE_ERROR,
{},
),
CONNECTION_INVALID_HOSTNAME_REGEX: (
__('Could not resolve hostname: "%(host)s".'),
SupersetErrorType.CONNECTION_INVALID_HOSTNAME_ERROR,
{},
),
CONNECTION_INVALID_PORT_ERROR: (
__("Port out of range 0-65535"),
SupersetErrorType.CONNECTION_INVALID_PORT_ERROR,
{},
),
INVALID_CONNECTION_STRING_REGEX: (
__(
"Invalid Connection String: Expecting String of"
" the form 'ocient://user:pass@host:port/database'."
),
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
),
SYNTAX_ERROR_REGEX: (
__('Syntax Error: %(qualifier)s input "%(input)s" expecting "%(expected)s'),
SupersetErrorType.SYNTAX_ERROR,
{},
),
TABLE_DOES_NOT_EXIST_REGEX: (
__('Table or View "%(table)s" does not exist.'),
SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR,
{},
),
COLUMN_DOES_NOT_EXIST_REGEX: (
__('Invalid reference to column: "%(column)s"'),
SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR,
{},
),
}
_time_grain_expressions = {
None: "{col}",
TimeGrain.SECOND: "ROUND({col}, 'SECOND')",
TimeGrain.MINUTE: "ROUND({col}, 'MINUTE')",
TimeGrain.HOUR: "ROUND({col}, 'HOUR')",
TimeGrain.DAY: "ROUND({col}, 'DAY')",
TimeGrain.WEEK: "ROUND({col}, 'WEEK')",
TimeGrain.MONTH: "ROUND({col}, 'MONTH')",
TimeGrain.QUARTER_YEAR: "ROUND({col}, 'QUARTER')",
TimeGrain.YEAR: "ROUND({col}, 'YEAR')",
}
@classmethod
def get_table_names(
cls, database: Database, inspector: Inspector, schema: Optional[str]
) -> set[str]:
return inspector.get_table_names(schema)
@classmethod
def fetch_data(
cls, cursor: Any, limit: Optional[int] = None
) -> list[tuple[Any, ...]]:
try:
rows: list[tuple[Any, ...]] = super().fetch_data(cursor, limit)
except Exception as exception:
with OcientEngineSpec.query_id_mapping_lock:
del OcientEngineSpec.query_id_mapping[
getattr(cursor, "superset_query_id")
]
raise exception
# TODO: Unsure if we need to verify that we are receiving rows:
if len(rows) > 0 and type(rows[0]).__name__ == "Row":
# Peek at the schema to determine which column values, if any,
# require sanitization.
columns_to_sanitize: list[PlacedSanitizeFunc] = _find_columns_to_sanitize(
cursor
)
if columns_to_sanitize:
# At least 1 column has to be sanitized.
def identity(x: Any) -> Any:
return x
# Use the identity function if the column type doesn't need to be
# sanitized.
sanitization_functions: list[SanitizeFunc] = [
identity for _ in range(len(cursor.description))
]
for info in columns_to_sanitize:
sanitization_functions[info.column_index] = info.sanitize_func
# pyocient returns a list of NamedTuple objects which represent a
# single row. We have to do this copy because that data type is
# NamedTuple's are immutable.
rows = [
tuple(
sanitize_func(val)
for sanitize_func, val in zip(sanitization_functions, row)
)
for row in rows
]
return rows
@classmethod
def epoch_to_dttm(cls) -> str:
return "DATEADD(S, {col}, '1970-01-01')"
@classmethod
def epoch_ms_to_dttm(cls) -> str:
return "DATEADD(MS, {col}, '1970-01-01')"
@classmethod
def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]:
# Return a Non-None value
# If None is returned, Superset will not call cancel_query
return "DUMMY_VALUE"
@classmethod
def handle_cursor(cls, cursor: Any, query: Query) -> None:
with OcientEngineSpec.query_id_mapping_lock:
OcientEngineSpec.query_id_mapping[query.id] = cursor.query_id
# Add the query id to the cursor
setattr(cursor, "superset_query_id", query.id)
return super().handle_cursor(cursor, query)
@classmethod
def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool:
with OcientEngineSpec.query_id_mapping_lock:
if query.id in OcientEngineSpec.query_id_mapping:
cursor.execute(f"CANCEL {OcientEngineSpec.query_id_mapping[query.id]}")
# Query has been cancelled, so we can safely remove the cursor from
# the cache
del OcientEngineSpec.query_id_mapping[query.id]
return True
# If the query is not in the cache, it must have either been cancelled
# elsewhere or completed
return False