data_taster 0.4.2 → 0.5.0

data/lib/data_taster/adapters/database_output.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class DatabaseOutput < Output
+    attr_reader :target_client
+
+    def sample!
+      source = DataTaster.config.source
+      table_names(source).each do |table_name|
+        collection = DataTaster::Collection.new(table_name)
+        import_table(collection, table_name)
+      end
+    end
+
+    def run_sanitization?
+      true
+    end
+
+    def default_data
+      { "schema_migrations" => "1 = 1" }
+    end
+
+    def import_table(collection, table_name)
+      payload = collection.assemble
+      return write_drop_table(table_name) if payload.empty?
+
+      export(collection, table_name, payload)
+    end
+
+    def target_database
+      target_client.query_options[:database]
+    end
+
+    def qualified_table_name(table_name)
+      "#{target_database}.#{table_name}"
+    end
+
+    def write_statement(sql)
+      safe_execute(sql)
+    end
+
+    def write_raw(line)
+      write_statement(line)
+    end
+
+  private
+
+    def safe_execute(sql)
+      foreign_key_check = target_client.query("SELECT @@FOREIGN_KEY_CHECKS").first["@@FOREIGN_KEY_CHECKS"]
+
+      begin
+        target_client.query("SET FOREIGN_KEY_CHECKS=0")
+        target_client.query(sql)
+      ensure
+        target_client.query("SET FOREIGN_KEY_CHECKS=#{foreign_key_check};")
+      end
+    end
+
+    def table_names(source)
+      source.table_names
+    end
+
+    def export(collection, table_name, payload)
+      write_statement("TRUNCATE TABLE #{qualified_table_name(table_name)}")
+      insert_table_name = "#{ExportContext.quote_ident(target_database)}.#{ExportContext.quote_ident(table_name)}"
+      Sanitizer.new(table_name, payload[:sanitize]).export_sanitized_rows(
+        collection,
+        insert_table_name: insert_table_name
+      ) { |header, values| write_sanitized_insert(header, values) }
+    end
+
+    def write_drop_table(table_name)
+      write_statement("DROP TABLE IF EXISTS #{qualified_table_name(table_name)}")
+    end
+
+    def write_sanitized_insert(header, values)
+      write_statement("#{header} #{values}")
+    end
+  end
+end

data/lib/data_taster/adapters/file_output.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class FileOutput < Output
+    attr_reader :path, :target_database
+
+    def sample!
+      start_export
+
+      process_export
+
+      finish_export
+    end
+
+    def default_data
+      {}
+    end
+
+    def run_sanitization?
+      false
+    end
+
+    def qualified_table_name(table_name)
+      DataTaster::ExportContext.quote_ident(table_name)
+    end
+
+    def write_statement(sql)
+      @io.puts "#{sql};"
+    end
+
+    def write_raw(line)
+      @io.puts line
+    end
+
+  private
+
+    def process_export
+      table_names.each do |table_name|
+        collection = DataTaster::Collection.new(table_name)
+        export_table(collection, table_name)
+      end
+    end
+
+    def start_export
+      DataTaster.logger.info("Writing SQL file to #{path}")
+      @io = File.open(path, "w")
+      @io.puts "SET FOREIGN_KEY_CHECKS=0;"
+    end
+
+    def finish_export
+      @io.puts "SET FOREIGN_KEY_CHECKS=1;"
+      @io.close
+    end
+
+    def export_table(collection, table_name)
+      payload = collection.assemble
+      return if payload.empty?
+
+      insert_table_name = ExportContext.quote_ident(table_name)
+      Sanitizer.new(table_name, payload[:sanitize]).export_sanitized_rows(
+        collection,
+        insert_table_name: insert_table_name
+      ) { |header, values| write_sanitized_insert(header, values) }
+    end
+
+    def table_names
+      DataTaster.confection.keys
+    end
+
+    def write_sanitized_insert(header, values)
+      write_raw(header)
+      write_raw("#{values};")
+    end
+  end
+end

data/lib/data_taster/adapters/mysql_source.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class MysqlSource
+    attr_reader :source_client
+
+    def initialize(source_client:)
+      @source_client = source_client
+    end
+
+    def query(sql)
+      source_client.query(sql)
+    end
+
+    def database
+      source_client.query_options[:database]
+    end
+
+    def table_names
+      query("SHOW tables").map { |row| row[row.keys.first] }
+    end
+  end
+end

