bigshift 0.3.1 → 0.3.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: b44f1ad6233e29754f8ae3b66f5857141660529e
-  data.tar.gz: a42978dfd9b983592e1051a727373b6bde9b73ff
+  metadata.gz: b0a1088f4c4d8c66c8af35a4e67a8377d8b6f805
+  data.tar.gz: bff611a9528b2b08a3587177cb3a448b7dbea4de
 SHA512:
-  metadata.gz: 82fc075bcc22f380ff8492acb83b6a73e5ef6a741be04b8078811d2587efe2bbecd900d5eab5f25a51307afbe3b3475f7768ff5605fc3a781cc70c518c68488d
-  data.tar.gz: 30b626b8b7fce0445a7c8ef613a446d769b04580803a4438558c876775ac99abd81c1f6f3cdff7b58343a1889222f289bff4976d56ca2bb0900813a9fbeba376
+  metadata.gz: 045ba2e30068a4259ac34763f3597d7087053ad8889443c077643576e5cb8df55ad02d0a01e50577b4865afa130a68311188bd4db05bc6c34a632d2ab9bfe39d
+  data.tar.gz: 31cf2ec5852d2a1c200398a089b1be9c3093abb5e17a90cab3f130f43ef50df5d0c5749f117651f4ddc123d6504397fa1602a02dcf0e99d4fed511d320b803cf

data/README.md

@@ -32,7 +32,7 @@ BigShift tells Redshift to compress the dumps by default, even if that means tha
 
 ## Arguments
 
-Running `bigshift` without any arguments, or with `--help` will show the options. All except `--s3-prefix`, `--bq-table`, `--max-bad-records`, `--steps` and `--[no-]compress` are required.
+Running `bigshift` without any arguments, or with `--help` will show the options. All except `--s3-prefix`, `--rs-schema`, `--bq-table`, `--max-bad-records`, `--steps` and `--[no-]compress` are required.
 
 ### GCP credentials
 
@@ -106,6 +106,10 @@ If you don't want to put the data dumped from Redshift directly into the root of
 
 Because of how GCS' Transfer Service works the transferred files will have exactly the same keys in the destination bucket, this cannot be configured.
 
+### Redshift schema
+
+By default the schema in Redshift is called `public`, but in case you're not using that one, you can use the argument `--rs-schema` to specify the schema your table is in.
+
 ### BigQuery table ID
 
 By default the BigQuery table ID will be the same as the Redshift table name, but the optional argument `--bq-table` can be used to tell BigShift to use another table ID.

data/lib/bigshift/cli.rb

@@ -4,6 +4,7 @@ require 'json'
 require 'stringio'
 require 'logger'
 require 'optparse'
+require 'socket'
 require 'bigshift'
 
 module BigShift
@@ -47,7 +48,7 @@ module BigShift
     def unload
       if run?(:unload)
         s3_uri = "s3://#{@config[:s3_bucket_name]}/#{s3_table_prefix}"
-        @factory.redshift_unloader.unload_to(@config[:rs_table_name], s3_uri, allow_overwrite: false, compression: @config[:compression])
+        @factory.redshift_unloader.unload_to(@config[:rs_schema_name], @config[:rs_table_name], s3_uri, allow_overwrite: false, compression: @config[:compression])
       else
         @logger.debug('Skipping unload')
       end
@@ -56,7 +57,7 @@ module BigShift
 
     def transfer
       if run?(:transfer)
-        description = "bigshift-#{@config[:rs_database_name]}-#{@config[:rs_table_name]}-#{Time.now.utc.strftime('%Y%m%dT%H%M')}"
+        description = "bigshift-#{@config[:rs_database_name]}-#{@config[:rs_schema_name]}-#{@config[:rs_table_name]}-#{Time.now.utc.strftime('%Y%m%dT%H%M')}"
         @factory.cloud_storage_transfer.copy_to_cloud_storage(@unload_manifest, @config[:cs_bucket_name], description: description, allow_overwrite: false)
       else
         @logger.debug('Skipping transfer')
@@ -99,6 +100,7 @@ module BigShift
       ['--aws-credentials', 'PATH', String, :aws_credentials_path, nil],
       ['--rs-credentials', 'PATH', String, :rs_credentials_path, :required],
       ['--rs-database', 'DB_NAME', String, :rs_database_name, :required],
+      ['--rs-schema', 'SCHEMA_NAME', String, :rs_schema_name, nil],
       ['--rs-table', 'TABLE_NAME', String, :rs_table_name, :required],
       ['--bq-dataset', 'DATASET_ID', String, :bq_dataset_id, :required],
       ['--bq-table', 'TABLE_ID', String, :bq_table_id, nil],
