219 lines
8.0 KiB
Python
219 lines
8.0 KiB
Python
"""Compatibility helpers for openGauss/Dolphin's MySQL protocol.
|
|
|
|
The Dolphin endpoint deliberately presents itself as MySQL, but a few pieces
|
|
of server metadata and DDL reflection differ from real MySQL. These patches
|
|
are installed only when the deployment explicitly selects the ``gauss_mysql``
|
|
database product. A normal MySQL deployment never imports or installs them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
GAUSS_MYSQL_PRODUCT = "gauss_mysql"
|
|
GAUSS_MYSQL_SQLALCHEMY_VERSION = (5, 7, 20)
|
|
_DOLPHIN_INIT_COMMAND = "SET standard_conforming_strings TO off"
|
|
|
|
|
|
def is_gauss_mysql_product() -> bool:
|
|
"""Return whether this process was started for the Dolphin endpoint."""
|
|
|
|
return os.getenv("DEERFLOW_DATABASE_PRODUCT", "").strip().lower() == GAUSS_MYSQL_PRODUCT
|
|
|
|
|
|
def is_gauss_mysql_dialect(dialect: object) -> bool:
|
|
"""Return whether the selected SQLAlchemy dialect has the Dolphin hook."""
|
|
|
|
return bool(getattr(dialect, "_deerflow_gauss_mysql_compat_installed", False))
|
|
|
|
|
|
def _decode_scalar(value: Any) -> str:
|
|
if isinstance(value, bytes):
|
|
return value.decode(errors="replace")
|
|
return str(value)
|
|
|
|
|
|
def _gauss_mysql_default_schema_name(self, connection): # noqa: ANN001, ARG001
|
|
"""Return Dolphin's current schema rather than its physical database."""
|
|
|
|
schema = connection.exec_driver_sql("SELECT current_schema()").scalar()
|
|
if schema is None:
|
|
raise RuntimeError("GaussDB did not return current_schema()")
|
|
return _decode_scalar(schema)
|
|
|
|
|
|
def _gauss_mysql_server_version_info(self, connection): # noqa: ANN001
|
|
"""Expose a stable MySQL capability version for a Dolphin connection."""
|
|
|
|
dbapi_connection = connection.connection
|
|
cursor = dbapi_connection.cursor()
|
|
try:
|
|
cursor.execute("SELECT VERSION()")
|
|
row = cursor.fetchone()
|
|
finally:
|
|
cursor.close()
|
|
|
|
raw_version = "" if row is None or row[0] is None else _decode_scalar(row[0])
|
|
version = GAUSS_MYSQL_SQLALCHEMY_VERSION
|
|
self._deerflow_gauss_server_version = raw_version
|
|
self._set_mariadb(False, version)
|
|
self._mariadb_normalized_version_info = version
|
|
self.server_version_info = version
|
|
logger.info(
|
|
"GaussDB MySQL compatibility enabled: server=%r, SQLAlchemy capability=%s",
|
|
raw_version,
|
|
".".join(map(str, version)),
|
|
)
|
|
return version
|
|
|
|
|
|
def _gauss_mysql_isolation_level(self, dbapi_connection): # noqa: ANN001, ARG001
|
|
"""Read isolation without comparing an unparseable vendor version."""
|
|
|
|
cursor = dbapi_connection.cursor()
|
|
try:
|
|
cursor.execute("SELECT @@transaction_isolation")
|
|
row = cursor.fetchone()
|
|
finally:
|
|
cursor.close()
|
|
|
|
if row is None or row[0] is None:
|
|
raise NotImplementedError("GaussDB did not return a transaction isolation level")
|
|
return _decode_scalar(row[0]).upper().replace("-", " ")
|
|
|
|
|
|
def _gauss_mysql_has_table(self, connection, table_name, schema=None, **kw): # noqa: ANN001, ARG001
|
|
"""Check relation existence without a failing ``DESCRIBE`` transaction."""
|
|
|
|
self._ensure_has_table_connection(connection)
|
|
if schema is None:
|
|
schema = self.default_schema_name
|
|
if schema is None:
|
|
raise ValueError("GaussDB MySQL table check requires a database name")
|
|
|
|
statement = "SELECT 1 FROM pg_catalog.pg_class AS c JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace WHERE n.nspname = %s AND c.relname = %s AND c.relkind IN ('r', 'p', 'v', 'm', 'f') LIMIT 1"
|
|
with connection.exec_driver_sql(
|
|
statement,
|
|
(schema, table_name),
|
|
execution_options={"skip_user_error_events": True},
|
|
) as result:
|
|
return result.fetchone() is not None
|
|
|
|
|
|
def _gauss_mysql_create_connect_args(self, url): # noqa: ANN001
|
|
"""Add the Dolphin session setting required by PyMySQL-compatible drivers."""
|
|
|
|
args, kwargs = self._deerflow_gauss_mysql_original_create_connect_args(url)
|
|
existing = kwargs.get("init_command")
|
|
if existing and existing.strip().rstrip(";").lower() != _DOLPHIN_INIT_COMMAND.lower():
|
|
raise ValueError(
|
|
"gauss_mysql requires init_command='SET standard_conforming_strings TO off'; "
|
|
"remove the conflicting init_command from the database URL"
|
|
)
|
|
kwargs["init_command"] = _DOLPHIN_INIT_COMMAND
|
|
return args, kwargs
|
|
|
|
|
|
def _unescape_gauss_mysql_json_container(value: str) -> Any:
|
|
"""Undo exactly one complete MySQL SQL-string escape layer."""
|
|
|
|
text = value.strip()
|
|
if not text.startswith(("[", "{")) or "\\" not in text:
|
|
raise ValueError("Not an escaped JSON container")
|
|
escapes = {
|
|
"0": "\x00",
|
|
"n": "\n",
|
|
"r": "\r",
|
|
"Z": "\x1a",
|
|
"\\": "\\",
|
|
"'": "'",
|
|
'"': '"',
|
|
}
|
|
output: list[str] = []
|
|
position = 0
|
|
while position < len(text):
|
|
character = text[position]
|
|
if character == "\\":
|
|
position += 1
|
|
if position == len(text) or text[position] not in escapes:
|
|
raise ValueError("Unknown or incomplete MySQL escape")
|
|
output.append(escapes[text[position]])
|
|
elif character in {'"', "'"}:
|
|
raise ValueError("Mixed escaped and unescaped quotes")
|
|
else:
|
|
output.append(character)
|
|
position += 1
|
|
decoded = json.loads("".join(output))
|
|
if not isinstance(decoded, (dict, list)):
|
|
raise ValueError("Escaped legacy JSON must be an object or array")
|
|
return decoded
|
|
|
|
|
|
def _is_json_domain(value: Any) -> bool:
|
|
"""Recursively reject Python-only values accepted by ``literal_eval``."""
|
|
|
|
if value is None or isinstance(value, (str, bool, int)):
|
|
return True
|
|
if isinstance(value, float):
|
|
return math.isfinite(value)
|
|
if isinstance(value, list):
|
|
return all(_is_json_domain(item) for item in value)
|
|
if isinstance(value, dict):
|
|
return all(isinstance(key, str) and _is_json_domain(item) for key, item in value.items())
|
|
return False
|
|
|
|
|
|
def decode_gauss_mysql_json(value: Any) -> Any:
|
|
"""Decode strict JSON plus the two legacy Dolphin encodings we shipped."""
|
|
|
|
if isinstance(value, memoryview):
|
|
value = value.tobytes()
|
|
if isinstance(value, bytes):
|
|
value = value.decode("utf-8")
|
|
if not isinstance(value, str):
|
|
return value
|
|
try:
|
|
return json.loads(value)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
try:
|
|
decoded = ast.literal_eval(value)
|
|
except (SyntaxError, ValueError):
|
|
try:
|
|
return _unescape_gauss_mysql_json_container(value)
|
|
except (TypeError, ValueError):
|
|
preview = value[:160].replace("\n", "\\n")
|
|
raise ValueError(f"GaussDB MySQL JSON value is neither valid JSON, a supported legacy literal, nor a single SQL-escaped JSON container: {preview!r}") from None
|
|
if not _is_json_domain(decoded):
|
|
raise ValueError(f"GaussDB MySQL JSON legacy literal decoded to an unsupported type: {type(decoded).__name__}")
|
|
return decoded
|
|
|
|
|
|
def install_gauss_mysql_dialect_compatibility(driver: str = "aiomysql") -> None:
|
|
"""Patch only the selected async MySQL dialect class, once per process."""
|
|
|
|
if driver == "aiomysql":
|
|
from sqlalchemy.dialects.mysql.aiomysql import MySQLDialect_aiomysql as dialect_class
|
|
elif driver == "asyncmy":
|
|
from sqlalchemy.dialects.mysql.asyncmy import MySQLDialect_asyncmy as dialect_class
|
|
else:
|
|
raise ValueError(f"Unsupported GaussDB MySQL driver: {driver}")
|
|
|
|
if getattr(dialect_class, "_deerflow_gauss_mysql_compat_installed", False):
|
|
return
|
|
|
|
dialect_class._deerflow_gauss_mysql_original_create_connect_args = dialect_class.create_connect_args
|
|
dialect_class.create_connect_args = _gauss_mysql_create_connect_args
|
|
dialect_class._get_server_version_info = _gauss_mysql_server_version_info
|
|
dialect_class._get_default_schema_name = _gauss_mysql_default_schema_name
|
|
dialect_class.get_isolation_level = _gauss_mysql_isolation_level
|
|
dialect_class.has_table = _gauss_mysql_has_table
|
|
dialect_class._deerflow_gauss_mysql_compat_installed = True
|