data/lib/data_taster/adapters/output.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class Output
+    def initialize(**options)
+      options.each { |key, value| instance_variable_set(:"@#{key}", value) }
+    end
+
+    def sample!
+      raise NotImplementedError
+    end
+
+    def default_data
+      raise NotImplementedError
+    end
+
+    def run_sanitization?
+      raise NotImplementedError
+    end
+
+    def target_database
+      raise NotImplementedError
+    end
+
+    def qualified_table_name(table_name)
+      raise NotImplementedError
+    end
+
+    def write_statement(_sql)
+      raise NotImplementedError
+    end
+
+    def write_raw(_line)
+      raise NotImplementedError
+    end
+  end
+end

data/lib/data_taster/collection.rb

@@ -12,7 +12,6 @@ module DataTaster
     def initialize(table_name)
       @table_name = table_name
       @ingredients = DataTaster.confection[table_name]
-      @include_insert = DataTaster.config.include_insert
     end
 
     def assemble
@@ -26,9 +25,16 @@ module DataTaster
       end
     end
 
+    def export_select_sql
+      <<-SQL.squish
+        SELECT * FROM #{source_db}.#{table_name}
+        WHERE #{where_clause}
+      SQL
+    end
+
   private
 
-    attr_reader :table_name, :ingredients, :include_insert
+    attr_reader :table_name, :ingredients
 
     def skippable?
       table_name.downcase.match(/^_/) ||
@@ -36,14 +42,7 @@ module DataTaster
     end
 
     def selection
-      insert = include_insert ? "INSERT INTO #{working_db}.#{table_name}" : ""
-
-      sql = <<-SQL.squish
-        #{insert}
-        SELECT * FROM #{source_db}.#{table_name}
-        WHERE #{where_clause}
-      SQL
-
+      sql = export_select_sql
       DataTaster.logger.info(sql)
       sql
     end
@@ -64,11 +63,7 @@ module DataTaster
     end
 
     def source_db
-      @source_db ||= DataTaster.config.source_client.query_options[:database]
-    end
-
-    def working_db
-      @working_db ||= DataTaster.config.working_client.query_options[:database]
+      @source_db ||= DataTaster.config.source.database
     end
   end
 end

data/lib/data_taster/confection.rb

@@ -30,9 +30,7 @@ module DataTaster
   private
 
     def default_data
-      {
-        "schema_migrations" => "1 = 1",
-      }
+      DataTaster.config.output.default_data
     end
   end
 end

data/lib/data_taster/detergent.rb

@@ -30,10 +30,35 @@ module DataTaster
       sql
     end
 
+    def insert_value_expression(row, client, column_types: {})
+      return DataTaster::SKIP_CODE if value == DataTaster::SKIP_CODE
+
+      if sanitize_function?
+        wash_values(value.to_s, row, client, column_types)
+      elsif value.blank?
+        "NULL"
+      else
+        DataTaster::SqlLiteral.format(client, value, column_type: column_types[column_name.to_s])
+      end
+    end
+
   private
 
     attr_reader :table_name, :column_name, :value
 
+    def wash_values(expression, row, client, column_types = {})
+      result = expression.dup
+      row.keys.map(&:to_s).grep(/\A[a-zA-Z_][a-zA-Z0-9_]*\z/).sort_by { |k| -k.size }.each do |key|
+        next unless result.match?(/\b#{Regexp.escape(key)}\b/)
+
+        result.gsub!(
+          /\b#{Regexp.escape(key)}\b/,
+          SqlLiteral.format(client, row[key], column_type: column_types[key])
+        )
+      end
+      result
+    end
+
     def parse_value(given_value)
       # yml files can't hold custom-set dates, they have to be converted to strings
       return given_value unless given_value.is_a?(String) && given_value.match?(/\d{4}-\d{2}-\d{2}/)
@@ -59,7 +84,7 @@ module DataTaster
 
     def sql_for_uncast_value
       <<-SQL.squish
-        UPDATE #{working_db}.#{table_name}
+        UPDATE #{qualified_table_name}
         SET #{column_name} = #{value}
         WHERE #{column_name} IS NOT NULL
         AND #{column_name} <> #{value}
@@ -68,7 +93,7 @@ module DataTaster
 
     def sql_for_date_value
       <<-SQL.squish
-        UPDATE #{working_db}.#{table_name}
+        UPDATE #{qualified_table_name}
         SET #{column_name} = '#{value}'
         WHERE #{column_name} IS NOT NULL
       SQL
@@ -76,7 +101,7 @@ module DataTaster
 
     def sql_for_nil_value
       <<-SQL.squish
-        UPDATE #{working_db}.#{table_name}
+        UPDATE #{qualified_table_name}
         SET #{column_name} = NULL
         WHERE #{column_name} IS NOT NULL
       SQL
@@ -84,15 +109,15 @@ module DataTaster
 
     def sql_for_cast_value
       <<-SQL.squish
-        UPDATE #{working_db}.#{table_name}
+        UPDATE #{qualified_table_name}
         SET #{column_name} = '#{value}'
         WHERE #{column_name} IS NOT NULL
         AND #{column_name} <> ''
       SQL
     end
 
-    def working_db
-      DataTaster.config.working_client.query_options[:database]
+    def qualified_table_name
+      DataTaster.config.output.qualified_table_name(table_name)
     end
   end
 end

data/lib/data_taster/export_context.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class ExportContext
+    def initialize(table_name, sanitize, insert_table_name: nil, column_types: {})
+      @table_name = table_name
+      @insert_table_name = insert_table_name || self.class.quote_ident(table_name)
+      @rules = DataTaster::Sanitizer.new(table_name, sanitize).sanitization_rules
+      @column_types = column_types.transform_keys(&:to_s)
+    end
+
+    def self.quote_ident(name)
+      "`#{name.to_s.gsub('`', '``')}`"
+    end
+
+    attr_reader :insert_table_name
+
+    def format_row_tuple(columns, row, client)
+      frags = columns.map do |col|
+        col_key = col.to_s
+        if @rules.key?(col_key) # the colunm has a sanitization rule
+          DataTaster::Detergent.new(@table_name, col_key, @rules[col_key])
+                               .insert_value_expression(row, client, column_types: @column_types)
+        else
+          DataTaster::SqlLiteral.format(client, row[col], column_type: @column_types[col_key])
+        end
+      end
+      "(#{frags.join(', ')})"
+    end
+  end
+end

data/lib/data_taster/helper.rb

@@ -14,7 +14,9 @@ module DataTaster
     end
 
     def db_config
-      ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: "primary").configuration_hash
+      ActiveRecord::Base.configurations
+                        .configs_for(env_name: Rails.env, name: "primary")
+                        .configuration_hash
     end
 
     def logg(message)

