import duckdb import time import logging from typing import List, Optional from .config import APP_ENV, DEV_PARQUET_URI, S3_ENDPOINT_TEMPLATE, S3_URL_STYLE from .config import ALICLOUD_ACCESS_KEY, ALICLOUD_SECRET_KEY, ALICLOUD_ENDPOINT_TEMPLATE, ALICLOUD_URL_STYLE from .config import LOG_CRED_DEBUG from .auth import get_s3_creds_by_github_token class DuckDBStore: def __init__(self): self.conn = duckdb.connect(database=":memory:", read_only=False) self.conn.execute("INSTALL httpfs;") self.conn.execute("LOAD httpfs;") self.conn.execute("INSTALL delta;") self.conn.execute("LOAD delta;") self.conn.execute("SET enable_object_cache=true;") self.conn.execute("SET enable_http_metadata_cache=true;") self.conn.execute("SET enable_external_file_cache=true;") self.conn.execute("SET parquet_metadata_cache=true;") self.conn.execute("SET preserve_insertion_order=false;") try: logger = logging.getLogger("slowquery") logger.info("event=db_init status=ok") except Exception: pass def close(self): try: self.conn.close() except Exception: pass def set_aws_s3(self, region: str): # Ensure aws extension is available try: self.conn.execute("INSTALL aws;") self.conn.execute("LOAD aws;") except Exception: pass logger = logging.getLogger("slowquery") logger.info(f"event=auth_begin region={region}") creds = get_s3_creds_by_github_token() if not creds.access_key_id or not creds.secret_access_key: raise RuntimeError("invalid aws credentials from auth") endpoint = ( S3_ENDPOINT_TEMPLATE.format(region=region) if S3_ENDPOINT_TEMPLATE else f"s3.{region}.amazonaws.com" ) url_style = S3_URL_STYLE or "virtual_host" # Prefer SECRET-based config to avoid credential_chain fallback secret_sql = ( f"CREATE OR REPLACE SECRET secret (" f"TYPE s3, PROVIDER config, KEY_ID '{creds.access_key_id}', SECRET '{creds.secret_access_key}', " f"SESSION_TOKEN '{creds.session_token or ''}', ENDPOINT '{endpoint}', REGION '{region}', URL_STYLE '{url_style}'" f");" ) self.conn.execute(secret_sql) logger.info(f"event=s3_secret_set provider=config region={region} endpoint={endpoint} url_style={url_style}") # Also set runtime parameters for completeness self.conn.execute(f"SET s3_access_key_id='{creds.access_key_id}';") self.conn.execute(f"SET s3_secret_access_key='{creds.secret_access_key}';") if creds.session_token: self.conn.execute(f"SET s3_session_token='{creds.session_token}';") self.conn.execute(f"SET s3_region='{region}';") self.conn.execute(f"SET s3_endpoint='{endpoint}';") self.conn.execute(f"SET s3_url_style='{url_style}';") logger.info(f"event=s3_params_set region={region} endpoint={endpoint} url_style={url_style}") if LOG_CRED_DEBUG: logger = logging.getLogger("slowquery") masked_key = creds.access_key_id[:4] + "****" logger.info( f"event=auth_ok key_id={masked_key} session_token={'true' if creds.session_token else 'false'} region={region} endpoint={endpoint} url_style={url_style}" ) self.s3_region = region self.s3_endpoint = endpoint self.s3_url_style = url_style def set_alicloud_s3(self, region: str): if not ALICLOUD_ACCESS_KEY or not ALICLOUD_SECRET_KEY: raise RuntimeError("missing alicloud access_key/secret_key in config.json") logger = logging.getLogger("slowquery") endpoint = ( ALICLOUD_ENDPOINT_TEMPLATE.format(region=region) if ALICLOUD_ENDPOINT_TEMPLATE else f"oss-{region}.aliyuncs.com" ) url_style = ALICLOUD_URL_STYLE or "path" secret_sql = ( f"CREATE OR REPLACE SECRET secret (" f"TYPE s3, PROVIDER config, KEY_ID '{ALICLOUD_ACCESS_KEY}', SECRET '{ALICLOUD_SECRET_KEY}', " f"ENDPOINT '{endpoint}', REGION '{region}', URL_STYLE '{url_style}'" f");" ) self.conn.execute(secret_sql) logger.info(f"event=s3_secret_set provider=config vendor=alicloud region={region} endpoint={endpoint} url_style={url_style}") self.conn.execute(f"SET s3_access_key_id='{ALICLOUD_ACCESS_KEY}';") self.conn.execute(f"SET s3_secret_access_key='{ALICLOUD_SECRET_KEY}';") self.conn.execute(f"SET s3_region='{region}';") self.conn.execute(f"SET s3_endpoint='{endpoint}';") self.conn.execute(f"SET s3_url_style='{url_style}';") logger.info(f"event=s3_params_set vendor=alicloud region={region} endpoint={endpoint} url_style={url_style}") if LOG_CRED_DEBUG: masked_key = ALICLOUD_ACCESS_KEY[:4] + "****" logger.info(f"event=auth_ok vendor=alicloud key_id={masked_key} region={region} endpoint={endpoint} url_style={url_style}") self.s3_region = region self.s3_endpoint = endpoint self.s3_url_style = url_style def create_view_delta(self, name: str, uri: str): logger = logging.getLogger("slowquery") logger.info(f"event=view_create type=delta name={name} uri={uri}") try: self.conn.execute(f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM delta_scan('{uri}');") except Exception as e: logger.exception(f"event=view_create_error type=delta name={name} uri={uri} error_class={e.__class__.__name__} error={str(e)}") raise def create_view_parquet(self, name: str, uri: str): logger = logging.getLogger("slowquery") logger.info(f"event=view_create type=parquet name={name} uri={uri}") try: self.conn.execute(f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM parquet_scan('{uri}');") except Exception as e: logger.exception(f"event=view_create_error type=parquet name={name} uri={uri} error_class={e.__class__.__name__} error={str(e)}") raise def execute_sql(self, sql: str): logger = logging.getLogger("slowquery") preview = sql.replace("\n", " ")[:200] # logger.info(f"event=sql_execute_start sql_preview={preview}") start = time.time() try: res = self.conn.execute(sql) except Exception as e: rg = getattr(self, "s3_region", "") ep = getattr(self, "s3_endpoint", "") us = getattr(self, "s3_url_style", "") logger.exception(f"event=sql_execute_error error_class={e.__class__.__name__} error={str(e)} s3_region={rg} s3_endpoint={ep} s3_url_style={us}") raise try: tbl = res.fetch_arrow_table() rows = tbl.to_pylist() columns = tbl.column_names except Exception: columns = [d[0] for d in res.description] raw = res.fetchall() rows = [dict(zip(columns, r)) for r in raw] meta = {"columns": columns, "row_count": len(rows), "elapsed_ms": int((time.time() - start) * 1000)} logger.info(f"event=sql_execute_end row_count={len(rows)} elapsed_ms={meta.get('elapsed_ms')}") return rows, meta def build_default_views(conn: DuckDBStore, vendor: str, region: Optional[str], tenant_id: Optional[str], o11y_id: Optional[str]): if APP_ENV == "dev": logger = logging.getLogger("slowquery") logger.info("event=build_views mode=dev") conn.create_view_parquet("slowlogs", DEV_PARQUET_URI) return if vendor == "aws": if not region: raise RuntimeError("missing region for aws vendor") r = region logger = logging.getLogger("slowquery") logger.info(f"event=build_views mode=prod vendor=aws region={r}") conn.set_aws_s3(r) if not tenant_id or not o11y_id: raise RuntimeError("missing tenant_id or o11y_id in prod") uri = f"s3://o11y-prod-shared-{r}/deltalake/{tenant_id}/{o11y_id}/slowlogs" logger.info(f"event=build_views_uri name=slowlogs uri={uri}") conn.create_view_delta("slowlogs", uri) return if vendor == "alicloud": if not region: raise RuntimeError("missing region for alicloud vendor") r = region logger = logging.getLogger("slowquery") logger.info(f"event=build_views mode=prod vendor=alicloud region={r}") conn.set_alicloud_s3(r) if not tenant_id or not o11y_id: raise RuntimeError("missing tenant_id or o11y_id in prod") uri = f"s3://o11y-prod-shared-{r}/deltalake/{tenant_id}/{o11y_id}/slowlogs" logger.info(f"event=build_views_uri name=slowlogs uri={uri}") conn.create_view_delta("slowlogs", uri) return raise RuntimeError("unsupported vendor for prod")