Merge pull request #5525 from rmosolgo/detailed-trace-active-record-backend

ActiveRecord support for DetailedTrace

Author: rmosolgo

lib/generators/graphql/detailed_trace_generator.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+require 'rails/generators/active_record'
+
+module Graphql
+  module Generators
+    class DetailedTraceGenerator < ::Rails::Generators::Base
+      include ::Rails::Generators::Migration
+      desc "Install GraphQL::Tracing::DetailedTrace for your schema"
+      source_root File.expand_path('../templates', __FILE__)
+
+      class_option :redis,
+        type: :boolean,
+        default: false,
+        desc: "Use Redis for persistence instead of ActiveRecord"
+
+      def self.next_migration_number(dirname)
+        ::ActiveRecord::Generators::Base.next_migration_number(dirname)
+      end
+
+      def install_detailed_traces
+
+        schema_glob = File.expand_path("app/graphql/*_schema.rb", destination_root)
+        schema_file = Dir.glob(schema_glob).first
+        if !schema_file
+          raise ArgumentError, "Failed to find schema definition file (checked: #{schema_glob.inspect})"
+        end
+        schema_file_match = /( *)class ([A-Za-z:]+) < GraphQL::Schema/.match(File.read(schema_file))
+        schema_name = schema_file_match[2]
+        indent = schema_file_match[1] + "  "
+
+        if !options.redis?
+          migration_template 'create_graphql_detailed_traces.erb', 'db/migrate/create_graphql_detailed_traces.rb'
+        end
+
+        log :add_detailed_traces_plugin
+        sentinel = /< GraphQL::Schema\s*\n/m
+        code = <<-RUBY
+#{indent}use GraphQL::Tracing::DetailedTrace#{options.redis? ? ", redis: raise(\"TODO: pass a connection to a persistent redis database\")" : ""}, limit: 50
+
+#{indent}# When this returns true, DetailedTrace will trace the query
+#{indent}# Could use `query.context`, `query.selected_operation_name`, `query.query_string` here
+#{indent}# Could call out to Flipper, etc
+#{indent}def self.detailed_trace?(query)
+#{indent}  rand <= 0.000_1 # one in ten thousand
+#{indent}end
+
+        RUBY
+
+        in_root do
+          inject_into_file schema_file, code, after: sentinel, force: false
+        end
+
+        routes_source = File.read(File.expand_path("config/routes.rb", destination_root))
+        already_has_dashboard = routes_source.include?("GraphQL::Dashboard") ||
+          routes_source.include?("Schema.dashboard") ||
+          routes_source.include?("GraphQL::Pro::Routes::Lazy")
+
+        if (!already_has_dashboard || behavior == :revoke)
+          log :route, "GraphQL::Dashboard"
+          shell.mute do
+            route <<~RUBY
+            # TODO: add authorization to this route and expose it in production
+            # See https://graphql-ruby.org/pro/dashboard.html#authorizing-the-dashboard
+            if Rails.env.development?
+              mount GraphQL::Dashboard, at: "/graphql/dashboard", schema: #{schema_name.inspect}
+            end
+
+            RUBY
+          end
+        end
+      end
+    end
+  end
+end

lib/generators/graphql/templates/create_graphql_detailed_traces.erb

@@ -0,0 +1,10 @@
+class CreateGraphqlDetailedTraces < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :graphql_detailed_traces, force: true do |t|
+      t.integer :begin_ms, null: false
+      t.float :duration_ms, null: false
+      t.text :trace_data, null: false
+      t.string :operation_name, null: false
+    end
+  end
+end

lib/graphql/tracing/detailed_trace.rb

@@ -1,10 +1,13 @@
 # frozen_string_literal: true
+if defined?(ActiveRecord)
+  require "graphql/tracing/detailed_trace/active_record_backend"
+end
 require "graphql/tracing/detailed_trace/memory_backend"
 require "graphql/tracing/detailed_trace/redis_backend"
 
 module GraphQL
   module Tracing
-    # `DetailedTrace` can make detailed profiles for a subset of production traffic.
+    # `DetailedTrace` can make detailed profiles for a subset of production traffic. Install it in Rails with `rails generate graphql:detailed_trace`.
     #
     # When `MySchema.detailed_trace?(query)` returns `true`, a profiler-specific `trace_mode: ...` will be used for the query,
     # overriding the one in `context[:trace_mode]`.
@@ -13,10 +16,19 @@ module Tracing
     # this behavior by extending {DetailedTrace} and overriding {#inspect_object}. You can opt out of debug annotations
     # entirely with `use ..., debug: false` or for a single query with `context: { detailed_trace_debug: false }`.
     #
