Skip to content

Commit

Permalink
feat: adds billing to opentel (#1889)
Browse files Browse the repository at this point in the history
  • Loading branch information
chalmerlowe authored Apr 11, 2024
1 parent 3634405 commit 38697fb
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 0 deletions.
8 changes: 8 additions & 0 deletions google/cloud/bigquery/opentelemetry_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,4 +153,12 @@ def _set_job_attributes(job_ref):
if job_ref.num_child_jobs is not None:
job_attributes["num_child_jobs"] = job_ref.num_child_jobs

total_bytes_billed = getattr(job_ref, "total_bytes_billed", None)
if total_bytes_billed is not None:
job_attributes["total_bytes_billed"] = total_bytes_billed

total_bytes_processed = getattr(job_ref, "total_bytes_processed", None)
if total_bytes_processed is not None:
job_attributes["total_bytes_processed"] = total_bytes_processed

return job_attributes
6 changes: 6 additions & 0 deletions tests/unit/test_opentelemetry_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,8 @@ def test_default_job_attributes(setup):
"timeEnded": ended_time.isoformat(),
"hasErrors": True,
"state": "some_job_state",
"total_bytes_billed": 42,
"total_bytes_processed": 13,
}
with mock.patch("google.cloud.bigquery.job._AsyncJob") as test_job_ref:
test_job_ref.job_id = "test_job_id"
Expand All @@ -154,6 +156,8 @@ def test_default_job_attributes(setup):
test_job_ref.ended = ended_time
test_job_ref.error_result = error_result
test_job_ref.state = "some_job_state"
test_job_ref.total_bytes_billed = 42
test_job_ref.total_bytes_processed = 13

with opentelemetry_tracing.create_span(
TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref
Expand All @@ -180,6 +184,8 @@ def test_optional_job_attributes(setup):
test_job_ref.state = "some_job_state"
test_job_ref.num_child_jobs = None
test_job_ref.parent_job_id = None
test_job_ref.total_bytes_billed = None
test_job_ref.total_bytes_processed = None

with opentelemetry_tracing.create_span(
TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref
Expand Down

0 comments on commit 38697fb

Please sign in to comment.