Add BigQuery support to evaluation config generation

plugin/skills/context-engineering-init/references/bigquery.md

@@ -0,0 +1,36 @@
+## BigQuery
+
+**Required Information:**
+- Data Source Name (e.g., `my-bigquery-db`)
+- Google Cloud Project ID
+- Dataset ID
+
+**Template:**
+
+```yaml
+kind: source
+name: <data_source_name>
+type: bigquery
+project: <project_id>
+dataset: <dataset_id>
+---
+kind: tool
+name: <data_source_name>-list-schemas
+type: bigquery-sql
+source: <data_source_name>
+description: |
+  Use this tool to list tables and their schemas in the <data_source_name> dataset.
+
+  Progressive Schema Discovery (Recommended):
+  1) Fetch structure first,
+  2) Go deep on specific parts if interested,
+  3) Use batching if info is too large.
+statement: |
+  SELECT table_name, column_name, data_type, description FROM `<project_id>`.`<dataset_id>`.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS ORDER BY table_name, column_name
+---
+kind: tool
+name: <data_source_name>-execute-sql
+type: bigquery-execute-sql
+source: <data_source_name>
+description: Use this tool to execute SQL statements against the <data_source_name> BigQuery dataset.
+```

src/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery.py

@@ -0,0 +1,80 @@
+from typing import Any
+
+import google.cloud.geminidataanalytics_v1beta as gda
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class BigQueryConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit BigQuery configuration
+    topologies utilized by both EvalBench binaries and GDA Context objects.
+    """
+
+    SOURCE_TYPE = "bigquery"
+    DIALECT = "googlesql"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "dataset",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.dataset = params.get("dataset")
+        self.location = params.get("location")
+        # Optional explicit table scoping; the public GDA proto references
+        # BigQuery at table granularity rather than dataset granularity.
+        self.tables = params.get("tables") or []
+
+    def generate_db_config(self) -> str:
+        db_type = "bigquery"
+        db_path = f"projects/{self.project}/datasets/{self.dataset}"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.dataset,
+            "database_path": db_path,
+            "gcp_project_id": self.project,
+            "max_executions_per_minute": 100,
+        }
+        if self.location:
+            db_config["location"] = self.location
+        return yaml.safe_dump(
+            db_config, sort_keys=False, default_flow_style=False
+        ).strip()
+
+    def build_datasource_reference(
+        self, context_set_id: str
+    ) -> gda.DatasourceReferences:
+        datasource_ref = gda.DatasourceReferences()
+
+        table_references = [
+            gda.BigQueryTableReference(
+                project_id=self.project,
+                dataset_id=self.dataset,
+                table_id=table_id,
+            )
+            for table_id in self.tables
+        ]
+
+        bq_references = gda.BigQueryTableReferences(
+            table_references=table_references
+        )
+
+        # The agent_context_reference field on BigQueryTableReferences is not
+        # yet available in the public google-cloud-geminidataanalytics SDK
+        # (restricted-visibility rollout; see the Wave2 FR). Attach it
+        # dynamically so this generator works with internal SDK builds and
+        # degrades gracefully on public ones.
+        if "agent_context_reference" in {
+            f.name for f in type(bq_references).pb(bq_references).DESCRIPTOR.fields
+        }:
+            bq_references.agent_context_reference = gda.AgentContextReference(
+                context_set_id=context_set_id
+            )
+
+        datasource_ref.bq = bq_references
+        return datasource_ref

src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py

@@ -10,6 +10,7 @@
 
 from .db_generators.alloydb import AlloyDBConfigGenerator
 from .db_generators.base import BaseDBConfigGenerator
+from .db_generators.bigquery import BigQueryConfigGenerator
 from .db_generators.mysql import MySQLConfigGenerator
 from .db_generators.postgres import PostgresConfigGenerator
 from .db_generators.spanner import SpannerConfigGenerator
@@ -130,6 +131,7 @@ def _get_db_generator(params: dict[str, Any]) -> BaseDBConfigGenerator:
 
     generators = {
         AlloyDBConfigGenerator.SOURCE_TYPE: AlloyDBConfigGenerator,
+        BigQueryConfigGenerator.SOURCE_TYPE: BigQueryConfigGenerator,
         PostgresConfigGenerator.SOURCE_TYPE: PostgresConfigGenerator,
         MySQLConfigGenerator.SOURCE_TYPE: MySQLConfigGenerator,
         SpannerConfigGenerator.SOURCE_TYPE: SpannerConfigGenerator,

tests/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery_test.py

@@ -0,0 +1,81 @@
+import pytest
+import yaml
+
+from google.cloud.db_context_enrichment.evaluate.db_generators.bigquery import (
+    BigQueryConfigGenerator,
+)
+
+
+@pytest.fixture
+def mock_params():
+    return {
+        "project": "test-project",
+        "dataset": "test-dataset",
+    }
+
+
+def test_generate_db_config(mock_params):
+    gen = BigQueryConfigGenerator(mock_params)
+    db_config_yaml = gen.generate_db_config()
+
+    assert gen.DIALECT == "googlesql"
+
+    config = yaml.safe_load(db_config_yaml)
+    assert config == {
+        "db_type": "bigquery",
+        "dialect": "googlesql",
+        "database_name": "test-dataset",
+        "database_path": "projects/test-project/datasets/test-dataset",
+        "gcp_project_id": "test-project",
+        "max_executions_per_minute": 100,
+    }
+
+
+def test_generate_db_config_with_location(mock_params):
+    gen = BigQueryConfigGenerator({**mock_params, "location": "US"})
+    config = yaml.safe_load(gen.generate_db_config())
+    assert config["location"] == "US"
+
+
+def test_missing_required_fields():
+    with pytest.raises(ValueError, match="dataset"):
+        BigQueryConfigGenerator({"project": "test-project"})
+
+
+def test_generate_model_config(mock_params):
+    gen = BigQueryConfigGenerator(mock_params)
+    model_config_yaml = gen.generate_model_config(
+        "projects/test-project/locations/us-west1/contextSets/my-context"
+    )
+    m_config = yaml.safe_load(model_config_yaml)
+
+    assert m_config["generator"] == "query_data_api"
+    assert m_config["project_id"] == "test-project"
+    assert m_config["location"] == "global"
+    # The public GDA SDK references BigQuery at table granularity; with no
+    # explicit tables configured the reference set is empty but present.
+    assert "bq" in m_config["context"]["datasource_references"]
+
+
+def test_generate_model_config_with_tables(mock_params):
+    gen = BigQueryConfigGenerator({**mock_params, "tables": ["t1", "t2"]})
+    model_config_yaml = gen.generate_model_config(
+        "projects/test-project/locations/us-west1/contextSets/my-context"
+    )
+    m_config = yaml.safe_load(model_config_yaml)
+
+    table_refs = m_config["context"]["datasource_references"]["bq"][
+        "table_references"
+    ]
+    assert table_refs == [
+        {
+            "project_id": "test-project",
+            "dataset_id": "test-dataset",
+            "table_id": "t1",
+        },
+        {
+            "project_id": "test-project",
+            "dataset_id": "test-dataset",
+            "table_id": "t2",
+        },
+    ]