Merge pull request #93 from potato2003/master

Add support for clustered table

Author: potato2003

README.md

@@ -108,6 +108,8 @@ Following options are same as [bq command-line tools](https://cloud.google.com/b
 |  time_partitioning.expiration_ms  | int      | optional  | nil     | Number of milliseconds for which to keep the storage for a partition. |
 |  time_partitioning.field          | string   | optional  | nil     | `DATE` or `TIMESTAMP` column used for partitioning |
 |  time_partitioning.require_partition_filter | boolean      | optional  | nil     | If true, valid partition filter is required when query |
+|  clustering                       | hash     | optional  | nil     | (Experimental) Currently, clustering is supported for partitioned tables, so must be used with `time_partitioning` option. NOTE: **clustered tables** is a beta release. See [clustered tables](https://cloud.google.com/bigquery/docs/clustered-tables) |
+|  clustering.fields                | array    | required  | nil     | One or more fields on which data should be clustered. The order of the specified columns determines the sort order of the data. |
 |  schema_update_options            | array    | optional  | nil     | (Experimental) List of `ALLOW_FIELD_ADDITION` or `ALLOW_FIELD_RELAXATION` or both. See [jobs#configuration.load.schemaUpdateOptions](https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.load.schemaUpdateOptions). NOTE for the current status: `schema_update_options` does not work for `copy` job, that is, is not effective for most of modes such as `append`, `replace` and `replace_backup`. `delete_in_advance` deletes origin table so does not need to update schema. Only `append_direct` can utilize schema update. |
 
 ### Example

lib/embulk/output/bigquery.rb

@@ -89,6 +89,7 @@ def self.configure(config, schema, task_count)
           'ignore_unknown_values'          => config.param('ignore_unknown_values',          :bool,    :default => false),
           'allow_quoted_newlines'          => config.param('allow_quoted_newlines',          :bool,    :default => false),
           'time_partitioning'              => config.param('time_partitioning',              :hash,    :default => nil),
+          'clustering'                     => config.param('clustering',                     :hash,    :default => nil), # google-api-ruby-client >= v0.21.0
           'schema_update_options'          => config.param('schema_update_options',          :array,   :default => nil),
 
           # for debug
@@ -234,6 +235,12 @@ def self.configure(config, schema, task_count)
           task['time_partitioning'] = {'type' => 'DAY'}
         end
 
+        if task['clustering']
+          unless task['clustering']['fields']
+            raise ConfigError.new "`clustering` must have `fields` key"
+          end
+        end
+
         if task['schema_update_options']
           task['schema_update_options'].each do |schema_update_option|
             unless %w[ALLOW_FIELD_ADDITION ALLOW_FIELD_RELAXATION].include?(schema_update_option)

lib/embulk/output/bigquery/bigquery_client.rb

@@ -441,6 +441,13 @@ def create_table(table, dataset: nil, options: nil)
               }
             end
 
+            options['clustering'] ||= @task['clustering']
+            if options['clustering']
+              body[:clustering] = {
+                fields: options['clustering']['fields'],
+              }
+            end
+
             opts = {}
             Embulk.logger.debug { "embulk-output-bigquery: insert_table(#{@project}, #{dataset}, #{@location_for_log}, #{body}, #{opts})" }
             with_network_retry { client.insert_table(@project, dataset, body, opts) }

test/test_configure.rb

@@ -84,6 +84,7 @@ def test_configure_default
         assert_equal false, task['ignore_unknown_values']
         assert_equal false, task['allow_quoted_newlines']
         assert_equal nil, task['time_partitioning']
+        assert_equal nil, task['clustering']
         assert_equal false, task['skip_load']
       end
 
@@ -277,6 +278,14 @@ def test_time_partitioning
         assert_equal 'DAY', task['time_partitioning']['type']
       end
 
+      def test_clustering
+        config = least_config.merge('clustering' => {'fields' => ['field_a']})
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('clustering' => {})
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
       def test_schema_update_options
         config = least_config.merge('schema_update_options' => ['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION'])
         assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }