Make generator do other DetailedTrace setup and support Redis, add tests, add docs

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/detailed_traces_generator.rb

@@ -7,7 +7,7 @@
 
 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]`.
@@ -16,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:
@@ -56,13 +65,14 @@ 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)
+          ActiveRecordBackend.new(limit: limit, model_class: model_class)
         else
           raise ArgumentError, "To store traces, install ActiveRecord or provide `redis: ...`"
         end

lib/graphql/tracing/detailed_trace.rb

@@ -7,12 +7,13 @@ class ActiveRecordBackend
         class GraphqlDetailedTrace < ActiveRecord::Base
         end
 
-        def initialize(limit: nil)
+        def initialize(limit: nil, model_class: nil)
           @limit = limit
+          @model_class = model_class || GraphqlDetailedTrace
         end
 
         def traces(last:, before:)
-          gdts = GraphqlDetailedTrace.all.order("begin_ms DESC")
+          gdts = @model_class.all.order("begin_ms DESC")
           if before
             gdts = gdts.where("begin_ms < ?", before)
           end
@@ -23,16 +24,16 @@ def traces(last:, before:)
         end
 
         def delete_trace(id)
-          GraphqlDetailedTrace.where(id: id).destroy_all
+          @model_class.where(id: id).destroy_all
           nil
         end
 
         def delete_all_traces
-          GraphqlDetailedTrace.all.destroy_all
+          @model_class.all.destroy_all
         end
 
         def find_trace(id)
-          gdt = GraphqlDetailedTrace.find_by(id: id)
+          gdt = @model_class.find_by(id: id)
           if gdt
             record_to_stored_trace(gdt)
           else
@@ -41,14 +42,14 @@ def find_trace(id)
         end
 
         def save_trace(operation_name, duration_ms, begin_ms, trace_data)
-          gdt = GraphqlDetailedTrace.create!(
+          gdt = @model_class.create!(
             begin_ms: begin_ms,
             operation_name: operation_name,
             duration_ms: duration_ms,
             trace_data: Base64.encode64(trace_data),
           )
           if @limit
-            GraphqlDetailedTrace
+            @model_class
               .where("id NOT IN(SELECT id FROM graphql_detailed_traces ORDER BY begin_ms DESC LIMIT ?)", @limit)
               .delete_all
           end

lib/graphql/tracing/detailed_trace/active_record_backend.rb

@@ -8,5 +8,21 @@
     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/active_record_backend_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