db_subsetter 0.4.1 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: abf5990b7e3fd33ac74ba829e56919a5d6cc84c1
-  data.tar.gz: 3c83e352af29441d90675d8f1c924bdc80dfad14
+SHA256:
+  metadata.gz: 45469d6d8b6183ae244553ee87165494cd0a5ecef1b843cd73f2be192a554c2c
+  data.tar.gz: 87f7ab413790695a3c5808817f8c502b793c028a7bfdbe6637ecf6c62831fb15
 SHA512:
-  metadata.gz: 0fb729b7e0db5e1de5133eaf082fb3e6067f4aaaa370b303a07c49d5a18bd3355f97a3a1ab43c3c41b508d470c9a126abac969d311c649fac9b0aa6bac23278e
-  data.tar.gz: f0637b18e81c8fb5414cac673b2f94d954bc8c10e1d7398ca8589c976069107ea22f5aff37d6aa7db7778f7173b76fc3056e6d3e0015931584c4f3d5d104ca0f
+  metadata.gz: a9c78f51a33be41e7c6f84a0e9bb35e48cb21041846c57eb3d01467f0a85664aa124f260661f73a0d95a73ea2a6ab729db1604e5371936a26f20410ebb2b3c63
+  data.tar.gz: b706e528ffb67f5e228440fbfb5eb774c73aa642787a21cf61105c382dfc2c82e38cd3431a55f325bf17602acdbfab0b3feefc97b98dad5d62d9d29db3da8dda

data/README.md

@@ -1,5 +1,12 @@
+<!-- vim: set nofoldenable: -->
 # db_subsetter
 
