data-query-agent/src/data_query_agent/odps_client.py

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Any

import pandas as pd
from odps import ODPS

from .config import Settings
from .models import QueryResult
from .sql_guard import SQLGuard, TableRef

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class TableMetadata:
    partitions: dict[str, list[str]]


class ODPSClient:
    def __init__(self, settings: Settings) -> None:
        kwargs: dict[str, Any] = {
            "access_id": settings.odps_access_id,
            "secret_access_key": settings.odps_access_key,
            "project": settings.odps_project,
            "endpoint": settings.odps_endpoint,
        }
        if settings.odps_tunnel_endpoint:
            kwargs["tunnel_endpoint"] = settings.odps_tunnel_endpoint
        self._client = ODPS(**kwargs)
        self.project = settings.odps_project
        self.timeout = settings.query_timeout_seconds
        self.max_rows = settings.query_max_rows

    async def partition_metadata(self, refs: list[TableRef]) -> TableMetadata:
        return await asyncio.to_thread(self._partition_metadata_sync, refs)

    def _partition_metadata_sync(self, refs: list[TableRef]) -> TableMetadata:
        result: dict[str, list[str]] = {}
        for ref in refs:
            full_name = f"{ref.project}.{ref.name}" if ref.project else ref.name
            table = self._client.get_table(full_name)
            result[full_name] = [column.name for column in table.schema.partitions]
        return TableMetadata(result)

    async def execute(self, sql: str) -> QueryResult:
        return await asyncio.to_thread(self._execute_sync, sql)

    def _execute_sync(self, sql: str) -> QueryResult:
        instance = self._client.run_sql(sql)
        instance_id = str(instance.id)
        logger.info("ODPS query started instance_id=%s", instance_id)
        try:
            instance.wait_for_success(timeout=self.timeout)
        except Exception:
            try:
                instance.stop()
            except Exception:
                logger.warning("Failed to stop ODPS instance %s", instance_id, exc_info=True)
            raise
        with instance.open_reader() as reader:
            frame = reader.to_pandas(start=0, count=self.max_rows + 1)
        truncated = len(frame.index) > self.max_rows
        if truncated:
            frame = frame.iloc[: self.max_rows].copy()
        if not isinstance(frame, pd.DataFrame):
            frame = pd.DataFrame(frame)
        logger.info("ODPS query finished instance_id=%s rows=%d truncated=%s", instance_id, len(frame), truncated)
        return QueryResult(frame, instance_id, truncated)


async def validate_for_odps(sql: str, guard: SQLGuard, odps: ODPSClient) -> list[TableRef]:
    refs = guard.validate(sql)
    metadata = await odps.partition_metadata(refs)
    guard.validate_partition_predicates(sql, metadata.partitions)
    return refs