From f5a253ebdf1ca072c1ea6f978d848523bf1b2848 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Mon, 8 Apr 2024 12:32:47 +0000 Subject: [PATCH] feat: adds billing to opentel --- google/cloud/bigquery/opentelemetry_tracing.py | 8 ++++++++ tests/unit/test_opentelemetry_tracing.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index e2a05e4d0..b5f6bf991 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -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 diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 579d7b1b7..546cc02bd 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -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" @@ -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 @@ -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