feat: capture BigQuery query stats for mktplace metering

mindsdb/api/http/namespaces/sql.py

@@ -28,6 +28,7 @@
 from mindsdb.utilities.exception import QueryError
 from mindsdb.utilities.functions import mark_process
 from mindsdb.interfaces.agents.chart_agent import ChartAgent
+from mindsdb.integrations.handlers.bigquery_handler import query_stats_registry
 
 logger = log.getLogger(__name__)
 
@@ -388,6 +389,21 @@ def find_constants_f(node, is_table, is_target, callstack, **kwargs):
         return response, 200
 
 
+@ns_conf.route("/query_stats/<string:query_id>")
+@ns_conf.param("query_id", "Correlation id supplied by the caller when executing the query")
+class QueryStats(Resource):
+    @ns_conf.doc("query_stats")
+    @api_endpoint_metrics("GET", "/sql/query_stats")
+    def get(self, query_id):
+        """Return and remove BigQuery execution stats for the given correlation id.
+
+        Returns a JSON object with total_bytes_billed, cache_hit, and project_id,
+        or an empty object if the id is not found.
+        """
+        stats = query_stats_registry.pop(query_id)
+        return stats, 200
+
+
 @ns_conf.route("/list_databases")
 @ns_conf.param("list_databases", "lists databases of mindsdb")
 class ListDatabases(Resource):

mindsdb/integrations/handlers/bigquery_handler/bigquery_handler.py

@@ -6,7 +6,9 @@
 import pandas as pd
 from sqlalchemy_bigquery.base import BigQueryDialect
 
+from mindsdb.integrations.handlers.bigquery_handler import query_stats_registry
 from mindsdb.utilities import log
+from mindsdb.utilities.context import context as ctx
 from mindsdb_sql_parser.ast.base import ASTNode
 from mindsdb.integrations.libs.base import MetaDatabaseHandler
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
@@ -287,8 +289,9 @@ def native_query(self, query: str) -> Response:
             job_config = QueryJobConfig(
                 default_dataset=f"{self.connection_data['project_id']}.{self.connection_data['dataset']}"
             )
-            query = connection.query(query, job_config=job_config)
-            result = query.to_dataframe()
+            query_job = connection.query(query, job_config=job_config)
+            result = query_job.to_dataframe()
+            self._record_query_stats(query_job)
             if not result.empty:
                 response = Response(RESPONSE_TYPE.TABLE, result)
             else:
@@ -298,6 +301,21 @@ def native_query(self, query: str) -> Response:
             response = Response(RESPONSE_TYPE.ERROR, error_message=str(e))
         return response
 