data/lib/data_taster/sanitizer.rb

@@ -2,12 +2,12 @@
 
 module DataTaster
   class Sanitizer
-    # Ensures the given tables are cleaned of
-    # information deemed sensitive
+    BATCH_SIZE = SanitizerExporter::BATCH_SIZE
+
+    # Ensures the given tables are cleaned of information deemed sensitive
     def initialize(table_name, custom_selections)
       @table_name = table_name
       @custom_selections = custom_selections || {}
-      @include_insert = DataTaster.config.include_insert
     end
 
     def clean!
@@ -23,19 +23,33 @@ module DataTaster
       end
     end
 
+    def sanitization_rules
+      return {} if skippable_table?
+
+      default_selections.merge(custom_selections).each_with_object({}) do |(column_name, sanitized_value), memo|
+        next if sanitized_value == DataTaster::SKIP_CODE
+
+        memo[column_name.to_s] = sanitized_value
+      end
+    end
+
+    def export_sanitized_rows(collection, insert_table_name:, &write_insert)
+      SanitizerExporter.new(table_name, custom_selections, insert_table_name: insert_table_name)
+                       .export_sanitized_rows(collection, &write_insert)
+    end
+
   private
 
-    attr_reader :table_name, :custom_selections, :include_insert
+    attr_reader :table_name, :custom_selections
 
     def process(sql)
       return if sql == DataTaster::SKIP_CODE
-      return sql unless include_insert
+      return sql unless output.run_sanitization?
 
-      DataTaster.safe_execute(sql)
+      output.write_statement(sql)
+      sql
     rescue => e
-      e.message << context_warning
-
-      raise e
+      raise e, e.message + context_warning
     end
 
     def context_warning
@@ -50,6 +64,10 @@ module DataTaster
       WARNING
     end
 
+    def output
+      DataTaster.config.output
+    end
+
     def skippable_table?
       DataTaster.confection[table_name].blank? ||
         DataTaster.confection[table_name] == DataTaster::SKIP_CODE