-    # __Redis__: The sampler stores its results in a provided Redis database. Depending on your needs,
-    # You can configure this database to retain all data (persistent) or to expire data according to your rules.
+    # You can store saved traces in two ways:
+    #
+    # - __ActiveRecord__: With `rails generate graphql:detailed_trace`, a new migration will be added to your app.
+    # That table will be used to store trace data.
+    #
+    # - __Redis__: Pass `redis: ...` to save trace data to a Redis database. Depending on your needs,
+    # you can configure this database to retain all data (persistent) or to expire data according to your rules.
+    #
     # If you need to save traces indefinitely, you can download them from Perfetto after opening them there.
     #
+    # @example Installing with Rails
+    #   rails generate graphql:detailed_trace # optional: --redis
+    #
     # @example Adding the sampler to your schema
     #   class MySchema < GraphQL::Schema
     #     # Add the sampler:
@@ -53,13 +65,16 @@ class DetailedTrace
       # @param redis [Redis] If provided, profiles will be stored in Redis for later review
       # @param limit [Integer] A maximum number of profiles to store
       # @param debug [Boolean] if `false`, it won't create `debug` annotations in Perfetto traces (reduces overhead)
-      def self.use(schema, trace_mode: :profile_sample, memory: false, debug: debug?, redis: nil, limit: nil)
+      # @param model_class [Class<ActiveRecord::Base>] Overrides {ActiveRecordBackend::GraphqlDetailedTrace} if present
+      def self.use(schema, trace_mode: :profile_sample, memory: false, debug: debug?, redis: nil, limit: nil, model_class: nil)
         storage = if redis
           RedisBackend.new(redis: redis, limit: limit)
         elsif memory
           MemoryBackend.new(limit: limit)
+        elsif defined?(ActiveRecord)
+          ActiveRecordBackend.new(limit: limit, model_class: model_class)
         else
-          raise ArgumentError, "Pass `redis: ...` to store traces in Redis for later review"
+          raise ArgumentError, "To store traces, install ActiveRecord or provide `redis: ...`"
         end
         detailed_trace = self.new(storage: storage, trace_mode: trace_mode, debug: debug)
         schema.detailed_trace = detailed_trace

lib/graphql/tracing/detailed_trace/active_record_backend.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Tracing
+    class DetailedTrace
+      class ActiveRecordBackend
+        class GraphqlDetailedTrace < ActiveRecord::Base
+        end
+
+        def initialize(limit: nil, model_class: nil)
+          @limit = limit
+          @model_class = model_class || GraphqlDetailedTrace
+        end
+
+        def traces(last:, before:)
+          gdts = @model_class.all.order("begin_ms DESC")
+          if before
+            gdts = gdts.where("begin_ms < ?", before)
+          end
+          if last
+            gdts = gdts.limit(last)
+          end
+          gdts.map { |gdt| record_to_stored_trace(gdt) }
+        end
+
+        def delete_trace(id)
+          @model_class.where(id: id).destroy_all
+          nil
+        end
+
+        def delete_all_traces
+          @model_class.all.destroy_all
+        end
+
+        def find_trace(id)
+          gdt = @model_class.find_by(id: id)
+          if gdt
+            record_to_stored_trace(gdt)
+          else
+            nil
+          end
+        end
+
+        def save_trace(operation_name, duration_ms, begin_ms, trace_data)
+          gdt = @model_class.create!(
+            begin_ms: begin_ms,
+            operation_name: operation_name,
+            duration_ms: duration_ms,
+            trace_data: Base64.encode64(trace_data),
+          )
+          if @limit
+            @model_class
+              .where("id NOT IN(SELECT id FROM graphql_detailed_traces ORDER BY begin_ms DESC LIMIT ?)", @limit)
+              .delete_all
+          end
+          gdt.id
+        end
+
+        private
+
+        def record_to_stored_trace(gdt)
+          StoredTrace.new(
+            id: gdt.id,
+            begin_ms: gdt.begin_ms,
+            operation_name: gdt.operation_name,
+            duration_ms: gdt.duration_ms,
+            trace_data: Base64.decode64(gdt.trace_data)
+          )
+
+        end
+      end
+    end
+  end
+end

spec/graphql/tracing/detailed_trace/active_record_backend_spec.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+require "spec_helper"
+require_relative "./backend_assertions"
+
+if testing_rails?
+  describe GraphQL::Tracing::DetailedTrace::ActiveRecordBackend do
+    include GraphQLTracingDetailedTraceBackendAssertions
+    def new_backend(**kwargs)
+      GraphQL::Tracing::DetailedTrace::ActiveRecordBackend.new(**kwargs)
+    end
+
+    class DummyModel
+      def self.find_by(id:)
+        OpenStruct.new(
+          trace_data: Base64.encode64("DummyModel##{id}")
+        )
+      end
+    end
+
+    it "can use a custom model class" do
+      schema = Class.new(GraphQL::Schema) do
+        use GraphQL::Tracing::DetailedTrace, model_class: DummyModel
+      end
+
+      assert_equal "DummyModel#1234", schema.detailed_trace.find_trace(1234).trace_data
+    end
+  end
+end

