support query()

Author: Linchin

google/cloud/bigquery/_job_helpers.py

@@ -49,6 +49,7 @@
 import google.api_core.exceptions as core_exceptions
 from google.api_core import retry as retries
 
+from google.cloud.bigquery import enums
 from google.cloud.bigquery import job
 import google.cloud.bigquery.job.query
 import google.cloud.bigquery.query
@@ -265,6 +266,7 @@ def _to_query_request(
     query: str,
     location: Optional[str] = None,
     timeout: Optional[float] = None,
+    timestamp_precision: Optional[enums.TimestampPrecision] = None,
 ) -> Dict[str, Any]:
     """Transform from Job resource to QueryRequest resource.
 
@@ -290,6 +292,12 @@ def _to_query_request(
     request_body.setdefault("formatOptions", {})
     request_body["formatOptions"]["useInt64Timestamp"] = True  # type: ignore
 
+    if timestamp_precision == enums.TimestampPrecision.PICOSECOND:
+        # Cannot specify both use_int64_timestamp and timestamp_output_format.
+        del request_body["formatOptions"]["useInt64Timestamp"]
+
+        request_body["formatOptions"]["timestampOutputFormat"] = "ISO8601_STRING"
+
     if timeout is not None:
         # Subtract a buffer for context switching, network latency, etc.
         request_body["timeoutMs"] = max(0, int(1000 * timeout) - _TIMEOUT_BUFFER_MILLIS)
@@ -370,14 +378,19 @@ def query_jobs_query(
     retry: retries.Retry,
     timeout: Optional[float],
     job_retry: Optional[retries.Retry],
+    timestamp_precision: Optional[enums.TimestampPrecision] = None,
 ) -> job.QueryJob:
     """Initiate a query using jobs.query with jobCreationMode=JOB_CREATION_REQUIRED.
 
     See: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
     """
     path = _to_query_path(project)
     request_body = _to_query_request(
-        query=query, job_config=job_config, location=location, timeout=timeout
+        query=query,
+        job_config=job_config,
+        location=location,
+        timeout=timeout,
+        timestamp_precision=timestamp_precision,
     )
 
     def do_query():

google/cloud/bigquery/client.py

@@ -3469,6 +3469,8 @@ def query(
         timeout: TimeoutType = DEFAULT_TIMEOUT,
         job_retry: Optional[retries.Retry] = DEFAULT_JOB_RETRY,
         api_method: Union[str, enums.QueryApiMethod] = enums.QueryApiMethod.INSERT,
+        *,
+        timestamp_precision: Optional[enums.TimestampPrecision] = None,
     ) -> job.QueryJob:
         """Run a SQL query.
 
