Support query result traverse pages

Author: ttanimichi

lib/bigquery-client/jobs.rb

@@ -3,10 +3,19 @@
 module BigQuery
   module Jobs
     def sql(query, options = {})
-      result = jobs_query(query, options)
-      names = result['schema']['fields'].map {|field| field['name'] }
-      types = result['schema']['fields'].map {|field| field['type'] }
-      records = (result['rows'] || []).map {|row| row['f'].map {|record| record['v'] } }
+      jobs_query_response = jobs_query(query, options)
+      names = jobs_query_response['schema']['fields'].map {|field| field['name'] }
+      types = jobs_query_response['schema']['fields'].map {|field| field['type'] }
+      records = extract_records(jobs_query_response)
+
+      page_token = jobs_query_response['pageToken']
+      job_id = jobs_query_response['jobReference']['jobId']
+      while page_token
+        query_results_response = query_results(job_id, pageToken: page_token)
+        records += extract_records(query_results_response)
+        page_token = query_results_response['pageToken']
+      end
+
       convert(records, types).map { |values| [names, values].transpose.to_h }
     end
 
@@ -52,6 +61,10 @@ def query_results(id, options = {})
 
     private
 
+    def extract_records(response)
+      (response['rows'] || []).map {|row| row['f'].map {|record| record['v'] } }
+    end
+
     def convert(records, types)
       records.map do |values|
         values.map.with_index do |value, index|

test/test_jobs.rb

@@ -1,7 +1,7 @@
 require 'helper'
 
 class JobsTest < Test::Unit::TestCase
-  @@query = <<-"EOS"
+  @@normal_query = <<-"EOS"
     SELECT
       born_alive_alive, mother_residence_state, is_male
     FROM
@@ -23,8 +23,17 @@ class JobsTest < Test::Unit::TestCase
       0
   EOS
 
+  @@huge_result_query = <<-"EOS"
+    SELECT
+      title
+    FROM
+      publicdata:samples.wikipedia
+    LIMIT
+      1234567
+  EOS
+
   def test_sql
-    result = $client.sql(@@query)
+    result = $client.sql(@@normal_query)
     assert { result.size == 100 }
     assert { result.sample["born_alive_alive"].is_a? Fixnum }
     assert { result.sample["mother_residence_state"].is_a? String }
@@ -35,4 +44,9 @@ def test_sql_when_no_rows
     result = $client.sql(@@no_rows_query)
     assert { result == [] }
   end
+
+  def test_sql_pagination
+    record_size = $client.sql(@@huge_result_query).size
+    assert { record_size == 1234567 }
+  end
 end