+[![Build Status](https://travis-ci.org/lostapathy/db_subsetter.svg?branch=master)](https://travis-ci.org/lostapathy/db_subsetter)
+[![Maintainability](https://api.codeclimate.com/v1/badges/26b61bf940b79bbfa529/maintainability)](https://codeclimate.com/github/lostapathy/db_subsetter/maintainability)
+[![Test Coverage](https://codeclimate.com/github/lostapathy/db_subsetter/badges/coverage.svg)](https://codeclimate.com/github/lostapathy/db_subsetter/coverage)
+
+![db_subsetter logo](/logo/db_subsetter_logo.png?raw=true "db_subsetter logo")
+
 Extract a subset of a relational database for use in development or testing.  Provides a simple API to filter rows and preserve referential integrity.  The extracted data is packed into a [SQLite](https://www.sqlite.org/) database to allow easy copying.
 
 Developing against a realistic dataset extracted from production provides a lot of advantages over starting with an empty database.  This tools was inspired by [rdbms-subsetter](https://github.com/18F/rdbms-subsetter) and [yaml_db](https://github.com/yamldb/yaml_db/) and combines some of the best attributes of both.
@@ -8,6 +15,21 @@ When working against a legacy database, automatic relationship management does n
 
 ActiveRecord is used for database access, however you *do not* need to have ActiveRecord models for all tables you wish to subset.  Any database supported by ActiveRecord should work.  In theory, you should be able to subset from database and import into another (i.e., MySQL -> Postgres), however in practice this may or may not work well depending on exactly what data types are used.
 
+## RDBMS Support
+
+db_subsetter requires a small RDBMS-specific adapter in order to deal with a few things during the export/import process, mainly related to foreign keys.  At present, the following dialects are supported.  Writing others is pretty straightforward, PRs welcome.
+
+* MySQL
+* MS SQL
+* Postgres
+* Sqlite
+
+## Limitations
+
+Over time we hope to remove some of these limitations.  Until then, tables affected by these limitations can either be skipped or processed manually.
+
+* Tables to be exported must have a single-column primary key unless they have less than SELECT_BATCH_SIZE (5000) rows
+* Foreign keys that do not point back to a primary key are not automatically filtered on
 
 ## Installation
 
@@ -27,7 +49,53 @@ Or install it yourself as:
 
 ## Usage
 
-TODO: Write usage instructions here
+db_subsetter is a toolset for creating export/import scripts to export and import your data.  There is no command to run, rather, you build your own scripts.  These instructions give an overview of how to build up a typical configuration to export a subset of data for typical development workflows, but should just be considered a starting point.
+
+### Prerequisites
+
+The examples provided here assume you are using db_subsetter in the context of a Rails app and that ActiveRecord is already configured and "just works."  This is just done for brevity in the example scripts, as db_subsetter absolutely does not require you to use Rails.  Using Rails just makes some operations a little more convenient.  If you aren't a Rails user, you'll need to add code (after the require statements) to connect ActiveRecord, such as:
+
+```ruby
+ActiveRecord::Base.establish_connection(
+  adapter: "mysql2",
+  host: "127.0.0.1",
+  username: "dbuser",
+  database: "huge_db"
+)
+```
+### A Minimal Start
+
+We'll start our example with a minimal export.rb and build up from there.  This
+
+```ruby
+#!/usr/bin/env ruby
+require 'db_subsetter'
+
+exporter = DbSubsetter::Exporter.new
+filename = "project-#{Rails.env}.sqlite3"
+FileUtils.rm(filename) if File.exists?(filename)
+
+exporter.export(filename)
+```
+Time to run it against our db and see what happens!
+
+
+
+
+TODO: These instructions are a work in progress.  More to come!
+
+## Applications
+
+The obvious application of db_subsetter is to provide a subset for development. There are many other non-obvious uses.
+
+* Capture state when an exception occurs to ease in reproducing the problem
+* Creating reproducible scenarios for complex integration tests
+* Exporting the underlying data used to generate a report, for compliance and audit purposes
+* Archival before deletion of data
+* Providing customers with their own data
+* Migration between RDBMS systems
+
+Come up with something else?  Please file an issue or submit a PR, we'd love to hear about it!
 
 ## Development
 
@@ -35,25 +103,29 @@ After checking out the repo, run `bin/setup` to install dependencies. Then, run
 
 To install this gem onto your local machine, run `bundle exec rake install`. To release a new version, update the version number in `version.rb`, and then run `bundle exec rake release`, which will create a git tag for the version, push git commits and tags, and push the `.gem` file to [rubygems.org](https://rubygems.org).
 
+## Roadmap
+
+* 0.4.x (released) - fully functional, requires manual filtering of all tables
+* 0.5.x (December 2017) - automating filtering of tables by foreign keys, requires much less configuration but will have small breaking API changes
+* 0.6.x (TBA) - improve/expand the scrambler API to allow much simpler filtering of tables, breaking changes to scrambler API likely
+
 ## TODO
 
-* Improve the dialect handling
 * Better example docs on usage and filtering examples
-* Implement a scrubber API to allow sanitizing or correcting data at export time.  This allows us to keep sensitive/personal data out of the export and also allows correction of broken data that won't re-insert.
+* (0.6.0) Implement a scrubber API to allow sanitizing or correcting data at export time.  This allows us to keep sensitive/personal data out of the export and also allows correction of broken data that won't re-insert.
 * Add an executable and/or rake task to perform export and import rather than requiring the API to used directly.  Will need a config file to specific custom plugins
 * Add pre-flight check on import to make sure all tables smell like they will load the data (right columns, at minimum)
-* Finish building and test checks to make sure foreign keys are valid after import
-* Have verify_exportability return all failures together rather than one at a time
-* Add a verbose mode to display more detailed stats while running an export or import (what table we're on, records exported, time taken
+* Examples of validating referential integrity after import
+* Add a verbose mode to display more detailed stats while running an export or import (what table we're on, records exported, time taken)
+* Decouple generating the subset from outputting it, so we could have alternate outputs - like sending direct to another db
+* Provide an alternate API to allow filtering without dealing directly with Arel.  Perhaps a method to pass in an array of IDs to filter from?
+* (0.6.0) Add API calls to allow columns to be skipped completely when subsetting
 
 ## Contributing
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/lostapathy/db_subsetter.
 
-
-
 ## License
 
-
 The gem is available as open source under the terms of the [MIT License](http://opensource.org/licenses/MIT).
 

data/lib/db_subsetter.rb

@@ -1,13 +1,15 @@
-require "db_subsetter/version"
-require "db_subsetter/filter"
-require "db_subsetter/exporter"
-require "db_subsetter/importer"
-require "db_subsetter/scrambler"
-require "db_subsetter/dialect/generic"
-require "db_subsetter/dialect/my_sql"
-require "db_subsetter/dialect/ms_sql"
-
-
-module DbSubsetter
-  # Your code goes here...
-end
+require 'db_subsetter/version'
+require 'db_subsetter/circular_relation_error'
+require 'db_subsetter/filter'
+require 'db_subsetter/database'
+require 'db_subsetter/table'
+require 'db_subsetter/exporter'
+require 'db_subsetter/importer'
+require 'db_subsetter/relation'
+require 'db_subsetter/scrambler'
+require 'db_subsetter/type_helper'
+require 'db_subsetter/dialect/generic'
+require 'db_subsetter/dialect/my_sql'
+require 'db_subsetter/dialect/ms_sql'
+require 'db_subsetter/dialect/postgres'
+require 'db_subsetter/dialect/sqlite'

data/lib/db_subsetter/circular_relation_error.rb

@@ -0,0 +1,4 @@
+module DbSubsetter
+  class CircularRelationError < StandardError
+  end
+end

data/lib/db_subsetter/database.rb

@@ -0,0 +1,46 @@
+module DbSubsetter
+  # A database to be exported from/to
+  class Database
+    def initialize(exporter)
+      @exporter = exporter
+      @tables = {}
+      all_table_names.each { |table_name| @tables[table_name] = Table.new(table_name, self, @exporter) }
+    end
+
+    def find_table(name)
+      @tables[name.to_s]
+    end
+
+    def tables
+      @tables.values
+    end
+
+    def exported_tables
+      tables.reject(&:ignored?)
+    end
+
+    # Raw list of names of all tables in the database.
+    def all_table_names
+      @all_table_names ||= ActiveRecord::Base.connection.tables - ['ar_internal_metadata']
+    end
+
+    # Used in debugging/reporting
+    def total_row_counts
+      tables.map { |table| [table.name, table.total_row_count] }.to_h
+    end
+
+    # Used in debugging/reporting
+    def filtered_row_counts
+      tables.map { |table| [table.name, table.filtered_row_count] }.to_h
+    end
+
+    def exportable?
+      puts "Verifying table exportability ...\n\n" if @exporter.verbose?
+      exported_tables.reject(&:exportable?).count.zero?
+    end
+
+    def exportability_issues
+      exported_tables.reject(&:exportable?).map { |table| [table.name, table.exportability_issues] }.to_h
+    end
+  end
+end

data/lib/db_subsetter/dialect/generic.rb

@@ -1,14 +1,25 @@
 module DbSubsetter
   module Dialect
+    # Dialect to subset to/from database without explicit support
     class Generic
+      INSERT_BATCH_SIZE = 500
+
       def self.import
-        yield
+        ActiveRecord::Base.connection.disable_referential_integrity do
+          yield
+        end
       end
 
       def self.integrity_problems
-        []
+        raise NotImplementedError, 'integrity_problems not implemented for this dialect'
+      end
+
+      def self.truncate_table(table)
+        ActiveRecord::Base.connection.truncate(table)
+      rescue NotImplementedError
+        table = ActiveRecord::Base.connection.quote_table_name(table)
+        ActiveRecord::Base.connection.execute("DELETE FROM #{table}")
       end
     end
   end
 end
-

data/lib/db_subsetter/dialect/ms_sql.rb

@@ -1,19 +1,8 @@
 module DbSubsetter
   module Dialect
+    # Dialect to subset to/from Microsoft SQL Server
     class MSSQL < Generic
-      def self.import
-        ActiveRecord::Base.connection.execute('EXEC sp_msforeachtable "ALTER TABLE ? NOCHECK CONSTRAINT all"')
-        ActiveRecord::Base.connection.execute('EXEC sp_msforeachtable "ALTER TABLE ? DISABLE TRIGGER all"')
-        ActiveRecord::Base.connection.execute("select 'ALTER INDEX ' + I.name + ' ON ' + T.name + ' DISABLE'
-            from sys.indexes I
-            inner join sys.tables T on I.object_id = T.object_id
-            where I.type_desc = 'NONCLUSTERED'
-            and I.name is not null")
-
-        yield
-        ActiveRecord::Base.connection.execute('EXEC sp_msforeachtable "ALTER TABLE ? ENABLE TRIGGER all"')
-        ActiveRecord::Base.connection.execute('EXEC sp_msforeachtable "ALTER TABLE ? WITH CHECK CHECK CONSTRAINT all"')
-      end
+      INSERT_BATCH_SIZE = 100
 
       def self.integrity_problems
         ActiveRecord::Base.connection.execute('EXEC sp_msforeachtable "DBCC CHECKCONSTRAINTS WITH ALL_CONSTRAINTS"')
@@ -21,4 +10,3 @@ module DbSubsetter
     end
   end
 end
-

data/lib/db_subsetter/dialect/my_sql.rb

@@ -1,16 +1,7 @@
 module DbSubsetter
   module Dialect
+    # Dialect to subset to/from MySQL
     class MySQL < Generic
-      def self.import
-        ActiveRecord::Base.connection.execute("SET FOREIGN_KEY_CHECKS=0;")
-        yield
-        ActiveRecord::Base.connection.execute("SET FOREIGN_KEY_CHECKS=1;")
-      end
-
-      def self.integrity_problems
-        raise NotImplementedError.new("integrity_problems not implemented for MySQL")
-      end
     end
   end
 end
-

data/lib/db_subsetter/dialect/postgres.rb

@@ -0,0 +1,7 @@
+module DbSubsetter
+  module Dialect
+    # Dialect to subset to/from postgres
+    class Postgres < Generic
+    end
+  end
+end

data/lib/db_subsetter/dialect/sqlite.rb

@@ -0,0 +1,10 @@
+module DbSubsetter
+  module Dialect
+    # Dialect to subset to/from sqlite
+    class Sqlite < Generic
+      def self.integrity_problems
+        ActiveRecord::Base.connection.execute('PRAGMA foreign_key_check')
+      end
+    end
+  end
+end

data/lib/db_subsetter/exporter.rb

@@ -2,154 +2,89 @@ require 'sqlite3'
 require 'active_record'
 
 module DbSubsetter
+  # Manages exporting a subset of data
   class Exporter
-    attr_writer :filter, :max_unfiltered_rows, :max_filtered_rows
+    attr_writer :max_filtered_rows
+    attr_reader :scramblers, :output, :database
+    attr_accessor :filter, :verbose
+    alias verbose? verbose
 
-    def all_tables
-      ActiveRecord::Base.connection.tables
-    end
-
-    def tables
-      filter.tables
-    end
-
-    def total_row_counts
-      tables.each.map do |table|
-        query = Arel::Table.new(table, ActiveRecord::Base).project("count(1) AS num_rows")
-        rows = ActiveRecord::Base.connection.select_one(query.to_sql)["num_rows"]
-        {table => rows}
-      end
-    end
-
-    def filtered_row_counts
-      tables.each.map do |table|
-        {table => filtered_row_count(table)}
-      end
-    end
-
-    def verify_exportability(verbose = true)
-      puts "Verifying table exportability ...\n\n" if verbose
-      errors = tables.map{|x| verify_table_exportability(x) }.flatten.compact
-      if errors.count > 0
-        puts errors.join("\n")
-        raise ArgumentError.new "Some tables are not exportable"
+    # this is the batch size we insert into sqlite, which seems to be a reasonable balance of speed and memory usage
+    INSERT_BATCH_SIZE = 250
+    SELECT_BATCH_SIZE = 5000
+
+    def export(filename)
+      unless @database.exportable?
+        if verbose?
+          STDERR.puts "\nExportability issues:\n"
+          @database.exportability_issues.each do |table, issues|
+            STDERR.puts table
+            issues.each { |issue| STDERR.puts "\t#{issue}" }
+          end
+        end
+        raise ArgumentError, 'Database is not exportable as filtered!'
       end
-      puts "\n\n" if verbose
-    end
-
-    def export(filename, verbose = true)
-      @verbose = verbose
-      verify_exportability(verbose)
 
       puts "Exporting data...\n\n" if @verbose
       @output = SQLite3::Database.new(filename)
-      @output.execute("CREATE TABLE tables (name TEXT, records_exported INTEGER, columns TEXT)")
-      tables.each do |table|
-        export_table(table)
-      end
+      @output.execute 'CREATE TABLE tables (name TEXT, records_exported INTEGER, columns TEXT)'
+      @database.exported_tables.each(&:export)
     end
 
     def add_scrambler(scrambler)
       @scramblers << scrambler
     end
 
+    def ignore_tables(ignored)
+      limit_tables('ignore!', ignored)
+    end
+
+    def subset_full_tables(full_tables)
+      limit_tables('subset_in_full!', full_tables)
+    end
+
     def initialize
       @scramblers = []
       @page_counts = {}
-    end
-
-    private
-    def max_unfiltered_rows
-      @max_unfiltered_rows || 1000
+      @database = Database.new(self)
+      @filter = Filter.new
+      @verbose = true
+      $stdout.sync
     end
 
     def max_filtered_rows
       @max_filtered_rows || 2000
     end
 
-    # this is the batch size we insert into sqlite, which seems to be a reasonable balance of speed and memory usage
-    def insert_batch_size
-      250
-    end
-
-    def select_batch_size
-      insert_batch_size * 20
-    end
-
-    def filter
-      @filter ||= Filter.new
-      @filter.exporter = self
-      @filter
-    end
-
-    def filtered_row_count(table)
-      query = Arel::Table.new(table, ActiveRecord::Base)
-      query = filter.filter(table, query).project( Arel.sql("count(1)") )
-      ActiveRecord::Base.connection.select_one(query.to_sql).values.first
-    end
-
-    def pages(table)
-      @page_counts[table] ||= ( filtered_row_count(table) / select_batch_size.to_f ).ceil
-    end
-
-    def order_by(table)
-      #TODO should probably allow the user to override this and manually set a sort order?
-      key = ActiveRecord::Base.connection.primary_key(table)
-      key || false
+    # FIXME: look at this API, passing a table name back seems wrong
+    def sanitize_row(table_name, row)
+      row = TypeHelper.cleanup_types(row)
+      scramble_row(table_name, row)
     end
 
-    def verify_table_exportability(table)
-      puts "Verifying: #{table}" if @verbose
-      errors = []
-      errors << "ERROR: Multiple pages but no primary key on: #{table}" if pages(table) > 1 && order_by(table).blank?
-      errors << "ERROR: Too many rows in: #{table} (#{filtered_row_count(table)})" if( filtered_row_count(table) > max_filtered_rows )
-      errors
-    end
-
-    def cleanup_types(row)
-      row.map do |field|
-        case field
-        when Date, Time then field.to_s(:db)
-        else
-          field
-        end
-      end
-    end
+    private
 
-    def scramble_data(table, data)
-      @scramblers.each do |scrambler|
-        data = scrambler.scramble(table, data)
+    def scramble_row(table_name, row)
+      scramblers.each do |scrambler|
+        row = scrambler.scramble(table_name, row)
       end
-      data
+      row
     end
 
-    def export_table(table)
-      print "Exporting: #{table} (#{pages(table)} pages)" if @verbose
-      $stdout.flush if @verbose
-      columns = ActiveRecord::Base.connection.columns(table).map{ |table| table.name }
-      rows_exported = 0
-      @output.execute("CREATE TABLE #{table.underscore} ( data TEXT )")
-      for i in 0..(pages(table) - 1)
-        arel_table = query = Arel::Table.new(table, ActiveRecord::Base)
-        query = filter.filter(table, query)
-        # Need to extend this to take more than the first batch_size records
-        query = query.order(arel_table[order_by(table)]) if order_by(table)
-
-
-        query = query.skip(i * select_batch_size).take(select_batch_size) if pages(table) > 1
-        sql = query.project( Arel.sql('*') ).to_sql
-
-        records = ActiveRecord::Base.connection.select_rows( sql )
-        records.each_slice(insert_batch_size) do |rows|
-          @output.execute("INSERT INTO #{table.underscore} (data) VALUES #{ Array.new(rows.size){"(?)"}.join(",")}", rows.map{|x| scramble_data(table, cleanup_types(x))}.map(&:to_json) )
-          rows_exported += rows.size
+    def limit_tables(operation, apply_to)
+      if apply_to.is_a?(Array)
+        apply_to.each do |t|
+          @database.find_table(t).send(operation)
         end
-        print "." if @verbose
-        $stdout.flush if @verbose
+      elsif apply_to.is_a?(Symbol) || apply_to.is_a?(String)
+        @database.find_table(apply_to).send(operation)
+      elsif apply_to.is_a?(Regexp)
+        @database.tables.each do |table|
+          table.send(operation) if table.name =~ apply_to
+        end
+      else
+        raise ArgumentError, "Don't know how to #{operation} a #{apply_to.class}"
       end
-      puts "" if @verbose
-      @output.execute("INSERT INTO tables VALUES (?, ?, ?)", [table, rows_exported, columns.to_json])
     end
   end
 end
-

data/lib/db_subsetter/filter.rb

@@ -1,26 +1,16 @@
 require 'active_record'
 
 module DbSubsetter
+  # Base class for defining a custom filter for defining how to create a subset
+  # of your database
   class Filter
-    attr_writer :exporter
-
-    def ignore_tables
-      []
-    end
-
-    def tables
-      @exporter.all_tables - ActiveRecord::SchemaDumper.ignore_tables - ignore_tables
-    end
-
-    def filter(table, query)
-      filter_method = "filter_#{table.downcase}"
-      if self.respond_to? filter_method
-        self.send(filter_method, query)
+    def apply(table, query)
+      filter_method = "filter_#{table.name.downcase}"
+      if respond_to? filter_method
+        send(filter_method, query)
       else
         query
       end
     end
-
   end
 end
-

data/lib/db_subsetter/importer.rb

@@ -1,24 +1,36 @@
 require 'sqlite3'
 
 module DbSubsetter
+  # Manages importing a subset of data
   class Importer
-
-    def initialize(filename, dialect = DbSubsetter::Dialect::Generic)
-      raise ArgumentError.new("invalid input file") unless File.exists?(filename)
+    def initialize(filename)
+      raise ArgumentError, 'invalid input file' unless File.exist?(filename)
 
       @data = SQLite3::Database.new(filename)
-      @dialect = dialect
+      @dialect = case ActiveRecord::Base.connection_config[:adapter]
+                 when 'mysql2'
+                   DbSubsetter::Dialect::MySQL
+                 when 'postgresql'
+                   DbSubsetter::Dialect::Postgres
+                 when 'sqlite3'
+                   DbSubsetter::Dialect::Sqlite
+                 when 'sqlserver'
+                   DbSubsetter::Dialect::MSSQL
+                 else
+                   DbSubsetter::Dialect::Generic
+                 end
     end
 
     def tables
       all_tables = []
-      @data.execute("SELECT name FROM tables") do |row|
+      @data.execute('SELECT name FROM tables') do |row|
         all_tables << row[0]
       end
       all_tables
     end
 
-    def import
+    def import(verbose = true)
+      @verbose = verbose
       @dialect.import do
         tables.each do |table|
           import_table(table)
@@ -26,39 +38,37 @@ module DbSubsetter
       end
     end
 
-    def insert_batch_size
-      100 # more like 500 for mysql
-    end
-
     private
+
     def import_table(table)
-      begin
-        ActiveRecord::Base.connection.truncate(table)
-      rescue NotImplementedError
-        ActiveRecord::Base.connection.execute("DELETE FROM #{quoted_table_name(table)}")
-      end
+      $stdout.sync
+      started_at = Time.now
+      print "Importing #{table}" if @verbose
+      @dialect.truncate_table(table)
 
       ActiveRecord::Base.connection.begin_db_transaction
 
       all_rows = @data.execute("SELECT data FROM #{table.underscore}")
-      all_rows.each_slice(insert_batch_size) do |rows|
-        quoted_rows = rows.map{ |row| "(" + quoted_values(row).join(",") + ")" }.join(",")
-        insert_sql = "INSERT INTO #{quoted_table_name(table)} (#{quoted_column_names(table).join(",")}) VALUES #{quoted_rows}"
+      all_rows.each_slice(@dialect::INSERT_BATCH_SIZE) do |rows|
+        quoted_rows = rows.map { |row| '(' + quoted_values(row).join(',') + ')' }.join(',')
+        insert_sql = "INSERT INTO #{quoted_table_name(table)} (#{quoted_column_names(table).join(',')}) VALUES #{quoted_rows}"
         ActiveRecord::Base.connection.execute(insert_sql)
+        print '.' if @verbose
       end
 
-       ActiveRecord::Base.connection.commit_db_transaction
+      ActiveRecord::Base.connection.commit_db_transaction
+      puts " (#{(Time.now - started_at).round(3)}s)" if @verbose
     end
 
     def quoted_values(row)
       out = JSON.parse(row[0])
-      out = out.map{|x| ActiveRecord::Base.connection.type_cast(x, nil) }
-      out = out.map{|x| ActiveRecord::Base.connection.quote(x) }
+      out = out.map { |x| ActiveRecord::Base.connection.type_cast(x, nil) }
+      out = out.map { |x| ActiveRecord::Base.connection.quote(x) }
       out
     end
 
     def columns(table)
-      raw = @data.execute("SELECT columns FROM tables WHERE name = ?", [table]).first[0]
+      raw = @data.execute('SELECT columns FROM tables WHERE name = ?', [table]).first[0]
       JSON.parse(raw)
     end
 
@@ -67,9 +77,7 @@ module DbSubsetter
     end
 
     def quoted_column_names(table)
-      columns(table).map{ |column| ActiveRecord::Base.connection.quote_column_name(column) }
+      columns(table).map { |column| ActiveRecord::Base.connection.quote_column_name(column) }
     end
-
   end
 end
-

data/lib/db_subsetter/relation.rb

@@ -0,0 +1,34 @@
+module DbSubsetter
+  # Wraps a foreign key relationship between two tables
+  class Relation
+    attr_reader :to_table, :column
+
+    def initialize(ar_association, database)
+      @column = ar_association.column
+      @other_column = ar_association.primary_key
+      @to_table = database.find_table ar_association.to_table
+      @from_table = database.find_table ar_association.from_table
+    end
+
+    # We cannot subset automatically if the relation points to a non-primary key
+    def can_subset_from?
+      @to_table.primary_key == @other_column
+    end
+
+    def apply_subset(query)
+      return query if !can_subset_from? || @to_table.subset_in_full?
+
+      # If the other table is ignored, we must not include any records that reference it
+      query = query.where(arel_table[@column].neq(nil)) if @to_table.ignored?
+
+      # If a related table will be exported in full, don't bother subsetting on that key
+      unless @to_table.subset_in_full?
+        other_ids = @to_table.filtered_ids
+        arel_table = @from_table.arel_table
+        conditions = arel_table[@column].in(other_ids).or(arel_table[@column].eq(nil))
+        query = query.where(conditions)
+      end
+      query
+    end
+  end
+end

data/lib/db_subsetter/scrambler.rb

@@ -1,11 +1,12 @@
 require 'random-word'
 
 module DbSubsetter
+  # Clean or redact data to be exported
   class Scrambler
     def scramble(table, row)
       scramble_method = "scramble_#{table.downcase}"
-      if self.respond_to? scramble_method
-        self.send(scramble_method, row)
+      if respond_to? scramble_method
+        send(scramble_method, row)
       else
         row
       end
@@ -16,13 +17,16 @@ module DbSubsetter
     end
 
     protected
+
     def scramble_column(table, column, row_data, value)
       row_data[column_index(table, column)] = value
     end
 
     private
+
     def column_index(table, column)
-      @column_index_cache["#{table}##{column}"] ||= ActiveRecord::Base.connection.columns(table).map{|table| table.name}.index(column.to_s)
+      @column_index_cache["#{table}##{column}"] ||=
+        ActiveRecord::Base.connection.columns(table).map.(&:name).index(column.to_s)
     end
   end
 end

data/lib/db_subsetter/table.rb

@@ -0,0 +1,138 @@
+module DbSubsetter
+  # A database in the database to be subset or imported
+  class Table
+    attr_accessor :name
+
+    def initialize(name, database, exporter)
+      @name = name
+      @exporter = exporter
+      @database = database
+      @full_table = @ignored = false
+    end
+
+    # FIXME: these 4 methods don't feel quite like the correct API yet
+    def ignore!
+      @ignored = true
+    end
+
+    def subset_in_full!
+      @subset_in_full = true
+    end
+
+    def subset_in_full?
+      @subset_in_full
+    end
+
+    def ignored?
+      @ignored
+    end
+
+    def total_row_count
+      query = arel_table.project('count(1) AS num_rows')
+      ActiveRecord::Base.connection.select_one(query.to_sql)['num_rows'].to_i # rails-4.2+pg needs to_i
+    end
+
+    def filtered_row_count
+      query = filtered_records.project(Arel.sql('count(1) AS num_rows'))
+      ActiveRecord::Base.connection.select_one(query.to_sql)['num_rows'].to_i # rails-4.2+pg needs to_i
+    end
+
+    # FIXME: move the raw SQL into another class
+    def export
+      print "Exporting: #{@name} (#{pages} pages)" if verbose
+
+      rows_exported = 0
+      @exporter.output.execute("CREATE TABLE #{@name.underscore} ( data TEXT )")
+      0.upto(pages - 1).each do |page|
+        records_for_page(page).each_slice(Exporter::INSERT_BATCH_SIZE) do |rows|
+          data = rows.map { |x| @exporter.sanitize_row(@name, x) }.map(&:to_json)
+
+          @exporter.output.execute("INSERT INTO #{@name.underscore} (data) VALUES #{Array.new(rows.size) { '(?)' }.join(',')}", data)
+          rows_exported += rows.size
+        end
+
+        print '.' if verbose
+      end
+      puts '' if verbose
+      columns = ActiveRecord::Base.connection.columns(@name).map(&:name)
+      @exporter.output.execute('INSERT INTO tables VALUES (?, ?, ?)', [@name, rows_exported, columns.to_json])
+    end
+
+    def exportable?
+      exportability_issues.empty?
+    end
+
+    def exportability_issues
+      return @exportability_issues if @exportability_issues
+
+      @exportability_issues = []
+      begin
+        puts "Verifying: #{@name} (#{filtered_row_count}/#{total_row_count})" if verbose
+        @exportability_issues << 'Multiple pages but no primary key' if pages > 1 && primary_key.blank?
+        @exportability_issues << "Too many rows (#{filtered_row_count})" if filtered_row_count > @exporter.max_filtered_rows
+      rescue CircularRelationError
+        @exportability_issues << 'Circular relations through this table'
+      end
+      @exportability_issues
+    end
+
+    def filtered_ids
+      return @id_cache if @id_cache
+
+      raise CircularRelationError if @loaded_ids
+      @loaded_ids = true
+
+      sql = filtered_records.project(:id).to_sql
+
+      @id_cache = ActiveRecord::Base.connection.select_rows(sql).flatten
+    end
+
+    def arel_table
+      @arel_table ||= Arel::Table.new(@name)
+    end
+
+    def primary_key
+      ActiveRecord::Base.connection.primary_key(@name)
+    end
+
+    def relations
+      ActiveRecord::Base.connection.foreign_keys(@name).map { |x| Relation.new(x, @database) }
+    end
+
+    private
+
+    def verbose
+      @exporter.verbose?
+    end
+
+    def filtered_records
+      return arel_table if @exporter.nil? || @exporter.filter.nil?
+      query = @exporter.filter.apply(self, arel_table)
+
+      if total_row_count > @exporter.max_filtered_rows
+        query = filter_foreign_keys(query)
+      end
+      query
+    end
+
+    def filter_foreign_keys(query)
+      relations.each do |relation|
+        query = relation.apply_subset(query)
+      end
+      query
+    end
+
+    def records_for_page(page)
+      query = filtered_records
+      query = query.order(arel_table[primary_key]) if primary_key
+
+      query = query.skip(page * Exporter::SELECT_BATCH_SIZE).take(Exporter::SELECT_BATCH_SIZE) if pages > 1
+      sql = query.project(Arel.sql('*')).to_sql
+      ActiveRecord::Base.connection.select_rows(sql)
+    end
+
+    def pages
+      @page_count ||= (filtered_row_count / Exporter::SELECT_BATCH_SIZE.to_f).ceil
+    end
+  end
+end

data/lib/db_subsetter/type_helper.rb

@@ -0,0 +1,15 @@
+module DbSubsetter
+  # Utility module to help safely serialize types
+  # FIXME: nothing about this seems named correctly
+  module TypeHelper
+    def self.cleanup_types(row)
+      row.map do |field|
+        case field
+        when Date, Time then field.to_s(:db)
+        else
+          field
+        end
+      end
+    end
+  end
+end

data/lib/db_subsetter/version.rb

@@ -1,3 +1,3 @@
 module DbSubsetter
-  VERSION = "0.4.1"
+  VERSION = '0.5.0'.freeze
 end

metadata

@@ -1,15 +1,29 @@
 --- !ruby/object:Gem::Specification
 name: db_subsetter
 version: !ruby/object:Gem::Version
-  version: 0.4.1
+  version: 0.5.0
 platform: ruby
 authors:
 - Joe Francis
 autorequire: 
-bindir: exe
+bindir: bin
 cert_chain: []
-date: 2016-10-21 00:00:00.000000000 Z
+date: 2017-12-10 00:00:00.000000000 Z
 dependencies:
+- !ruby/object:Gem::Dependency
+  name: appraisal
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: bundler
   requirement: !ruby/object:Gem::Requirement
@@ -25,49 +39,77 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '1.12'
 - !ruby/object:Gem::Dependency
-  name: rake
+  name: minitest
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '5.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '5.0'
 - !ruby/object:Gem::Dependency
-  name: minitest
+  name: mysql2
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.0'
+        version: 0.4.10
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.0'
+        version: 0.4.10
 - !ruby/object:Gem::Dependency
-  name: activerecord
+  name: pg
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '4.2'
-  type: :runtime
+        version: 0.21.0
+  type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '4.2'
+        version: 0.21.0
 - !ruby/object:Gem::Dependency
-  name: sqlite3
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+- !ruby/object:Gem::Dependency
+  name: activerecord
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 4.2.6
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 4.2.6
+- !ruby/object:Gem::Dependency
+  name: random-word
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
@@ -81,7 +123,7 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '1.3'
 - !ruby/object:Gem::Dependency
-  name: random-word
+  name: sqlite3
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
@@ -101,23 +143,23 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
-- ".gitignore"
-- ".travis.yml"
-- Gemfile
 - LICENSE
 - README.md
-- Rakefile
-- bin/console
-- bin/setup
-- db_subsetter.gemspec
 - lib/db_subsetter.rb
+- lib/db_subsetter/circular_relation_error.rb
+- lib/db_subsetter/database.rb
 - lib/db_subsetter/dialect/generic.rb
 - lib/db_subsetter/dialect/ms_sql.rb
 - lib/db_subsetter/dialect/my_sql.rb
+- lib/db_subsetter/dialect/postgres.rb
+- lib/db_subsetter/dialect/sqlite.rb
 - lib/db_subsetter/exporter.rb
 - lib/db_subsetter/filter.rb
 - lib/db_subsetter/importer.rb
+- lib/db_subsetter/relation.rb
 - lib/db_subsetter/scrambler.rb
+- lib/db_subsetter/table.rb
+- lib/db_subsetter/type_helper.rb
 - lib/db_subsetter/version.rb
 homepage: https://github.com/lostapathy/db_subsetter
 licenses:
@@ -139,10 +181,9 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 2.5.1
+rubygems_version: 2.7.3
 signing_key: 
 specification_version: 4
 summary: Extract a subset of a relational database for use in development or testing.  Provides
   a simple API to filter rows and preserve referential integrity.
 test_files: []
-has_rdoc: 

data/.gitignore

@@ -1,9 +0,0 @@
-/.bundle/
-/.yardoc
-/Gemfile.lock
-/_yardoc/
-/coverage/
-/doc/
-/pkg/
-/spec/reports/
-/tmp/

data/.travis.yml

@@ -1,5 +0,0 @@
-sudo: false
-language: ruby
-rvm:
-  - 2.3.1
-before_install: gem install bundler -v 1.12.3

data/Gemfile

@@ -1,4 +0,0 @@
-source 'https://rubygems.org'
-
-# Specify your gem's dependencies in db_subsetter.gemspec
-gemspec

data/Rakefile

@@ -1,10 +0,0 @@
-require "bundler/gem_tasks"
-require "rake/testtask"
-
-Rake::TestTask.new(:test) do |t|
-  t.libs << "test"
-  t.libs << "lib"
-  t.test_files = FileList['test/**/*_test.rb']
-end
-
-task :default => :test

data/bin/console

@@ -1,14 +0,0 @@
-#!/usr/bin/env ruby
-
-require "bundler/setup"
-require "db_subsetter"
-
-# You can add fixtures and/or initialization code here to make experimenting
-# with your gem easier. You can also use a different console, if you like.
-
-# (If you use this, don't forget to add pry to your Gemfile!)
-# require "pry"
-# Pry.start
-
-require "irb"
-IRB.start

data/bin/setup

@@ -1,8 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-IFS=$'\n\t'
-set -vx
-
-bundle install
-
-# Do any other automated setup that you need to do here

data/db_subsetter.gemspec

@@ -1,32 +0,0 @@
-# coding: utf-8
-lib = File.expand_path('../lib', __FILE__)
-$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
-require 'db_subsetter/version'
-
-Gem::Specification.new do |spec|
-  spec.name          = "db_subsetter"
-  spec.version       = DbSubsetter::VERSION
-  spec.authors       = ["Joe Francis"]
-  spec.email         = ["joe@lostapathy.com"]
-
-  spec.summary       = %q{Extract a subset of a relational database for use in development or
-                          testing.  Provides a simple API to filter rows and preserve referential
-                          integrity.}
-  #spec.description   = %q{TODO: Write a longer description or delete this line.}
-  spec.homepage      = "https://github.com/lostapathy/db_subsetter"
-  spec.license       = "MIT"
-
-  spec.files         = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
-  spec.bindir        = "exe"
-  spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
-  spec.require_paths = ["lib"]
-
-  spec.add_development_dependency "bundler", "~> 1.12"
-  spec.add_development_dependency "rake", "~> 10.0"
-  spec.add_development_dependency "minitest", "~> 5.0"
-
-  spec.add_dependency "activerecord", "~> 4.2"
-  spec.add_dependency "sqlite3", "~> 1.3"
-  spec.add_dependency "random-word", "~> 1.3"
-end
-