feat: improve ta_storage.bigquery

database/models/reports.py

@@ -117,7 +117,7 @@ class Upload(CodecovBaseModel, MixinBaseClass):
     upload_type_id = Column(types.Integer)
 
     @cached_property
-    def flag_names(self):
+    def flag_names(self) -> list[str]:
         return [f.flag_name for f in self.flags]
 
 

protobuf/ta_testrun.proto

@@ -11,6 +11,7 @@ message TestRun {
         PASSED = 0;
         FAILED = 1;
         SKIPPED = 2;
+        FLAKY_FAILED = 3;
     }
 
     optional Outcome outcome = 6;
@@ -27,4 +28,8 @@ message TestRun {
 
     optional string filename = 14;
     optional string framework = 15;
+
+    optional int64 upload_id = 16;
+    optional bytes flags_hash = 17;
+    optional bytes test_id = 18;
 }

requirements.in

@@ -20,6 +20,7 @@ grpcio>=1.66.2
 httpx
 jinja2>=3.1.3
 lxml>=5.3.0
+mmh3>=5.0.1
 mock
 multidict>=6.1.0
 openai

requirements.txt

@@ -208,8 +208,10 @@ markupsafe==2.1.3
     # via jinja2
 minio==7.1.13
     # via shared
-mmh3==4.0.1
-    # via shared
+mmh3==5.0.1
+    # via
+    #   -r requirements.in
+    #   shared
 mock==4.0.3
     # via -r requirements.in
 monotonic==1.5

services/bigquery.py

@@ -1,5 +1,5 @@
 from types import ModuleType
-from typing import Dict, List, cast
+from typing import Dict, List, Sequence, cast
 
 import polars as pl
 from google.api_core import retry
@@ -108,7 +108,17 @@
 
         self.write_client.batch_commit_write_streams(batch_commit_write_streams_request)
 
-    def query(self, query: str, params: dict | None = None) -> List[Dict]:
+    def query(
+        self,
+        query: str,
+        params: Sequence[
+            bigquery.ScalarQueryParameter
+            | bigquery.RangeQueryParameter
+            | bigquery.ArrayQueryParameter
+            | bigquery.StructQueryParameter
+        ]
+        | None = None,
+    ) -> List[Dict]:
         """Execute a BigQuery SQL query and return results.
         Try not to write INSERT statements and use the write method instead.
 
@@ -125,9 +135,7 @@
         job_config = bigquery.QueryJobConfig()
 
         if params:
-            job_config.query_parameters = [
-                bigquery.ScalarQueryParameter(k, "STRING", v) for k, v in params.items()
-            ]
+            job_config.query_parameters = params
 
         row_iterator = self.client.query_and_wait(
             query, job_config=job_config, retry=retry.Retry(deadline=30)

services/tests/test_bigquery.py

@@ -3,6 +3,7 @@
 
 import polars as pl
 import pytest
+from google.cloud.bigquery import ScalarQueryParameter
 
 import generated_proto.testrun.ta_testrun_pb2 as ta_testrun_pb2
 from services.bigquery import BigQueryService
@@ -63,6 +64,33 @@
     assert {row["id"] for row in results} == {1, 2}
 
 
+sql = """
+WITH sample_data AS (
+    SELECT * FROM UNNEST([
+        STRUCT(TIMESTAMP '2025-01-01T00:00:00Z' AS timestamp, 1 AS id, 'name' AS name),
+        STRUCT(TIMESTAMP '2024-12-30T00:00:00Z' AS timestamp, 2 AS id, 'name2' AS name)
+    ])
+)
+SELECT * FROM sample_data where id = @id
+"""
+
+
+@pytest.mark.skip(reason="This test requires being run using actual working creds")
+def test_bigquery_service_params():
+    bigquery_service = BigQueryService(gcp_config)
+
+    results = bigquery_service.query(
+        sql, params=[ScalarQueryParameter("id", "INT64", 2)]
+    )
+
+    assert len(results) == 1
+    assert {row["timestamp"] for row in results} == {
+        datetime.fromisoformat("2024-12-30T00:00:00Z"),
+    }
+    assert {row["name"] for row in results} == {"name2"}
+    assert {row["id"] for row in results} == {2}
+
+
 @pytest.mark.skip(reason="This test requires being run using actual working creds")
 def test_bigquery_service_polars():
     bigquery_service = BigQueryService(gcp_config)

ta_storage/base.py

@@ -18,5 +18,6 @@ def write_testruns(
         upload: Upload,
         framework: str | None,
         testruns: list[test_results_parser.Testrun],
+        flaky_test_set: set[str],
     ):
         pass