@@ -136,6 +138,7 @@ module BigShift
         end
       end
       config[:bq_table_id] ||= config[:rs_table_name]
+      config[:rs_schema_name] ||= 'public'
       if config[:steps] && !config[:steps].empty?
         config[:steps] = STEPS.select { |s| config[:steps].include?(s.to_s) }
       else
@@ -150,8 +153,9 @@ module BigShift
     def s3_table_prefix
       @s3_table_prefix ||= begin
         db_name = @config[:rs_database_name]
+        schema_name = @config[:rs_schema_name]
         table_name = @config[:rs_table_name]
-        prefix = "#{db_name}/#{table_name}/#{db_name}-#{table_name}-"
+        prefix = "#{db_name}/#{schema_name}/#{table_name}/#{db_name}-#{schema_name}-#{table_name}-"
         if (s3_prefix = @config[:s3_prefix])
           s3_prefix = s3_prefix.gsub(%r{\A/|/\Z}, '')
           prefix = "#{s3_prefix}/#{prefix}"
@@ -175,7 +179,7 @@ module BigShift
     end
 
     def redshift_table_schema
-      @redshift_table_schema ||= RedshiftTableSchema.new(@config[:rs_table_name], rs_connection)
+      @redshift_table_schema ||= RedshiftTableSchema.new(@config[:rs_schema_name], @config[:rs_table_name], rs_connection)
     end
 
     def big_query_dataset
@@ -212,6 +216,13 @@ module BigShift
         password: @config[:rs_credentials]['password'],
         sslmode: 'require'
       )
+      socket = Socket.for_fd(@rs_connection.socket)
+      socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, 1)
+      socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPCNT, 5)
+      socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPINTVL, 2)
+      socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPIDLE, 2) if defined?(Socket::TCP_KEEPIDLE)
+      @rs_connection.exec("SET search_path = \"#{@config[:rs_schema_name]}\"")
+      @rs_connection
     end
 
     def cs_transfer_service

data/lib/bigshift/redshift_table_schema.rb

@@ -1,15 +1,16 @@
 module BigShift
   class RedshiftTableSchema
-    def initialize(table_name, redshift_connection)
+    def initialize(schema_name, table_name, redshift_connection)
+      @schema_name = schema_name
       @table_name = table_name
       @redshift_connection = redshift_connection
     end
 
     def columns
       @columns ||= begin
-        rows = @redshift_connection.exec_params(%|SELECT "column", "type", "notnull" FROM "pg_table_def" WHERE "schemaname" = 'public' AND "tablename" = $1|, [@table_name])
+        rows = @redshift_connection.exec_params(%|SELECT "column", "type", "notnull" FROM "pg_table_def" WHERE "schemaname" = $1 AND "tablename" = $2|, [@schema_name, @table_name])
         if rows.count == 0
-          raise sprintf('Table not found: %s', @table_name.inspect)
+          raise sprintf('Table %s for schema %s not found', @table_name.inspect, @schema_name.inspect)
         else
           columns = rows.map do |row|
             name = row['column']

data/lib/bigshift/redshift_unloader.rb

@@ -6,12 +6,12 @@ module BigShift
       @logger = options[:logger] || NullLogger::INSTANCE
     end
 
-    def unload_to(table_name, s3_uri, options={})
-      table_schema = RedshiftTableSchema.new(table_name, @redshift_connection)
+    def unload_to(schema_name, table_name, s3_uri, options={})
+      table_schema = RedshiftTableSchema.new(schema_name, table_name, @redshift_connection)
       credentials_string = "aws_access_key_id=#{@aws_credentials.access_key_id};aws_secret_access_key=#{@aws_credentials.secret_access_key}"
       select_sql = 'SELECT '
       select_sql << table_schema.columns.map(&:to_sql).join(', ')
-      select_sql << %Q< FROM "#{table_name}">
+      select_sql << %Q< FROM "#{schema_name}"."#{table_name}">
       select_sql.gsub!('\'') { |s| '\\\'' }
       unload_sql = %Q<UNLOAD ('#{select_sql}')>
       unload_sql << %Q< TO '#{s3_uri}'>

data/lib/bigshift/version.rb

@@ -1,3 +1,3 @@
 module BigShift
-  VERSION = '0.3.1'.freeze
+  VERSION = '0.3.2'.freeze
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: bigshift
 version: !ruby/object:Gem::Version
-  version: 0.3.1
+  version: 0.3.2
 platform: ruby
 authors:
 - Theo Hultberg
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2016-05-12 00:00:00.000000000 Z
+date: 2016-08-19 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pg
@@ -112,7 +112,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 2.4.8
+rubygems_version: 2.4.5
 signing_key: 
 specification_version: 4
 summary: A tool for moving tables from Redshift to BigQuery