BigQuery sample: import data from Cloud Storage file

Author: remi Taylor

bigquery_sample/Gemfile

@@ -20,4 +20,5 @@ group :test do
   gem "rake"
   gem "rubocop"
   gem "rspec"
+  gem "google-cloud-storage"
 end

bigquery_sample/Gemfile.lock

@@ -4,6 +4,7 @@ GEM
     addressable (2.4.0)
     ast (2.3.0)
     diff-lcs (1.2.5)
+    digest-crc (0.4.1)
     faraday (0.9.2)
       multipart-post (>= 1.2, < 3)
     google-api-client (0.9.13)
@@ -19,6 +20,10 @@ GEM
       google-api-client (~> 0.9.11)
       google-cloud-core (~> 0.20.0)
     google-cloud-core (0.20.1)
+    google-cloud-storage (0.20.0)
+      digest-crc (~> 0.4)
+      google-api-client (~> 0.9.11)
+      google-cloud-core (~> 0.20.0)
     googleauth (0.5.1)
       faraday (~> 0.9)
       jwt (~> 1.4)
@@ -82,6 +87,7 @@ PLATFORMS
 
 DEPENDENCIES
   google-cloud-bigquery
+  google-cloud-storage
   rake
   rspec
   rubocop

bigquery_sample/bigquery_samples.rb

@@ -138,9 +138,10 @@ def list_table_data project_id:, dataset_id:, table_id:
 def import_table_data_from_file project_id:, dataset_id:, table_id:,
                                 local_file_path:
   # [START import_table_data_from_file]
-  # project_id = "Your Google Cloud project ID"
-  # dataset_id = "ID of the dataset delete table from"
-  # table_id   = "ID of the table to import file data into"
+  # project_id      = "Your Google Cloud project ID"
+  # dataset_id      = "ID of the dataset delete table from"
+  # table_id        = "ID of the table to import file data into"
+  # local_file_path = "Path to local file to import into BigQuery table"
 
   require "google/cloud"
 
@@ -159,6 +160,31 @@ def import_table_data_from_file project_id:, dataset_id:, table_id:,
   # [END import_table_data_from_file]
 end
 
+def import_table_data_from_cloud_storage project_id:, dataset_id:, table_id:,
+                                         storage_path:
+  # [START import_table_data_from_file]
+  # project_id   = "Your Google Cloud project ID"
+  # dataset_id   = "ID of the dataset delete table from"
+  # table_id     = "ID of the table to import file data into"
+  # storage_path = "Storage path to file to import, eg. gs://bucket/file.csv"
+
+  require "google/cloud"
+
+  gcloud   = Google::Cloud.new project_id
+  bigquery = gcloud.bigquery
+  dataset  = bigquery.dataset dataset_id
+  table    = dataset.table table_id
+
+  puts "Importing data from Cloud Storage file: #{storage_path}"
+  load_job = table.load storage_path
+
+  puts "Waiting for load job to complete: #{load_job.job_id}"
+  load_job.wait_until_done!
+
+  puts "Data imported"
+  # [END import_table_data_from_file]
+end
+
 # TODO: separate sample into separate executable files
 #
 if __FILE__ == $PROGRAM_NAME

bigquery_sample/spec/bigquery_sample_spec.rb

@@ -17,12 +17,16 @@
 require "google/cloud"
 require "csv"
 
+# TODO: refactor CSV creation to simplify & slightly DRY specs
+
 RSpec.describe "Google Cloud BigQuery samples" do
 
   before do
     @project_id = ENV["GOOGLE_PROJECT_ID"]
     @gcloud     = Google::Cloud.new @project_id
     @bigquery   = @gcloud.bigquery
+    @storage    = @gcloud.storage
+    @bucket     = @storage.bucket ENV["STORAGE_BUCKET"]
 
     # Examples assume that test_dataset does not exist
     test_dataset = @bigquery.dataset "test_dataset"
@@ -214,11 +218,64 @@ def capture &block
       end
     end
 
-    example "import data from Cloud Storage"
+    example "import data from Cloud Storage" do
+      begin
+        dataset = @bigquery.create_dataset "test_dataset"
+
+        table = dataset.create_table "test_table" do |schema|
+          schema.string  "name"
+          schema.integer "value"
+        end
+
+        csv_file = Tempfile.new %w[ bigquery-test csv ]
+
+        CSV.open csv_file.path, "w" do |csv|
+          csv << [ "Alice", 5 ]
+          csv << [ "Bob",   10 ]
+        end
+
+        file = @bucket.create_file csv_file.path, "bigquery-test.csv"
+
+        expect(table.data).to be_empty
+
+        capture do
+          import_table_data_from_cloud_storage(
+            project_id:   @project_id,
+            dataset_id:   "test_dataset",
+            table_id:     "test_table",
+            storage_path: "gs://#{@bucket.name}/bigquery-test.csv"
+          )
+        end
+
+        expect(captured_output).to include(
+          "Importing data from Cloud Storage file: " +
+          "gs://#{@bucket.name}/bigquery-test.csv"
+        )
+        expect(captured_output).to match(
+          /Waiting for load job to complete: job_\w+/
+        )
+        expect(captured_output).to include "Data imported"
+
+        loaded_data = table.data
+
+        expect(loaded_data).not_to be_empty
+        expect(loaded_data.count).to eq 2
+        expect(loaded_data.first["name"]).to eq "Alice"
+        expect(loaded_data.first["value"]).to eq 5
+        expect(loaded_data.last["name"]).to eq "Bob"
+        expect(loaded_data.last["value"]).to eq 10
+      ensure
+        csv_file.flush
+        csv_file.close
+      end
+    end
+
     example "stream data import"
   end
 
   describe "Exporting data" do
+    # Needs a CSV file to import into the table before running
+    # the export command, so refactor CSV code before writing this
     example "export data to Cloud Storage"
   end