Merge pull request #58 from koichiro/feature/fix-dml-stat

Fix dml_stat error.

Author: koichiro

lib/airbnb_payous/bigquery_gateway.rb

@@ -92,18 +92,7 @@ def load_and_merge!(rows:)
         query_job.wait_until_done!
         raise query_job.error if query_job.failed?
 
-        # Capture DML stats with robust fallbacks
-        if query_job.dml_stats
-          result[:inserted_count] = query_job.dml_stats.inserted_row_count.to_i
-          result[:updated_count] = query_job.dml_stats.updated_row_count.to_i
-          @logger.info("DML Stats - Inserted: #{result[:inserted_count]}, Updated: #{result[:updated_count]}")
-        elsif (affected = query_job.num_dml_affected_rows.to_i) >= 0
-          # For our current MERGE query (which only has INSERT), all affected rows are inserts
-          result[:inserted_count] = affected
-          @logger.info("DML Stats not found, but num_dml_affected_rows is #{affected}. Assuming all are inserts.")
-        else
-          @logger.warn("No DML statistics available for query job #{query_job.job_id}.")
-        end
+        result.merge!(extract_dml_counts(query_job))
 
         result[:mode] = :merge
       end
@@ -172,6 +161,38 @@ def build_merge_query(columns)
       SQL
     end
 
+    def extract_dml_counts(query_job)
+      inserted_count = extract_query_job_metric(query_job, :inserted_row_count)
+      updated_count = extract_query_job_metric(query_job, :updated_row_count)
+
+      if inserted_count || updated_count
+        inserted_count ||= 0
+        updated_count ||= 0
+        @logger.info("DML row counts - Inserted: #{inserted_count}, Updated: #{updated_count}")
+        return { inserted_count:, updated_count: }
+      end
+
+      affected = extract_query_job_metric(query_job, :num_dml_affected_rows)
+      if affected
+        # Current MERGE only inserts new rows, so affected rows can be treated as inserts.
+        @logger.info("DML row counts unavailable; using num_dml_affected_rows=#{affected} as inserted count.")
+        return { inserted_count: affected, updated_count: 0 }
+      end
+
+      @logger.warn("No DML statistics available for query job #{query_job.job_id}.")
+      { inserted_count: 0, updated_count: 0 }
+    end
+
+    def extract_query_job_metric(query_job, method_name)
+      return nil unless query_job.respond_to?(method_name)
+
+      value = query_job.public_send(method_name)
+      value.nil? ? nil : value.to_i
+    rescue StandardError => e
+      @logger.warn("Failed to read #{method_name} from query job #{query_job.job_id}: #{e.message}")
+      nil
+    end
+
     def delete_staging_table(dataset)
       return if dataset.nil?
 

test/bigquery_gateway_test.rb

@@ -91,20 +91,49 @@ def error; nil; end
   class FakeQueryJob < FakeLoadJob
     attr_reader :sql
 
-    def initialize(inserted: 1, updated: 0)
+    def initialize(inserted: 1, updated: 0, supports_dml_stats: true, supports_row_counts: true, supports_affected_rows: true)
       super()
+      @inserted = inserted
+      @updated = updated
+      @supports_dml_stats = supports_dml_stats
+      @supports_row_counts = supports_row_counts
+      @supports_affected_rows = supports_affected_rows
       @dml_stats = Struct.new(:inserted_row_count, :updated_row_count).new(inserted, updated)
     end
 
     def dml_stats
+      raise NoMethodError, "undefined method 'dml_stats'" unless @supports_dml_stats
+
       @dml_stats
     end
 
     def num_dml_affected_rows
-      @dml_stats.inserted_row_count + @dml_stats.updated_row_count
+      raise NoMethodError, "undefined method 'num_dml_affected_rows'" unless @supports_affected_rows
+
+      @inserted + @updated
+    end
+
+    def inserted_row_count
+      raise NoMethodError, "undefined method 'inserted_row_count'" unless @supports_row_counts
+
+      @inserted
+    end
+
+    def updated_row_count
+      raise NoMethodError, "undefined method 'updated_row_count'" unless @supports_row_counts
+
+      @updated
     end
 
     def error; nil; end
+
+    def respond_to?(method_name, include_private = false)
+      return @supports_dml_stats if method_name == :dml_stats
+      return @supports_row_counts if [:inserted_row_count, :updated_row_count].include?(method_name)
+      return @supports_affected_rows if method_name == :num_dml_affected_rows
+
+      super
+    end
   end
 
   class FakeTable
@@ -160,9 +189,10 @@ def table(name)
   class FakeBigquery
     attr_reader :dataset_calls, :create_dataset_calls, :copy_job_calls, :query_job_calls
 
-    def initialize(dataset:, created_dataset: nil)
+    def initialize(dataset:, created_dataset: nil, query_job: nil)
       @dataset = dataset
       @created_dataset = created_dataset || dataset
+      @query_job = query_job || FakeQueryJob.new
       @dataset_calls = []
       @create_dataset_calls = []
       @copy_job_calls = []
@@ -186,7 +216,7 @@ def copy_job(source, destination, write:)
 
     def query_job(sql)
       @query_job_calls << sql
-      FakeQueryJob.new
+      @query_job
     end
   end
 
@@ -256,6 +286,48 @@ def test_merges_into_an_existing_target_table
     assert_empty @bigquery.copy_job_calls
   end
 
+  def test_merges_with_query_job_that_exposes_row_counts_but_not_dml_stats
+    bigquery = FakeBigquery.new(
+      dataset: @dataset,
+      query_job: FakeQueryJob.new(inserted: 2, updated: 3, supports_dml_stats: false, supports_row_counts: true)
+    )
+    gateway = AirbnbPayous::BigqueryGateway.new(
+      project_id: "project",
+      dataset_id: "dataset",
+      table_id: "table",
+      logger: @logger,
+      bigquery: bigquery,
+      storage: @storage
+    )
+
+    result = gateway.load_and_merge!(rows: @rows)
+
+    assert_equal :merge, result[:mode]
+    assert_equal 2, result[:inserted_count]
+    assert_equal 3, result[:updated_count]
+  end
+
+  def test_merges_with_query_job_that_only_exposes_affected_rows
+    bigquery = FakeBigquery.new(
+      dataset: @dataset,
+      query_job: FakeQueryJob.new(inserted: 4, updated: 0, supports_dml_stats: false, supports_row_counts: false, supports_affected_rows: true)
+    )
+    gateway = AirbnbPayous::BigqueryGateway.new(
+      project_id: "project",
+      dataset_id: "dataset",
+      table_id: "table",
+      logger: @logger,
+      bigquery: bigquery,
+      storage: @storage
+    )
+
+    result = gateway.load_and_merge!(rows: @rows)
+
+    assert_equal :merge, result[:mode]
+    assert_equal 4, result[:inserted_count]
+    assert_equal 0, result[:updated_count]
+  end
+
   def test_serializes_dates_and_decimals_for_bigquery_load_jobs
     @gateway.load_and_merge!(rows: @rows)