+    def _record_query_stats(self, query_job) -> None:
+        """Capture BigQuery QueryJob stats into the in-process registry if a correlation id is present."""
+        try:
+            query_id = ctx.params.get("correlation_id") if isinstance(ctx.params, dict) else None
+            if not query_id:
+                return
+            query_stats_registry.accumulate(
+                query_id,
+                bytes_billed=int(query_job.total_bytes_billed or 0),
+                cache_hit=bool(query_job.cache_hit),
+                project_id=self.connection_data.get("project_id", ""),
+            )
+        except Exception:  # noqa: BLE001
+            pass
+
     def query(self, query: ASTNode) -> Response:
         """
         Executes a SQL query represented by an ASTNode and retrieves the data.

mindsdb/integrations/handlers/bigquery_handler/query_stats_registry.py

@@ -0,0 +1,62 @@
+"""Thread-safe in-process registry for BigQuery query execution stats.
+
+Stats are stored keyed by a caller-supplied query_id, accumulated across
+multiple handler invocations for the same logical query (e.g. JOINs that
+hit the handler once per BQ table), then popped by the caller after the
+query returns.
+"""
+import threading
+import time
+
+_lock = threading.Lock()
+_registry: dict[str, dict] = {}
+_MAX_ENTRIES = 10_000
+_TTL_SECONDS = 300.0
+_EVICT_INTERVAL_SECONDS = 60.0
+_last_evict = 0.0
+
+
+def accumulate(query_id: str, bytes_billed: int, cache_hit: bool, project_id: str) -> None:
+    """Accumulate BigQuery stats for query_id.
+
+    Sums bytes_billed across multiple invocations (JOIN across BQ tables).
+    cache_hit stays True only when ALL sub-queries were cache hits.
+    """
+    now = time.monotonic()
+    with _lock:
+        _evict(now)
+        if len(_registry) >= _MAX_ENTRIES:
+            return
+        if query_id in _registry:
+            _registry[query_id]["total_bytes_billed"] += bytes_billed
+            _registry[query_id]["cache_hit"] = _registry[query_id]["cache_hit"] and cache_hit
+        else:
+            _registry[query_id] = {
+                "total_bytes_billed": bytes_billed,
+                "cache_hit": cache_hit,
+                "project_id": project_id,
+                "_ts": now,
+            }
+
+
+def pop(query_id: str) -> dict:
+    """Pop and return stats for query_id, or empty dict if not found."""
+    with _lock:
+        entry = _registry.pop(query_id, {})
+    entry.pop("_ts", None)
+    return entry
+
+
+def _evict(now: float) -> None:
+    """Remove TTL-expired entries. Must be called with _lock held.
+
+    Throttled to scan at most once per _EVICT_INTERVAL_SECONDS so that the
+    O(n) sweep does not run on every accumulate() call.
+    """
+    global _last_evict
+    if now - _last_evict < _EVICT_INTERVAL_SECONDS:
+        return
+    _last_evict = now
+    expired = [k for k, v in _registry.items() if now - v["_ts"] > _TTL_SECONDS]
+    for k in expired:
+        del _registry[k]

tests/unit/handlers/test_bigquery_query_stats_registry.py

@@ -0,0 +1,55 @@
+import pytest
+
+try:
+    from mindsdb.integrations.handlers.bigquery_handler import query_stats_registry as reg
+except ImportError:
+    pytestmark = pytest.mark.skip("Bigquery handler not installed")
+
+
+@pytest.fixture(autouse=True)
+def _clear_registry():
+    reg._registry.clear()
+    reg._last_evict = 0.0
+    yield
+    reg._registry.clear()
+    reg._last_evict = 0.0
+
+
+def test_accumulate_then_pop_returns_stats_without_internal_ts():
+    reg.accumulate("q1", bytes_billed=1024, cache_hit=False, project_id="p")
+
+    stats = reg.pop("q1")
+
+    assert stats == {"total_bytes_billed": 1024, "cache_hit": False, "project_id": "p"}
+    # The id is consumed on pop.
+    assert reg.pop("q1") == {}
+
+
+def test_accumulate_sums_bytes_and_ands_cache_hit_across_calls():
+    # Mimics a JOIN that hits the handler once per BigQuery table.
+    reg.accumulate("q1", bytes_billed=1000, cache_hit=True, project_id="p")
+    reg.accumulate("q1", bytes_billed=500, cache_hit=False, project_id="p")
+
+    stats = reg.pop("q1")
+
+    assert stats["total_bytes_billed"] == 1500
+    # cache_hit only stays True when ALL sub-queries were cache hits.
+    assert stats["cache_hit"] is False
+
+
+def test_pop_unknown_id_returns_empty_dict():
+    assert reg.pop("nope") == {}
+
+
+def test_evict_removes_expired_entries(monkeypatch):
+    clock = {"now": 1000.0}
+    monkeypatch.setattr(reg.time, "monotonic", lambda: clock["now"])
+
+    reg.accumulate("old", bytes_billed=1, cache_hit=False, project_id="p")
+
+    # Advance past TTL + eviction interval, then trigger another accumulate.
+    clock["now"] = 1000.0 + reg._TTL_SECONDS + reg._EVICT_INTERVAL_SECONDS + 1
+    reg.accumulate("new", bytes_billed=2, cache_hit=False, project_id="p")
+
+    assert reg.pop("old") == {}
+    assert reg.pop("new")["total_bytes_billed"] == 2