fix: validate opentelemetry span job attributes have values (#1327)

shollyman · web-flow · commit 8287af129916 · 2022-08-29T11:59:29.000-07:00
fix: validate opentelemetry span job attributes have values

There are several job properties that are optional.  Existing
opentelemetry instrumentation disallows span attribute keys without
appropriate values, so this change validates field presence before
propagating.
diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py
@@ -107,10 +107,7 @@ def _set_client_attributes(client):
 def _set_job_attributes(job_ref):
     job_attributes = {
         "db.name": job_ref.project,
-        "location": job_ref.location,
-        "num_child_jobs": job_ref.num_child_jobs,
         "job_id": job_ref.job_id,
-        "parent_job_id": job_ref.parent_job_id,
         "state": job_ref.state,
     }
 
@@ -125,4 +122,13 @@ def _set_job_attributes(job_ref):
     if job_ref.ended is not None:
         job_attributes["timeEnded"] = job_ref.ended.isoformat()
 
+    if job_ref.location is not None:
+        job_attributes["location"] = job_ref.location
+
+    if job_ref.parent_job_id is not None:
+        job_attributes["parent_job_id"] = job_ref.parent_job_id
+
+    if job_ref.num_child_jobs is not None:
+        job_attributes["num_child_jobs"] = job_ref.num_child_jobs
+
     return job_attributes
diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py
@@ -164,6 +164,32 @@ def test_default_job_attributes(setup):
             assert span.attributes == expected_attributes
 
 
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
+def test_optional_job_attributes(setup):
+    # This test ensures we don't propagate unset values into span attributes
+    import google.cloud._helpers
+
+    time_created = datetime.datetime(
+        2010, 5, 19, 16, 0, 0, tzinfo=google.cloud._helpers.UTC
+    )
+
+    with mock.patch("google.cloud.bigquery.job._AsyncJob") as test_job_ref:
+        test_job_ref.job_id = "test_job_id"
+        test_job_ref.location = None
+        test_job_ref.project = "test_project_id"
+        test_job_ref.created = time_created
+        test_job_ref.state = "some_job_state"
+        test_job_ref.num_child_jobs = None
+        test_job_ref.parent_job_id = None
+
+        with opentelemetry_tracing.create_span(
+            TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref
+        ) as span:
+            assert span is not None
+            for val in span.attributes.values():
+                assert val is not None
+
+
 @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 def test_default_no_data_leakage(setup):
     import google.auth.credentials