data/lib/data_taster/sanitizer_exporter.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module DataTaster
+  class SanitizerExporter
+    BATCH_SIZE = 100
+
+    def initialize(table_name, custom_selections, insert_table_name:)
+      @table_name = table_name
+      @custom_selections = custom_selections
+      @insert_table_name = insert_table_name
+    end
+
+    def export_sanitized_rows(collection, &write_insert)
+      query_result = export_source.query(collection.export_select_sql)
+
+      columns = query_result.fields
+      return if columns.empty?
+
+      export_context = build_export_context(columns, query_result)
+      process_export_in_batches(export_context, columns, query_result, &write_insert)
+    end
+
+  private
+
+    attr_reader :table_name, :custom_selections, :insert_table_name
+
+    def process_export_in_batches(export_context, columns, query_result, &write_insert)
+      batch = []
+      query_result.each do |row|
+        batch << row
+        if batch.size >= BATCH_SIZE
+          write_export_batch(export_context, columns, batch, &write_insert)
+          batch.clear
+        end
+      end
+      write_export_batch(export_context, columns, batch, &write_insert) if batch.any?
+    end
+
+    def write_export_batch(export_context, columns, rows, &_write_insert)
+      return if rows.empty?
+
+      client = export_source.source_client
+      col_list = columns.map { |c| ExportContext.quote_ident(c) }.join(", ")
+      tuples = rows.map { |row| export_context.format_row_tuple(columns, row, client) }
+      header = "INSERT INTO #{export_context.insert_table_name} (#{col_list}) VALUES"
+      values = tuples.join(",\n")
+
+      yield(header, values)
+    end
+
+    def export_source
+      DataTaster.config.source
+    end
+
+    def build_export_context(columns, query_result)
+      column_types = columns.zip(query_result.field_types).to_h.transform_keys(&:to_s)
+      ExportContext.new(
+        table_name,
+        custom_selections,
+        insert_table_name: insert_table_name,
+        column_types: column_types
+      )
+    end
+  end
+end

data/lib/data_taster/sql_literal.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+require "date"
+
+module DataTaster
+  # Necessary to format Ruby values (the non sanitized values) for use as MySQL INSERT literal fragments
+  module SqlLiteral
+    BINARY_COLUMN_TYPE = /\A(?:tiny|medium|long)?blob|varbinary|binary|bit|geometry/i
+    TEXT_COLUMN_TYPE = /\Ajson|(?:var)?char|(?:tiny|medium|long)?text|enum|set/i
+
+    def self.format(client, value, column_type: nil)
+      case value
+      when String
+        format_string_value(client, value, column_type: column_type)
+      when nil, true, false, Integer, Float, BigDecimal
+        format_scalar(value)
+      when Time, DateTime, Date
+        format_temporal(client, value)
+      else
+        "'#{client.escape(value.to_s)}'"
+      end
+    end
+
+    def self.format_temporal(client, value)
+      case value
+      when Time
+        format_escaped_time(client, value, "%Y-%m-%d %H:%M:%S")
+      when DateTime
+        format_escaped_time(client, value.to_time, "%Y-%m-%d %H:%M:%S")
+      when Date
+        format_escaped_date(client, value)
+      end
+    end
+    private_class_method :format_temporal
+
+    def self.format_scalar(value)
+      return "NULL" if value.nil?
+      return "TRUE" if value == true
+      return "FALSE" if value == false
+      return value.to_s if value.is_a?(Integer)
+      return value.finite? ? value.to_s : "NULL" if value.is_a?(Float)
+
+      value.to_s("F")
+    end
+    private_class_method :format_scalar
+
+    def self.format_escaped_time(client, time, strftime_format)
+      "'#{client.escape(time.strftime(strftime_format))}'"
+    end
+    private_class_method :format_escaped_time
+
+    def self.format_escaped_date(client, value)
+      "'#{client.escape(value.strftime('%Y-%m-%d'))}'"
+    end
+    private_class_method :format_escaped_date
+
+    def self.format_string_value(client, value, column_type: nil)
+      if binary_column?(value, column_type)
+        "X'#{value.unpack1('H*')}'"
+      else
+        "'#{client.escape(encode_as_text(value))}'"
+      end
+    end
+    private_class_method :format_string_value
+
+    def self.binary_column?(value, column_type)
+      return true if binary_column_type?(column_type)
+      return false if text_column_type?(column_type)
+
+      value.encoding == Encoding::ASCII_8BIT && !utf8_text?(value)
+    end
+    private_class_method :binary_column?
+
+    def self.binary_column_type?(column_type)
+      return false if column_type.nil?
+
+      column_type.to_s.match?(BINARY_COLUMN_TYPE)
+    end
+    private_class_method :binary_column_type?
+
+    def self.text_column_type?(column_type)
+      return false if column_type.nil?
+
+      column_type.to_s.match?(TEXT_COLUMN_TYPE)
+    end
+    private_class_method :text_column_type?
+
+    def self.utf8_text?(value)
+      value.dup.force_encoding(Encoding::UTF_8).valid_encoding?
+    end
+    private_class_method :utf8_text?
+
+    def self.encode_as_text(value)
+      return value unless value.encoding == Encoding::ASCII_8BIT
+
+      candidate = value.dup.force_encoding(Encoding::UTF_8)
+      candidate.valid_encoding? ? candidate : value
+    end
+    private_class_method :encode_as_text
+  end
+end

data/lib/data_taster/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module DataTaster
-  VERSION = "0.4.2"
+  VERSION = "0.5.0"
 end