Reparent tracing spans in sqlite

Signed-off-by: Caleb Schoepp <[email protected]>

crates/factor-sqlite/Cargo.toml

@@ -11,6 +11,7 @@ rust-version.workspace = true
 [dependencies]
 async-trait = "0.1"
 serde = { version = "1.0", features = ["rc"] }
+spin-factor-observe = { path = "../factor-observe" }
 spin-factors = { path = "../factors" }
 spin-locked-app = { path = "../locked-app" }
 spin-world = { path = "../world" }
@@ -26,5 +27,3 @@ tokio = { version = "1", features = ["macros", "rt"] }
 
 [lints]
 workspace = true
-
-# TODO(Caleb): Setup observe reparenting

crates/factor-sqlite/src/host.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 
+use spin_factor_observe::ObserveContext;
 use spin_factors::wasmtime::component::Resource;
 use spin_factors::{anyhow, SelfInstanceBuilder};
 use spin_world::v1::sqlite as v1;
@@ -16,6 +17,7 @@ pub struct InstanceState {
     allowed_databases: Arc<HashSet<String>>,
     connections: table::Table<Box<dyn Connection>>,
     get_connection_creator: ConnectionCreatorGetter,
+    observe_context: ObserveContext,
 }
 
 impl InstanceState {
@@ -35,11 +37,13 @@ impl InstanceState {
     pub fn new(
         allowed_databases: Arc<HashSet<String>>,
         get_connection_creator: ConnectionCreatorGetter,
+        observe_context: ObserveContext,
     ) -> Self {
         Self {
             allowed_databases,
             connections: table::Table::new(256),
             get_connection_creator,
+            observe_context,
         }
     }
 
@@ -66,6 +70,8 @@ impl v2::Host for InstanceState {
 impl v2::HostConnection for InstanceState {
     #[instrument(name = "spin_sqlite.open", skip(self), err(level = Level::INFO), fields(otel.kind = "client", db.system = "sqlite", sqlite.backend = Empty))]
     async fn open(&mut self, database: String) -> Result<Resource<v2::Connection>, v2::Error> {
+        self.observe_context.reparent_tracing_span();
+
         if !self.allowed_databases.contains(&database) {
             return Err(v2::Error::AccessDenied);
         }
@@ -90,6 +96,8 @@ impl v2::HostConnection for InstanceState {
         query: String,
         parameters: Vec<v2::Value>,
     ) -> Result<v2::QueryResult, v2::Error> {
+        self.observe_context.reparent_tracing_span();
+
         let conn = match self.get_connection(connection) {
             Ok(c) => c,
             Err(err) => return Err(err),

crates/factor-sqlite/src/lib.rs

@@ -7,6 +7,7 @@ use std::sync::Arc;
 use host::InstanceState;
 
 use async_trait::async_trait;
+use spin_factor_observe::ObserveContext;
 use spin_factors::{anyhow, Factor};
 use spin_locked_app::MetadataKey;
 use spin_world::v1::sqlite as v1;
@@ -86,7 +87,7 @@ impl Factor for SqliteFactor {
 
     fn prepare<T: spin_factors::RuntimeFactors>(
         &self,
-        ctx: spin_factors::PrepareContext<T, Self>,
+        mut ctx: spin_factors::PrepareContext<T, Self>,
     ) -> spin_factors::anyhow::Result<Self::InstanceBuilder> {
         let allowed_databases = ctx
             .app_state()
@@ -95,9 +96,11 @@ impl Factor for SqliteFactor {
             .cloned()
             .unwrap_or_default();
         let get_connection_creator = ctx.app_state().get_connection_creator.clone();
+        let observe_context = ObserveContext::from_prepare_context(&mut ctx)?;
         Ok(InstanceState::new(
             allowed_databases,
             get_connection_creator,
+            observe_context,
         ))
     }
 }