feat: Query API for SQL over extracted entities (sync spec v0.7.15)

cloudglue/client/main.py

@@ -22,6 +22,7 @@
 from cloudglue.sdk.api.share_api import ShareApi
 from cloudglue.sdk.api.data_connectors_api import DataConnectorsApi
 from cloudglue.sdk.api.deep_search_api import DeepSearchApi
+from cloudglue.sdk.api.query_api import QueryApi
 from cloudglue.sdk.configuration import Configuration
 from cloudglue.sdk.api_client import ApiClient
 
@@ -46,6 +47,7 @@
     Share,
     DataConnectors,
     DeepSearch,
+    Query,
 )
 from cloudglue._version import __version__
 
@@ -99,6 +101,7 @@ def __init__(
         self.share_api = ShareApi(self.api_client)
         self.data_connectors_api = DataConnectorsApi(self.api_client)
         self.deep_search_api = DeepSearchApi(self.api_client)
+        self.query_api = QueryApi(self.api_client)
 
         # Set up resources with their respective API clients
         self.chat = Chat(self.chat_api)
@@ -120,6 +123,7 @@ def __init__(
         self.share = Share(self.share_api)
         self.data_connectors = DataConnectors(self.data_connectors_api)
         self.deep_search = DeepSearch(self.deep_search_api)
+        self.query = Query(self.query_api)
 
     def close(self):
         """Close the API client."""

cloudglue/client/resources/__init__.py

@@ -21,8 +21,10 @@
 from cloudglue.client.resources.share import Share
 from cloudglue.client.resources.data_connectors import DataConnectors
 from cloudglue.client.resources.deep_search import DeepSearch
+from cloudglue.client.resources.query import Query
 
 __all__ = [
+    "Query",
     "CloudglueError",
     "Chat",
     "Completions",

cloudglue/client/resources/query.py

@@ -0,0 +1,236 @@
+# cloudglue/client/resources/query.py
+"""Query resource for Cloudglue API."""
+import time
+from typing import List, Optional
+
+from cloudglue.sdk.models.run_query_request import RunQueryRequest
+from cloudglue.sdk.rest import ApiException
+
+from cloudglue.client.resources.base import CloudglueError
+
+
+class Query:
+    """Client for the Cloudglue Query API.
+
+    Runs read-only SQL (or natural-language) queries over the structured data
+    extracted from collections, against three virtual tables — files,
+    entities, and segment_entities — built from each file's most recent
+    completed extraction.
+    """
+
+    def __init__(self, api):
+        """Initialize the Query client.
+
+        Args:
+            api: The QueryApi instance.
+        """
+        self.api = api
+
+    def run(
+        self,
+        collections: List[str],
+        sql: Optional[str] = None,
+        query: Optional[str] = None,
+        format: Optional[str] = None,
+        max_rows: Optional[int] = None,
+        background: Optional[bool] = None,
+        dry_run: Optional[bool] = None,
+    ):
+        """Run a read-only SQL query over one or more collections.
+
+        Provide exactly one of `sql` (2 credits) or `query` — a
+        natural-language question that Cloudglue compiles to SQL against the
+        same virtual schema (4 credits; the compiled statement is returned in
+        the result's `sql` field). Results are returned inline and stored, so
+        completed runs can be re-fetched via `get()`.
+
+        Args:
+            collections: Collection IDs (1-20) whose extracted data to query.
+            sql: A single read-only SELECT statement over the virtual tables.
+                Exactly one of sql or query must be provided.
+            query: A natural-language question to run instead of sql.
+            format: Output format ('json', 'csv', 'jsonl') for exports.
+            max_rows: Maximum number of rows to return (1-10000).
+            background: Run as a background export; poll with `get()` (or
+                `wait_for_ready()`) and download via the result's
+                download_url.
+            dry_run: Validate (and for natural-language queries, compile)
+                without executing.
+
+        Returns:
+            QueryResult object (inline rows, or export state when background).
+
+        Raises:
+            CloudglueError: On invalid/rejected SQL (400), insufficient
+                credits (402), unknown collections (404), execution timeout
+                (408), oversized datasets (409), uncompilable natural-language
+                queries (422), or rate limits (429).
+        """
+        try:
+            request = RunQueryRequest(
+                collections=collections,
+                sql=sql,
+                query=query,
+                format=format,
+                max_rows=max_rows,
+                background=background,
+                dry_run=dry_run,
+            )
+            return self.api.run_query(run_query_request=request)
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except Exception as e:
+            raise CloudglueError(str(e))
+
+    def get_schema(self, collections: List[str]):
+        """Introspect the virtual tables and per-collection extracted fields.
+
+        Returns column names, entity field names, types, and levels, plus each
+        collection's verbatim extract schema and prompt — use before writing a
+        query.
+
+        Args:
+            collections: Collection IDs (1-20) to introspect.
+
+        Returns:
+            QuerySchema object.
+
+        Raises:
+            CloudglueError: If there is an error introspecting the schema.
+        """
+        try:
+            return self.api.get_query_schema(collections=",".join(collections))
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except Exception as e:
+            raise CloudglueError(str(e))
+
+    def list(
+        self,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        status: Optional[str] = None,
+        created_before: Optional[str] = None,
+        created_after: Optional[str] = None,
+    ):
+        """List query runs with pagination and filtering.
+
+        List items omit the columns and rows payloads — fetch an individual
+        run via `get()` for the full result.
+
+        Args:
+            limit: Maximum number of runs to return (1-100).
+            offset: Number of runs to skip.
+            status: Filter by status ('completed', 'failed', 'in_progress',
+                'cancelled').
+            created_before: Filter runs created before a date (YYYY-MM-DD, UTC).
+            created_after: Filter runs created after a date (YYYY-MM-DD, UTC).
+
+        Returns:
+            QueryListResponse object.
+
+        Raises:
+            CloudglueError: If there is an error listing query runs.
+        """
+        try:
+            return self.api.list_queries(
+                limit=limit,
+                offset=offset,
+                status=status,
+                created_before=created_before,
+                created_after=created_after,
+            )
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except Exception as e:
+            raise CloudglueError(str(e))
+
+    def get(self, query_id: str):
+        """Retrieve a stored query run by ID, including its result rows.
+
+        Results larger than the inline storage cap are replayed truncated
+        (`truncated: true`).
+
+        Args:
+            query_id: The ID of the query run.
+
+        Returns:
+            QueryResult object.
+
+        Raises:
+            CloudglueError: If there is an error retrieving the run.
+        """
+        try:
+            return self.api.get_query(id=query_id)
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except Exception as e:
+            raise CloudglueError(str(e))
+
+    def cancel(self, query_id: str):
+        """Cancel an in-progress background export.
+
+        The run is marked cancelled synchronously, the export stream is
+        aborted mid-flight (the partial upload is discarded), and reserved
+        credits are refunded. A run that has already completed or failed is
+        returned unchanged.
+
+        Args:
+            query_id: The ID of the query run to cancel.
+
+        Returns:
+            QueryResult object.
+
+        Raises:
+            CloudglueError: If there is an error cancelling the run.
+        """
+        try:
+            return self.api.cancel_query_export(id=query_id)
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except Exception as e:
+            raise CloudglueError(str(e))
+
+    def wait_for_ready(
+        self,
+        query_id: str,
+        poll_interval: int = 5,
+        timeout: int = 180,
+    ):
+        """Wait for a background query run to reach a terminal state.
+
+        Polls `get()` until the run is completed, failed, or cancelled, and
+        returns the final QueryResult either way — a failed run is returned
+        with its `error` payload (code/message) intact so callers can inspect
+        why it failed, matching the other wait helpers in this client.
+
+        Args:
+            query_id: The ID of the query run to wait for.
+            poll_interval: How often to check the run status (in seconds).
+            timeout: Maximum time to wait (in seconds).
+
+        Returns:
+            The final QueryResult object (status 'completed', 'failed', or
+            'cancelled'; on failure, `error` carries the details).
+
+        Raises:
+            CloudglueError: If polling itself errors or the timeout is
+                reached.
+        """
+        try:
+            elapsed = 0
+            while elapsed < timeout:
+                run = self.get(query_id)
+                if run.status in ("completed", "failed", "cancelled"):
+                    return run
+                time.sleep(poll_interval)
+                elapsed += poll_interval
+            raise TimeoutError(
+                f"Query run did not complete within {timeout} seconds"
+            )
+        except ApiException as e:
+            raise CloudglueError(str(e), e.status, e.data, e.headers, e.reason)
+        except CloudglueError:
+            raise
+        except Exception as e:
+            raise CloudglueError(str(e))