lib/elastic_apm/spies/mongo.rb in elastic-apm-1.1.0 vs lib/elastic_apm/spies/mongo.rb in elastic-apm-2.0.0

@@ -12,11 +12,11 @@
         )
       end
 
       # @api private
       class Subscriber
-        TYPE = 'db.mongodb.query'.freeze
+        TYPE = 'db.mongodb.query'
 
         def initialize
           @events = {}
         end
 
@@ -35,24 +35,35 @@
         private
 
         def push_event(event)
           return unless ElasticAPM.current_transaction
 
-          ctx = Span::Context.new(
-            instance: event.database_name,
-            statement: nil,
-            type: 'mongodb'.freeze,
-            user: nil
-          )
-          span = ElasticAPM.span(event.command_name.to_s, TYPE, context: ctx)
+          span =
+            ElasticAPM.start_span(
+              event.command_name.to_s,
+              TYPE,
+              context: build_context(event)
+            )
+
           @events[event.operation_id] = span
         end
 
         def pop_event(event)
-          return unless ElasticAPM.current_transaction
-
+          return unless (curr = ElasticAPM.current_span)
           span = @events.delete(event.operation_id)
-          span && span.done
+
+          curr == span && ElasticAPM.end_span
+        end
+
+        def build_context(event)
+          Span::Context.new(
+            db: {
+              instance: event.database_name,
+              statement: nil,
+              type: 'mongodb',
+              user: nil
+            }
+          )
         end
       end
     end
 
     register 'Mongo', 'mongo', MongoSpy.new