spec/graphql/tracing/detailed_trace_spec.rb

@@ -44,13 +44,23 @@ def self.detailed_trace?(query)
     SamplerSchema.detailed_trace.delete_all_traces
   end
 
-  it "raises when no storage is configured" do
-    err = assert_raises ArgumentError do
-      Class.new(GraphQL::Schema) do
+  if testing_rails?
+    it "defaults to ActiveRecord" do
+      schema = Class.new(GraphQL::Schema) do
         use GraphQL::Tracing::DetailedTrace
       end
+
+      assert_instance_of GraphQL::Tracing::DetailedTrace::ActiveRecordBackend, schema.detailed_trace.instance_variable_get(:@storage)
+    end
+  else
+    it "raises when no storage is configured" do
+      err = assert_raises ArgumentError do
+        Class.new(GraphQL::Schema) do
+          use GraphQL::Tracing::DetailedTrace
+        end
+      end
+      assert_equal "To store traces, install ActiveRecord or provide `redis: ...`", err.message
     end
-    assert_equal "Pass `redis: ...` to store traces in Redis for later review", err.message
   end
 
   it "calls detailed_profile? on a Multiplex" do

spec/integration/rails/generators/graphql/detailed_trace_generator_spec.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+require "spec_helper"
+require "generators/graphql/install_generator"
+require "generators/graphql/detailed_trace_generator"
+
+class GraphQLGeneratorsDetailedTraceGeneratorTest < Rails::Generators::TestCase
+  tests Graphql::Generators::DetailedTraceGenerator
+  destination File.expand_path("../../../tmp/dummy", File.dirname(__FILE__))
+
+  setup do
+    prepare_destination
+    FileUtils.cd(File.join(destination_root, '..')) do
+      `rails new dummy --skip-active-record --skip-test-unit --skip-spring --skip-bundle --skip-webpack-install`
+      Graphql::Generators::InstallGenerator.start(["--skip-graphiql"], { destination_root: destination_root })
+    end
+  end
+
+  test "it creates a migration, installs a route, and adds schema configuration" do
+    run_generator
+    assert_migration "db/migrate/create_graphql_detailed_traces"
+    assert_file "app/graphql/dummy_schema.rb" do |content|
+      assert_includes content, "
+  use GraphQL::Tracing::DetailedTrace, limit: 50
+
+  # When this returns true, DetailedTrace will trace the query
+  # Could use `query.context`, `query.selected_operation_name`, `query.query_string` here
+  # Could call out to Flipper, etc
+  def self.detailed_trace?(query)
+    rand <= 0.000_1 # one in ten thousand
+  end
+"
+    end
+
+    assert_file "config/routes.rb" do |content|
+      assert_includes content, "mount GraphQL::Dashboard, at: \"/graphql/dashboard\", schema: \"DummySchema\""
+    end
+  end
+
+  test "it doesn't duplicate dashboard setup" do
+    routes_path = File.expand_path("config/routes.rb", destination_root)
+    existing_routes = File.read(routes_path)
+    new_routes = existing_routes.sub("draw do\n", "draw do\n  mount DummySchema.dashboard\n")
+    File.write(routes_path, new_routes)
+    run_generator
+    assert_file "config/routes.rb" do |content|
+      refute_includes content, "mount GraphQL::Dashboard"
+    end
+  end
+
+  test "it sets up Redis, too" do
+    run_generator(["--redis"])
+    assert_no_migration "db/migrate/create_graphql_detailed_traces"
+    assert_file "app/graphql/dummy_schema.rb" do |content|
+      assert_includes content, "use GraphQL::Tracing::DetailedTrace, redis: raise(\"TODO: pass a connection to a persistent redis database\"), limit: 50\n\n"
+    end
+
+    assert_file "config/routes.rb" do |content|
+      assert_includes content, "mount GraphQL::Dashboard, at: \"/graphql/dashboard\", schema: \"DummySchema\""
+    end
+  end
+end

spec/support/active_record_setup.rb

@@ -107,6 +107,13 @@
       t.integer :points, null: false
       t.decimal :rating, null: false
     end
+
+    create_table :graphql_detailed_traces, force: true do |t|
+      t.integer :begin_ms
+      t.float :duration_ms
+      t.text :trace_data
+      t.string :operation_name
+    end
   end
 
   class Food < ActiveRecord::Base