Refactor OpenTelemetry context propagation and cleanup specs

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    umbrellio-sequel-plugins (0.18.1)
+    umbrellio-sequel-plugins (0.19.0)
       concurrent-ruby
       sequel
 

lib/sequel/extensions/concurrent_thread_pool.rb

@@ -120,14 +120,14 @@ def choose_executor(opts)
       end
 
       def async_run(&)
-        if defined?(OpenTelemetry)
+        if defined?(OpenTelemetry::Context)
           otel_context = OpenTelemetry::Context.current
-          async_job_class.new(async_thread_executor) do
-            token = OpenTelemetry::Context.attach(otel_context)
-            begin
-              yield
-            ensure
-              OpenTelemetry::Context.detach(token)
+
+          if otel_context.equal?(OpenTelemetry::Context::ROOT)
+            async_job_class.new(async_thread_executor, &)
+          else
+            async_job_class.new(async_thread_executor) do
+              OpenTelemetry::Context.with_current(otel_context, &)
             end
           end
         else

lib/sequel/plugins/upsert.rb

@@ -25,10 +25,10 @@ def upsert_dataset(target: primary_key)
     # Executes the upsert request
     #
     # @param row [Hash] values
-    # @param options [Hash] options
+    # @param target [Symbol] target column
     #
     # @example
-    #   User.upsert(name: "John", email: "jd@test.com", target: :email)
+    #   User.upsert({name: "John", email: "jd@test.com"}, target: :email)
     # @return [Sequel::Model]
     def upsert(row, **)
       upsert_dataset(**).insert(sequel_values(row))

spec/extensions/concurrent_thread_pool_helpers.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+ASYNC_DB_URL = ENV.fetch("DB_URL", "postgres:///sequel_plugins")
+
+def make_concurrent_db(**opts)
+  Sequel.connect(ASYNC_DB_URL, **opts).tap { |d| d.extension(:concurrent_thread_pool) }
+end

spec/extensions/concurrent_thread_pool_opentelemetry_spec.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "opentelemetry-api"
+require_relative "concurrent_thread_pool_helpers"
+
+RSpec.describe "concurrent_thread_pool extension OpenTelemetry context propagation" do
+  let(:db) { make_concurrent_db }
+  let(:otel_key) { OpenTelemetry::Context.create_key("sequel-test") }
+
+  after { db.disconnect rescue nil }
+
+  it "does not propagate context when calling thread context is ROOT" do
+    captured = :not_set
+
+    db.send(:async_run) { captured = OpenTelemetry::Context.current.value(otel_key) }.__value
+
+    expect(captured).to be_nil
+  end
+
+  it "does not error when OpenTelemetry is defined without Context" do
+    db = make_concurrent_db
+    stub_const("OpenTelemetry", Object.new)
+
+    expect { db.send(:async_run) { :ok }.__value }.not_to raise_error
+  ensure
+    db.disconnect rescue nil
+  end
+
+  it "propagates calling-thread context to worker thread" do
+    captured = nil
+
+    OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+      db.send(:async_run) { captured = OpenTelemetry::Context.current.value(otel_key) }.__value
+    end
+
+    expect(captured).to eq("sentinel")
+  end
+
+  context "with single-thread executor" do
+    let(:db) { make_concurrent_db(num_async_threads: 1) }
+
+    it "restores context in worker thread after block completes" do
+      OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+        db.send(:async_run) { :done }.__value
+
+        value_after = Concurrent::Promises.future_on(db.async_thread_executor) do
+          OpenTelemetry::Context.current.value(otel_key)
+        end.value!
+
+        expect(value_after).to be_nil
+      end
+    end
+  end
+
+  it "propagates context even when block raises" do
+    captured = nil
+
+    OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+      result = db.send(:async_run) do
+        captured = OpenTelemetry::Context.current.value(otel_key)
+        raise "boom"
+      end
+
+      expect { result.__value }.to raise_error(RuntimeError, "boom")
+      expect(captured).to eq("sentinel")
+    end
+  end
+end

spec/extensions/concurrent_thread_pool_spec.rb

@@ -1,12 +1,6 @@
 # frozen_string_literal: true
 
-require "concurrent"
-
-ASYNC_DB_URL = ENV.fetch("DB_URL", "postgres:///sequel_plugins")
-
-def make_concurrent_db(**opts)
-  Sequel.connect(ASYNC_DB_URL, **opts).tap { |d| d.extension(:concurrent_thread_pool) }
-end
+require_relative "concurrent_thread_pool_helpers"
 
 RSpec.describe "concurrent_thread_pool extension" do
   let(:db) { make_concurrent_db(**db_opts) }
@@ -172,61 +166,6 @@ def make_concurrent_db(**opts)
     end
   end
 
-  describe "OpenTelemetry context propagation" do
-    require "opentelemetry-api"
-
-    let(:otel_key) { OpenTelemetry::Context.create_key("sequel-test") }
-
-    after do
-      # ensure no context leak between examples
-      OpenTelemetry::Context.detach(OpenTelemetry::Context.attach(OpenTelemetry::Context::ROOT))
-    end
-
-    it "propagates calling-thread context to worker thread" do
-      ctx = OpenTelemetry::Context.current.set_value(otel_key, "sentinel")
-      token = OpenTelemetry::Context.attach(ctx)
-      captured = nil
-
-      db.send(:async_run) { captured = OpenTelemetry::Context.current.value(otel_key) }.__value
-
-      expect(captured).to eq("sentinel")
-    ensure
-      OpenTelemetry::Context.detach(token)
-    end
-
-    it "restores context in worker thread after block completes" do
-      ctx = OpenTelemetry::Context.current.set_value(otel_key, "sentinel")
-      token = OpenTelemetry::Context.attach(ctx)
-      value_after = :not_set
-
-      db.send(:async_run) do
-        inner_token = OpenTelemetry::Context.attach(OpenTelemetry::Context::ROOT)
-        OpenTelemetry::Context.detach(inner_token)
-        value_after = OpenTelemetry::Context.current.value(otel_key)
-      end.__value
-
-      expect(value_after).to eq("sentinel")
-    ensure
-      OpenTelemetry::Context.detach(token)
-    end
-
-    it "propagates context even when block raises" do
-      ctx = OpenTelemetry::Context.current.set_value(otel_key, "sentinel")
-      token = OpenTelemetry::Context.attach(ctx)
-      captured = nil
-
-      result = db.send(:async_run) do
-        captured = OpenTelemetry::Context.current.value(otel_key)
-        raise "boom"
-      end
-
-      expect { result.__value }.to raise_error(RuntimeError, "boom")
-      expect(captured).to eq("sentinel")
-    ensure
-      OpenTelemetry::Context.detach(token)
-    end
-  end
-
   describe "async_job_class" do
     it "is Proxy by default" do
       expect(db.async_job_class).to eq(Sequel::Database::ConcurrentThreadPool::Proxy)

umbrellio-sequel-plugins.gemspec

@@ -5,7 +5,7 @@ $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 
 Gem::Specification.new do |spec|
   spec.name = "umbrellio-sequel-plugins"
-  spec.version = "0.18.1"
+  spec.version = "0.19.0"
   spec.required_ruby_version = ">= 3.3"
 
   spec.authors = ["Team Umbrellio"]