diff --git a/example/migrate_partitioned_table.yml b/example/migrate_partitioned_table.yml index f932704..b220dcf 100644 --- a/example/migrate_partitioned_table.yml +++ b/example/migrate_partitioned_table.yml @@ -19,5 +19,8 @@ actions: - action: migrate_partitioned_table <<: *bigquery schema_file: example/schema.json + options: + time_partitioning: + field: timestamp - action: delete_table <<: *bigquery diff --git a/lib/bigquery_migration/bigquery_wrapper.rb b/lib/bigquery_migration/bigquery_wrapper.rb index 56f2ffd..fd2a160 100644 --- a/lib/bigquery_migration/bigquery_wrapper.rb +++ b/lib/bigquery_migration/bigquery_wrapper.rb @@ -188,6 +188,7 @@ def insert_table(dataset: nil, table: nil, columns:, options: {}) body[:time_partitioning] = { type: options['time_partitioning']['type'], expiration_ms: options['time_partitioning']['expiration_ms'], + field: options['time_partitioning']['field'], } end @@ -217,7 +218,8 @@ def insert_table(dataset: nil, table: nil, columns:, options: {}) alias :create_table :insert_table def insert_partitioned_table(dataset: nil, table: nil, columns:, options: {}) - options['time_partitioning'] = {'type'=>'DAY'} + options['time_partitioning'] ||= {} + options['time_partitioning']['type'] = 'DAY' insert_table(dataset: dataset, table: table, columns: columns, options: options) end alias :create_partitioned_table :insert_partitioned_table diff --git a/test/test_bigquery_wrapper.rb b/test/test_bigquery_wrapper.rb index a010f94..27bebcd 100644 --- a/test/test_bigquery_wrapper.rb +++ b/test/test_bigquery_wrapper.rb @@ -661,6 +661,7 @@ def teardown def test_create_partitioned_table columns = [ + { name: 'partition_column', type: 'TIMESTAMP' }, { name: 'remained_column', type: 'INTEGER' }, { name: 'record', type: 'RECORD', fields: [ { name: 'record', type: 'RECORD', fields: [ @@ -668,12 +669,18 @@ def test_create_partitioned_table ] } ] } ] + options = { + 'time_partitioning' => { 'expiration_ms' => 86400000, 'field' => 'partition_column' }, + } + expected = columns.dup - result = instance.migrate_partitioned_table(columns: columns) + result = instance.migrate_partitioned_table(columns: columns, options: options) after_columns = result[:after_columns] assert { result[:responses][:insert_table].time_partitioning.type == 'DAY' } + assert { result[:responses][:insert_table].time_partitioning.expiration_ms == 86400000 } + assert { result[:responses][:insert_table].time_partitioning.field == 'partition_column' } assert { Schema.diff_columns(expected, after_columns) == [] } assert { Schema.diff_columns(after_columns, expected) == [] } end