import abc from typing import List, Union from sqlalchemy import inspect, text import pandas as pd from errors.errors import BadRequestException, PermissionDeniedException from svc.crossregion import CrossRegionClient from utils.database import create_engine_with_timeout from utils.credential import user_has_permission from config import config from svc.measuring import measuring_network from utils.logging import logging def get_sql_executor(data_context: dict, app_id: Union[int, None] = None, public_key: Union[str, None] = None, user_token: Union[str, None] = None, trace_id: str = "") -> "DBExecutorABC": if data_context.get("database_uri"): return SQLAlchemyExecutor(data_context["database_uri"]) else: return DataAPIExecutor(data_context, app_id, public_key, user_token, trace_id) class DBExecutorABC(abc.ABC): @abc.abstractmethod def execute_sql(self, sql: str, limit=0) -> dict: pass @abc.abstractmethod def explain_sql(self, sql: str): pass @abc.abstractmethod def sql_to_dataframe(self, sql, chunksize=None): pass @abc.abstractmethod def get_database_name(self): pass @abc.abstractmethod def get_table_names(self) -> List[str]: pass @abc.abstractmethod def get_columns(self, table_name) -> List: pass @abc.abstractmethod def get_foreign_keys(self, table_name): pass @abc.abstractmethod def get_primary_key(self, table_name): pass class SQLAlchemyExecutor(DBExecutorABC): def __init__(self, database_uri: str) -> None: self.database_uri = database_uri self.engine = create_engine_with_timeout(database_uri) self.inspector = inspect(self.engine) def get_database_name(self): dbname = self.engine.url.database if dbname and "sqlite" in dbname: return dbname.split("/")[-2] else: return dbname def execute_sql(self, sql: str, limit=0) -> dict: raise NotImplementedError("sqlalchemy executor not support execute_sql") def explain_sql(self, sql: str): return self.engine.execute(text(f"EXPLAIN {sql}")).fetchall() def sql_to_dataframe(self, sql, limit=None) -> pd.DataFrame: # limit is not used in sqlalchemy executor df = pd.read_sql_query(text(sql), self.engine.connect()) return df def get_table_names(self) -> List[str]: # ['car_makers', 'car_names', 'cars_data', 'continents', 'countries', 'model_list'] return self.inspector.get_table_names() + self.inspector.get_view_names() def get_columns(self, table_name) -> List: # [{'name': 'Id', 'type': INTEGER(display_width=11), 'default': None, 'comment': None, 'nullable': False, 'autoincrement': True}, {'name': 'Maker', 'type': TEXT(collation='utf8mb4_general_ci'), 'default': None, 'comment': None, 'nullable': True}, {'name': 'FullName', 'type': TEXT(collation='utf8mb4_general_ci'), 'default': None, 'comment': None, 'nullable': True}, {'name': 'country', 'type': INTEGER(display_width=11), 'default': None, 'comment': None, 'nullable': True, 'autoincrement': False}] columns = self.inspector.get_columns(table_name) return columns def get_foreign_keys(self, table_name): # [{'name': 'fk_country_id', 'constrained_columns': ['Country'], 'referred_schema': 'new_car_1', 'referred_table': 'countries', 'referred_columns': ['CountryId'], 'options': {}}] foreign_keys = self.inspector.get_foreign_keys(table_name) return foreign_keys def get_primary_key(self, table_name): # {'constrained_columns': ['Id'], 'name': None} primary_keys = self.inspector.get_pk_constraint(table_name) return primary_keys class DataAPIExecutor(DBExecutorABC): def __init__(self, data_context: dict, app_id: Union[int, None] = None, public_key: Union[str, None] = None, user_token: Union[str, None] = None, trace_id: str = "") -> None: self.data_context = data_context self.database = data_context["database"] self.org_id = data_context["org_id"] self.app_id = app_id self.app_key = public_key self.user_token = user_token self.trace_id = trace_id self.dataapi_client = CrossRegionDataAPI(self.app_key, self.app_id, self.user_token, trace_id) def get_database_name(self): return self.database def measure_network(self, resp_json): if self.user_token: logging.info("use user_token, skip measuring_network") return if not (self.app_id and self.app_key): logging.info("no app_id or app_key, skip measuring_network") return bytes_count = resp_json["bytes_count"] try: measuring_network(self.org_id, self.app_id, self.app_key, bytes_count) except Exception as e: logging.error("failed to measuring_network: %s", e) def execute_sql(self, sql: str, limit=0) -> dict: resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql, limit_rows=limit) resp_json = resp.json() self.measure_network(resp_json) resp_json = resp_json["data"] df_dict = {} rows = resp_json["rows"] for column in resp_json["columns"]: col = column["col"] nullable = column["nullable"] data = [] for row in rows: item = row[col] if nullable: item = item or None data.append(item) df_dict[col] = data return df_dict def explain_sql(self, sql: str): resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, f"EXPLAIN {sql}") resp_json = resp.json() if resp.status_code != 200: raise BadRequestException(resp_json["result"]["message"]) if resp_json["data"]["result"]["code"] not in (0, 200): raise BadRequestException(resp_json["data"]["result"]["message"]) return resp_json["data"]["rows"] def sql_to_dataframe(self, sql: str, limit=0) -> pd.DataFrame: df_dict = self.execute_sql(sql, limit) return pd.DataFrame(data=df_dict) def count_tables(self) -> int: sql = f"SELECT COUNT(*) as count FROM information_schema.tables WHERE table_schema = '{self.database}'" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql) resp_json = resp.json() if resp.status_code != 200: raise BadRequestException(resp_json["result"]["message"]) return int(resp_json["data"]["rows"][0]["count"]) def count_columns(self, table_name) -> int: sql = f"SELECT COUNT(*) as count FROM information_schema.columns WHERE table_schema = '{self.database}' and table_name = '{table_name}'" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql) resp_json = resp.json() if resp.status_code != 200: raise BadRequestException(resp_json["result"]["message"]) return int(resp_json["data"]["rows"][0]["count"]) def get_table_names(self) -> List[str]: try: table_count = self.count_tables() logging.info("table_count: %s", table_count) except Exception as e: table_count = 200 logging.info("failed to count_tables: %s, use default value %s", e, table_count) sql = f"SHOW FULL TABLES FROM `{self.database}`" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql, limit_rows=table_count) resp_json = resp.json()["data"] if resp.status_code != 200: raise BadRequestException(resp_json["result"]["message"]) table_names = [] for i in resp_json["rows"]: for k, v in i.items(): if k.startswith("Tables_in_"): table_names.append(v) return table_names def get_columns(self, table_name) -> List: try: column_count = self.count_columns(table_name) logging.info("column count of table %s: %s", table_name, column_count) except Exception as e: column_count = 1000 logging.info("failed to count columns of table %s: %s, use default value %s", table_name, e, column_count) sql = f"SHOW COLUMNS FROM `{table_name}`;" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql, limit_rows=column_count) resp_json = resp.json()["data"] column_names = [] for i in resp_json["rows"]: column = { 'name': i["Field"], 'type': i["Type"], 'default': i.get("Default") or None, # 'comment': None, 'nullable': i.get("Null") == "YES", # 'autoincrement': True, } column_names.append(column) return column_names def get_foreign_keys(self, table_name): # [{'name': 'fk_country_id', 'constrained_columns': ['Country'], 'referred_schema': 'new_car_1', 'referred_table': 'countries', 'referred_columns': ['CountryId'], 'options': {}}] sql = f"SELECT * FROM information_schema.key_column_usage WHERE referenced_table_name IS NOT NULL and TABLE_SCHEMA='{self.database}' and table_name='{table_name}'" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql) resp_json = resp.json()["data"] foreign_keys = [] for i in resp_json["rows"] or []: foreign_key = { 'name': i['CONSTRAINT_NAME'], 'constrained_columns': [i['COLUMN_NAME']], 'referred_schema': i['REFERENCED_TABLE_SCHEMA'], 'referred_table': i['REFERENCED_TABLE_NAME'], 'referred_columns': [i['REFERENCED_COLUMN_NAME']], 'options': {}, } foreign_keys.append(foreign_key) return foreign_keys def get_primary_key(self, table_name): sql = f"SHOW KEYS FROM `{table_name}` WHERE Key_name = 'PRIMARY'" resp = self.dataapi_client.exec_sql_in_o11y_with_data_context(self.data_context, sql) resp_json = resp.json()["data"] constrained_columns = [] for i in resp_json["rows"]: constrained_columns.append(i["Column_name"]) primary_key = {'constrained_columns': ['Id'], 'name': None} return primary_key def get_provider_for_regional_host(cloud_provider_name: str) -> str: """ Get provider for regional host prefix. Only AliCloud shared/serverless uses a provider-prefixed regional host. For AWS and others, returns empty string to use the legacy cross-region domain. Args: cloud_provider_name: Cloud provider name from cluster info (e.g., 'alibaba', 'aliyun', 'alicloud', 'aws') Returns: Provider string for regional host ('alicloud' for AliCloud, '' for others) """ from svc.crossregion import CLOUD_PROVIDER_ALICLOUD provider = (cloud_provider_name or "").strip().lower() # Map various alibaba cloud naming conventions to 'alicloud' if provider in ["alibaba", "aliyun", "alicloud", "alibabacloud"]: return CLOUD_PROVIDER_ALICLOUD return "" class CrossRegionDataAPI: def __init__(self, api_key: Union[str, None], app_id: Union[int, None], user_token: Union[str, None] = None, trace_id: str = "") -> None: self.api_key = api_key self.app_id = app_id self.user_token = user_token self.trace_id = trace_id self.cross_region_client = CrossRegionClient(config.cross_region_env, config.cross_region_auth, trace_id) def check_app_id_cluster_permission(self, region, app_id, cluster_id, provider: str = ""): """ Check if app_id has permission to access cluster. Note: AliCloud does not support app_id_cluster_permission_check endpoint. """ from svc.crossregion import CLOUD_PROVIDER_ALICLOUD # Skip permission check for AliCloud as it doesn't have this endpoint if provider == CLOUD_PROVIDER_ALICLOUD: logging.info("skip app_id_cluster_permission_check for AliCloud provider") return url = "/inner/dataservice/app_id_cluster_permission_check" payload = { "app_id": int(app_id), "cluster_id": int(cluster_id), } resp = self.cross_region_client.request_o11y(region, "post", url, json=payload) resp_json = resp.json() if resp.status_code != 200: raise BadRequestException(resp_json["message"]) if not resp_json["pass"]: raise PermissionDeniedException("no permission to access this cluster") def exec_sql_in_o11y(self, region, project_id, cluster_id, database, sql, session_id, limit_rows=0, timeout_ms=0, provider: str = ""): """ Execute SQL in o11y regional endpoint. Args: region: The region to send the request to project_id: Project ID cluster_id: Cluster ID database: Database name sql: SQL to execute session_id: Session ID limit_rows: Limit rows (default 0) timeout_ms: Timeout in milliseconds (default 0) provider: Cloud provider (e.g., 'alicloud', ''). When set to 'alicloud', uses AliCloud domain format and adds /dataservice path prefix. """ logging.info( "exec_sql_in_o11y, region: %s, project_id: %s, cluster_id: %s, database: %s, sql: %s, session_id: %s, limit_rows: %s, timeout_ms: %s, user_token %s, app_id: %s, app_key: %s, provider: %s", region, project_id, cluster_id, database, sql, session_id, limit_rows, timeout_ms, self.user_token, self.app_id, self.api_key, provider, ) # before exec sql, check permission if self.user_token: pass # had already check user token permission in outside else: self.check_app_id_cluster_permission(region, self.app_id, cluster_id, provider) url = "/inner/dataservice/exec/sql" credential = { "type": "APIKEY", "key": self.api_key, } if self.user_token: credential = { "type": "AUTH0", "key": self.user_token, } payload = { "project_id": project_id, "cluster_id": int(cluster_id), "database": database, "sql": sql, "session_id": session_id, "limit_rows": limit_rows, "timeout_ms": timeout_ms, "return_all_stmt_result": True, "credential": credential, } # Use provider-aware request method if provider is specified if provider: return self.cross_region_client.request_o11y_with_provider(region, provider, "post", url, json=payload) return self.cross_region_client.request_o11y(region, "post", url, json=payload) def exec_sql_in_o11y_with_data_context(self, data_context: dict, sql: str, session_id=None, limit_rows=0): """ Execute SQL in o11y regional endpoint using data context. Automatically detects cloud provider from data_context and uses appropriate endpoint. Args: data_context: Data context dict containing region, project_id, cluster_id, database, and optionally cloud_provider sql: SQL to execute session_id: Session ID (optional) limit_rows: Limit rows (default 0) """ session_id = session_id or "" # TODO # Get provider from data_context if available provider = get_provider_for_regional_host(data_context.get("cloud_provider", "")) resp = self.exec_sql_in_o11y( data_context["region"], data_context["project_id"], data_context["cluster_id"], data_context["database"], sql, session_id, limit_rows=limit_rows, timeout_ms=60000, provider=provider, ) resp_json = resp.json() if resp.status_code != 200: # {"data":{"columns":null,"rows":null,"result":{"code":49900001,"message":"Failed to get cluster information, Please check if the cluster is available.","start_ms":1700022656059,"end_ms":1700022656069,"latency":"9.281083ms","row_count":0,"row_affect":0,"limit":0}},"bytes_count":0} raise BadRequestException(resp_json["data"]["result"]["message"]) return resp