|
| 1 | +// Copyright 2026 Google LLC |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use crate::e2e::wait_for_trace; |
| 16 | +use google_cloud_spanner::client::Spanner; |
| 17 | +use google_cloud_test_utils::runtime_config::project_id; |
| 18 | +use opentelemetry::trace::TraceContextExt; |
| 19 | +use std::collections::BTreeSet; |
| 20 | +use tracing_opentelemetry::OpenTelemetrySpanExt; |
| 21 | + |
| 22 | +const ROOT_SPAN_NAME: &str = "e2e-spanner-test"; |
| 23 | + |
| 24 | +pub async fn spanner_e2e_tracing() -> anyhow::Result<()> { |
| 25 | + let project_id = project_id()?; |
| 26 | + // Create a trace with a number of interesting spans from the |
| 27 | + // `google-cloud-spanner` client. |
| 28 | + let trace_id = send_trace(&project_id).await?; |
| 29 | + let required = BTreeSet::from_iter([ |
| 30 | + ROOT_SPAN_NAME, |
| 31 | + "google.spanner.v1.Spanner/CreateSession", |
| 32 | + "google.spanner.v1.Spanner/BeginTransaction", |
| 33 | + "Spanner.CreateSession", |
| 34 | + "Spanner.BeginTransaction", |
| 35 | + ]); |
| 36 | + let trace = wait_for_trace(&project_id, &trace_id, &required).await?; |
| 37 | + |
| 38 | + println!("TRACE SPANS DUMP:"); |
| 39 | + for span in &trace.spans { |
| 40 | + println!("Span: {:?}", span); |
| 41 | + } |
| 42 | + |
| 43 | + // Verify the expected spans appear in the trace: |
| 44 | + let span_names = trace |
| 45 | + .spans |
| 46 | + .iter() |
| 47 | + .map(|s| s.name.as_str()) |
| 48 | + .collect::<BTreeSet<_>>(); |
| 49 | + let missing = required.difference(&span_names).collect::<Vec<_>>(); |
| 50 | + assert!(missing.is_empty(), "missing={missing:?}\n\n{trace:?}"); |
| 51 | + |
| 52 | + Ok(()) |
| 53 | +} |
| 54 | + |
| 55 | +async fn send_trace(project_id: &str) -> anyhow::Result<String> { |
| 56 | + // 1. Setup Telemetry (Google Cloud Destination) |
| 57 | + let creds = google_cloud_auth::credentials::Builder::default().build()?; |
| 58 | + let (provider, _, _) = crate::e2e::set_up_providers(project_id, "e2e-telemetry-test", "spanner-test".to_string(), creds).await?; |
| 59 | + |
| 60 | + // 2. Generate Trace |
| 61 | + // Start a root span |
| 62 | + let root_span = tracing::info_span!("e2e_root", { "otel.name" } = ROOT_SPAN_NAME); |
| 63 | + let trace_id = root_span |
| 64 | + .context() |
| 65 | + .span() |
| 66 | + .span_context() |
| 67 | + .trace_id() |
| 68 | + .to_string(); |
| 69 | + |
| 70 | + use tracing::Instrument; |
| 71 | + let _ = client_library_operations().instrument(root_span).await; |
| 72 | + |
| 73 | + println!( |
| 74 | + "View generated trace in Console: https://console.cloud.google.com/traces/explorer;traceId={}?project={}", |
| 75 | + trace_id, project_id |
| 76 | + ); |
| 77 | + |
| 78 | + // 4. Force flush to ensure spans are sent. |
| 79 | + if let Err(e) = provider.force_flush() { |
| 80 | + tracing::error!("error flushing provider: {e:}"); |
| 81 | + } |
| 82 | + Ok(trace_id) |
| 83 | +} |
| 84 | + |
| 85 | +async fn client_library_operations() -> anyhow::Result<()> { |
| 86 | + // Explicitly opt-in to E2E tracing headers for the test |
| 87 | + unsafe { |
| 88 | + std::env::set_var("GOOGLE_CLOUD_TEST_EXTRA_HEADERS", "x-goog-spanner-end-to-end-tracing=true"); |
| 89 | + } |
| 90 | + let project = |
| 91 | + std::env::var("GOOGLE_CLOUD_PROJECT").unwrap_or_else(|_| "westarle-78rabs".to_string()); |
| 92 | + let instance = std::env::var("GOOGLE_CLOUD_SPANNER_TEST_INSTANCE") |
| 93 | + .unwrap_or_else(|_| "trace-propagation-test-instance".to_string()); |
| 94 | + let db_id = std::env::var("GOOGLE_CLOUD_SPANNER_TEST_DATABASE") |
| 95 | + .unwrap_or_else(|_| "test-database".to_string()); |
| 96 | + |
| 97 | + let db_path = format!( |
| 98 | + "projects/{}/instances/{}/databases/{}", |
| 99 | + project, instance, db_id |
| 100 | + ); |
| 101 | + |
| 102 | + use google_cloud_auth::credentials::Builder as CredentialsBuilder; |
| 103 | + let creds = CredentialsBuilder::default().build()?; |
| 104 | + let spanner_client = Spanner::builder() |
| 105 | + .with_credentials(creds.clone()) |
| 106 | + .with_tracing() |
| 107 | + .build() |
| 108 | + .await?; |
| 109 | + |
| 110 | + // Calling `build()` on the database client triggers a `CreateSession` RPC |
| 111 | + let db_client = spanner_client.database_client(db_path).build().await?; |
| 112 | + |
| 113 | + use google_cloud_spanner::model::{ |
| 114 | + BeginTransactionRequest, TransactionOptions, transaction_options, |
| 115 | + }; |
| 116 | + let mut req = BeginTransactionRequest::default(); |
| 117 | + req.session = db_client.session.name.clone(); |
| 118 | + |
| 119 | + let mut options = TransactionOptions::default(); |
| 120 | + options.mode = Some(transaction_options::Mode::ReadOnly(Box::new( |
| 121 | + transaction_options::ReadOnly::default(), |
| 122 | + ))); |
| 123 | + req.options = Some(options); |
| 124 | + |
| 125 | + let _ = db_client |
| 126 | + .spanner |
| 127 | + .begin_transaction(req, google_cloud_gax::options::RequestOptions::default()) |
| 128 | + .await; |
| 129 | + |
| 130 | + Ok(()) |
| 131 | +} |
0 commit comments