From d019ed9c23f132d36b8a48d86780f73fda4b42db Mon Sep 17 00:00:00 2001
From: Caio Garcia <caioif4@gmail.com>
Date: Tue, 17 Sep 2024 15:17:29 -0300
Subject: [PATCH] metrics_recorder

---
 lib/dbt/metrics_recorder.rb | 58 +++++++++++++++++++++++++++++++++++++
 lib/dbt/runner.rb           | 11 ++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 lib/dbt/metrics_recorder.rb

diff --git a/lib/dbt/metrics_recorder.rb b/lib/dbt/metrics_recorder.rb
new file mode 100644
index 0000000..d674c73
--- /dev/null
+++ b/lib/dbt/metrics_recorder.rb
@@ -0,0 +1,58 @@
+module Dbt
+  class MetricsRecorder
+    def initialize(schema)
+      @schema = schema
+      @metrics = []
+    end
+
+    def record(model_name, model_path, start_time, end_time)
+      @metrics << {
+        model_name: model_name,
+        model_path: model_path,
+        start_time: start_time,
+        end_time: end_time,
+        duration: end_time - start_time
+      }
+    end
+
+    def save
+      begin
+        # Create the metrics table if it doesn't exist
+        ActiveRecord::Base.connection.execute <<-SQL
+          CREATE TABLE IF NOT EXISTS #{@schema}.dbt_internal_metrics (
+            id SERIAL PRIMARY KEY,
+            run_id INTEGER,
+            model_name TEXT,
+            model_path TEXT,
+            start_time TIMESTAMP,
+            end_time TIMESTAMP,
+            duration FLOAT
+          )
+        SQL
+
+        # Generate a new run_id
+        result = ActiveRecord::Base.connection.execute("SELECT MAX(run_id) AS last_run_id FROM #{@schema}.dbt_internal_metrics")
+        last_run_id = result.first["last_run_id"] || 0
+        run_id = last_run_id + 1
+
+        # Insert all metrics into the database
+        @metrics.each do |metric|
+          ActiveRecord::Base.connection.execute <<-SQL
+            INSERT INTO #{@schema}.dbt_internal_metrics (run_id, model_name, model_path, start_time, end_time, duration)
+            VALUES (
+              #{run_id},
+              '#{metric[:model_name]}',
+              '#{metric[:model_path]}',
+              '#{metric[:start_time]}',
+              '#{metric[:end_time]}',
+              #{metric[:duration]}
+            )
+          SQL
+        end
+      rescue => e
+        # Handle exceptions within the class
+        puts "Warning: MetricsRecorder encountered an error during save: #{e.message}"
+      end
+    end
+  end
+end
diff --git a/lib/dbt/runner.rb b/lib/dbt/runner.rb
index 680c609..d6a2b0b 100644
--- a/lib/dbt/runner.rb
+++ b/lib/dbt/runner.rb
@@ -12,9 +12,18 @@ def run(custom_schema = nil, glob_path = "app/sql/**/*.sql")
         graph = Dagwood::DependencyGraph.new dependencies
         md = Mermaid.markdown_for dependencies
         Mermaid.generate_file md
+        metrics_recorder = MetricsRecorder.new(schema)
         graph.order.each do |model_name|
-          models.find { |m| m.name == model_name }.build
+          model = models.find { |m| m.name == model_name }
+          start_time = Time.now
+          model.build
+          end_time = Time.now
+          # Record the timing data
+          metrics_recorder.record(model.name, model.filepath, start_time, end_time)
         end
+        # Save all metrics after the run completes
+        metrics_recorder.save
+        graph.order
       end
 
       def test