@@ -3524,6 +3526,11 @@ def query(
 
                 See :class:`google.cloud.bigquery.enums.QueryApiMethod` for
                 details on the difference between the query start methods.
+            timestamp_precision (Optional[enums.TimestampPrecision]):
+                [Private Preview] If set to `enums.TimestampPrecision.PICOSECOND`,
+                timestamp columns of picosecond precision will be returned with
+                full precision. Otherwise, will truncate to microsecond
+                precision. Only applies when api_method == `enums.QueryApiMethod.QUERY`.
 
         Returns:
             google.cloud.bigquery.job.QueryJob: A new query job instance.
@@ -3543,6 +3550,15 @@ def query(
                 "`job_id` was provided, but the 'QUERY' `api_method` was requested."
             )
 
+        if (
+            timestamp_precision == enums.TimestampPrecision.PICOSECOND
+            and api_method == enums.QueryApiMethod.INSERT
+        ):
+            raise ValueError(
+                "Picosecond Timestamp is only supported when `api_method "
+                "== enums.QueryApiMethod.QUERY`."
+            )
+
         if project is None:
             project = self.project
 
@@ -3568,6 +3584,7 @@ def query(
                 retry,
                 timeout,
                 job_retry,
+                timestamp_precision=timestamp_precision,
             )
         elif api_method == enums.QueryApiMethod.INSERT:
             return _job_helpers.query_jobs_insert(

tests/system/test_query.py

@@ -21,6 +21,7 @@
 import pytest
 
 from google.cloud import bigquery
+from google.cloud.bigquery import enums
 from google.cloud.bigquery.query import ArrayQueryParameter
 from google.cloud.bigquery.query import ScalarQueryParameter
 from google.cloud.bigquery.query import ScalarQueryParameterType
@@ -546,3 +547,15 @@ def test_session(bigquery_client: bigquery.Client, query_api_method: str):
 
     assert len(rows) == 1
     assert rows[0][0] == 5
+
+
+def test_query_picosecond(bigquery_client: bigquery.Client):
+    job = bigquery_client.query(
+        "SELECT CAST('2025-10-20' AS TIMESTAMP(12));",
+        api_method="QUERY",
+        timestamp_precision=enums.TimestampPrecision.PICOSECOND,
+    )
+
+    result = job.result()
+    rows = list(result)
+    assert rows[0][0] == "2025-10-20T00:00:00.000000000000Z"

tests/unit/test__job_helpers.py

@@ -335,6 +335,7 @@ def test_query_jobs_query_defaults():
     assert request["location"] == "asia-northeast1"
     assert request["formatOptions"]["useInt64Timestamp"] is True
     assert "timeoutMs" not in request
+    assert "timestampOutputFormat" not in request["formatOptions"]
 
 
 def test_query_jobs_query_sets_format_options():
@@ -400,6 +401,35 @@ def test_query_jobs_query_sets_timeout(timeout, expected_timeout):
     assert request["timeoutMs"] == expected_timeout
 
 
+def test_query_jobs_query_picosecond():
+    mock_client = mock.create_autospec(Client)
+    mock_retry = mock.create_autospec(retries.Retry)
+    mock_job_retry = mock.create_autospec(retries.Retry)
+    mock_client._call_api.return_value = {
+        "jobReference": {
+            "projectId": "test-project",
+            "jobId": "abc",
+            "location": "asia-northeast1",
+        }
+    }
+    _job_helpers.query_jobs_query(
+        mock_client,
+        "SELECT * FROM test",
+        None,
+        "asia-northeast1",
+        "test-project",
+        mock_retry,
+        None,
+        mock_job_retry,
+        enums.TimestampPrecision.PICOSECOND,
+    )
+
+    _, call_kwargs = mock_client._call_api.call_args
+    request = call_kwargs["data"]
+    assert "useInt64Timestamp" not in request["formatOptions"]
+    assert request["formatOptions"]["timestampOutputFormat"] == "ISO8601_STRING"
+
+
 def test_query_and_wait_uses_jobs_insert():
     """With unsupported features, call jobs.insert instead of jobs.query."""
     client = mock.create_autospec(Client)

tests/unit/test_client.py

@@ -5214,6 +5214,56 @@ def test_query_w_query_parameters(self):
             },
         )
 
+    def test_query_pico_timestamp(self):
+        query = "select *;"
+        response = {
+            "jobReference": {
+                "projectId": self.PROJECT,
+                "location": "EU",
+                "jobId": "abcd",
+            },
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(response)
+
+        client.query(
+            query,
+            location="EU",
+            api_method="QUERY",
+            timestamp_precision=TimestampPrecision.PICOSECOND,
+        )
+
+        # Check that query actually starts the job.
+        expected_resource = {
+            "query": query,
+            "useLegacySql": False,
+            "location": "EU",
+            "formatOptions": {"timestampOutputFormat": "ISO8601_STRING"},
+            "requestId": mock.ANY,
+        }
+        conn.api_request.assert_called_once_with(
+            method="POST",
+            path=f"/projects/{self.PROJECT}/queries",
+            data=expected_resource,
+            timeout=None,
+        )
+
+    def test_query_pico_timestamp_insert_error(self):
+        query = "select *;"
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+        with pytest.raises(ValueError):
+            client.query(
+                query,
+                location="EU",
+                api_method="INSERT",
+                timestamp_precision=TimestampPrecision.PICOSECOND,
+            )
+
     def test_query_job_rpc_fail_w_random_error(self):
         from google.api_core.exceptions import Unknown
         from google.cloud.bigquery.job import QueryJob