pg_shrink 0.0.2

data/lib/pg_shrink/table_filter.rb

@@ -0,0 +1,14 @@
+module PgShrink
+  class TableFilter
+    attr_accessor :table
+    def initialize(table, opts, &block)
+      self.table = table
+      @opts = opts # Currently not used, but who knows
+      @block = block
+    end
+
+    def apply(hash)
+      @block.call(hash)
+    end
+  end
+end

data/lib/pg_shrink/table_sanitizer.rb

@@ -0,0 +1,14 @@
+module PgShrink
+  class TableSanitizer
+    attr_accessor :table
+    def initialize(table, opts, &block)
+      self.table = table
+      @opts = opts # Currently not used, but who knows
+      @block = block
+    end
+
+    def apply(hash)
+      @block.call(hash)
+    end
+  end
+end

data/lib/pg_shrink/version.rb

@@ -0,0 +1,3 @@
+module PgShrink
+  VERSION = "0.0.2"
+end

data/lib/pg_shrink.rb

@@ -0,0 +1,59 @@
+require "uri"
+require "active_support"
+require "active_support/inflector"
+require 'active_support/core_ext/enumerable'
+require 'active_support/core_ext/hash'
+require "pg_shrink/version"
+require "pg_shrink/database"
+require "pg_shrink/database/postgres"
+require "pg_shrink/table_filter"
+require "pg_shrink/table_sanitizer"
+require "pg_shrink/sub_table_operator"
+require "pg_shrink/sub_table_filter"
+require "pg_shrink/sub_table_sanitizer"
+require "pg_shrink/table"
+module PgShrink
+
+  def self.blank_options
+    {
+      url: nil,
+      config: 'Shrinkfile',
+      force: false
+    }
+  end
+
+  def self.valid_pg_url?(url)
+    uri = URI.parse(url)
+    uri.scheme == 'postgres' && !uri.user.blank? && uri.path != '/'
+  rescue => ex
+    false
+  end
+
+  def self.run(options)
+    unless File.exists?(options[:config])
+      raise "Could not find file: #{options[:config]}"
+    end
+
+    unless valid_pg_url?(options[:url])
+      raise "Invalid postgres url: #{options[:url]}"
+    end
+
+    database = Database::Postgres.new(:postgres_url => options[:url])
+
+    database.instance_eval(File.read(options[:config]), options[:config], 1)
+
+    # TODO: Figure out how to write a spec for this.
+    unless options[:force] == true
+      puts 'WARNING:  pg_shrink is destructive!  It will change this database in place.'
+      puts 'Are you sure you want to continue? (y/N)'
+      cont = gets
+      cont = cont.strip
+      unless cont == 'y' || cont == 'Y'
+        puts 'Aborting!'
+        exit 1
+      end
+    end
+
+    database.shrink!
+  end
+end

data/pg_shrink.gemspec

@@ -0,0 +1,34 @@
+# coding: utf-8
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'pg_shrink/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = "pg_shrink"
+  spec.version       = PgShrink::VERSION
+  spec.authors       = ["Kevin Ball"]
+  spec.email         = ["kmball11@gmail.com"]
+  spec.description   = "pg_shrink makes it simple to shrink and sanitize a psql database"
+  spec.summary       = ""
+  spec.homepage      = ""
+  spec.license       = ""
+
+  spec.files         = `git ls-files`.split($/)
+  spec.executables   = spec.files.grep(%r{^bin/}) { |f| File.basename(f) }
+  spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
+  spec.require_paths = ["lib"]
+
+  spec.add_runtime_dependency 'pg'
+  spec.add_runtime_dependency 'activesupport'
+  spec.add_runtime_dependency 'sequel'
+  spec.add_development_dependency "bundler", "~> 1.6"
+  spec.add_development_dependency "rake"
+  spec.add_development_dependency "rspec"
+  spec.add_development_dependency "rspec-nc"
+  spec.add_development_dependency "rspec-mocks"
+  spec.add_development_dependency "guard"
+  spec.add_development_dependency "guard-rspec"
+  spec.add_development_dependency "pry"
+  spec.add_development_dependency "pry-remote"
+  spec.add_development_dependency "pry-nav"
+end

data/spec/Shrinkfile.basic

@@ -0,0 +1,6 @@
+filter_table :users do |f|
+   f.filter_by do |u|
+      u[:name] == "test 1"
+  end
+  f.filter_subtable(:user_preferences, :foreign_key => :user_id)
+end

data/spec/pg_config.yml

@@ -0,0 +1,6 @@
+test:
+  database: test_pg_shrink
+  user: postgres
+  password:
+  host: localhost
+  port:

data/spec/pg_shrink/database/postgres_spec.rb

@@ -0,0 +1,86 @@
+require 'spec_helper'
+require 'pg_spec_helper'
+
+describe PgShrink::Database::Postgres do
+  let(:db) do
+    PgShrink::Database::Postgres.new(PgSpecHelper.pg_config.merge(:batch_size => 5))
+  end
+
+  before(:all) do
+    PgSpecHelper.reset_database
+    PgSpecHelper.create_table(db.connection, :test_table,
+                              {'name' => 'character(128)', 'test' => 'integer'})
+  end
+
+  before(:each) do
+    PgSpecHelper.clear_table(db.connection, :test_table)
+  end
+
+  context "A simple postgres database" do
+    it "sets up a Sequel connection" do
+      expect(db.connection.class).to eq(Sequel::Postgres::Database)
+    end
+
+    context "with 20 simple records" do
+      before(:each) do
+        (1..20).each do |i|
+          db.connection.run(
+            "insert into test_table (name, test) values ('test', #{i})"
+          )
+        end
+      end
+
+      it "can fetch records in batches" do
+        batch_sizes = []
+        db.records_in_batches(:test_table) do |batch|
+          batch_sizes << batch.size
+        end
+        expect(batch_sizes).to eq([5, 5, 5, 5])
+      end
+
+      it "throws an error if records change their primary keys during update" do
+        old_records = db.connection["select * from test_table where test <= 5"].
+          all
+        new_records = old_records.map {|r| r.merge(:id => r[:id] * 10)}
+        expect {db.update_records(:test_table, old_records, new_records)}.
+          to raise_error
+      end
+
+      it "can delete records based on a condition" do
+        db.delete_records(:test_table, {:test => 1..5})
+
+        results = db.connection["select * from test_table"].all
+        expect(results.size).to eq(15)
+        expect(results.map {|r| r[:test]}).to match_array((6..20).to_a)
+      end
+
+      it "can update records" do
+        old_records = db.connection["select * from test_table where test <= 5"].
+          all
+        new_records = old_records.map {|r| r.merge(:test => r[:test] * 10)}
+        db.update_records(:test_table, old_records, new_records)
+        expect(
+          db.connection["select * from test_table where test <= 5"].all.size
+        ).to eq(0)
+        updated_records = db.connection[:test_table].
+          where(:id => old_records.map {|r| r[:id]}).all
+        expect(updated_records.size).to eq(old_records.size)
+        expect(updated_records).to eq(new_records)
+      end
+
+      it "throws an error if you try to delete records in update" do
+        old_records = db.connection["select * from test_table where test <= 5"].
+          all
+        new_records = old_records.first(2)
+        expect {db.update_records(:test_table, old_records, new_records)}.
+          to raise_error
+      end
+
+      it "deletes the whole table" do
+        db.remove_table(:test_table)
+        db.filter!
+        expect(db.connection["select * from test_table"].all.size).to eq(0)
+      end
+    end
+  end
+end

data/spec/pg_shrink/database_spec.rb

@@ -0,0 +1,26 @@
+require 'spec_helper'
+
+describe PgShrink::Database do
+  let(:db) {PgShrink::Database.new}
+
+  it "should yield a table to filter_table" do
+    db.filter_table(:test_table) do |tb|
+      expect(tb.class).to eq(PgShrink::Table)
+      expect(tb.database).to eq(db)
+      expect(tb.table_name).to eq(:test_table)
+    end
+  end
+
+  it "should allow options to set a different primary_key in filter_table" do
+    db.filter_table(:test_table, :primary_key => :foo) do |tb|
+      expect(tb.primary_key).to eq(:foo)
+    end
+  end
+
+  it "should retain options in later invocations" do
+    db.filter_table(:test_table, :primary_key => :foo) {}
+    db.filter_table(:test_table) do |tb|
+      expect(tb.primary_key).to eq(:foo)
+    end
+  end
+end

data/spec/pg_shrink/table_spec.rb

@@ -0,0 +1,158 @@
+require 'spec_helper'
+
+describe PgShrink::Table do
+  context "when a filter is specified" do
+    let(:database) {PgShrink::Database.new}
+    let(:table) { PgShrink::Table.new(database, :test_table) }
+
+    before(:each) do
+      table.filter_by {|test| test[:u] == 1 }
+    end
+
+    it "should add filter to filters array" do
+      expect(table.filters.size).to eq(1)
+    end
+
+    it "should accept values that match the block" do
+      expect(table.filters.first.apply({:u => 1})).to eq(true)
+    end
+
+    it "should reject values that don't match the block" do
+      expect(table.filters.first.apply({:u => 2})).to eq(false)
+    end
+
+    context "when running filters" do
+      it "should return matching subset" do
+        test_data = [{:u => 1}, {:u => 2}]
+        expect(table).to receive(:records_in_batches).and_yield(test_data)
+        expect(table).to receive(:delete_records) do |old_batch, new_batch|
+          expect(old_batch.size).to eq(2)
+          expect(new_batch.size).to eq(1)
+          expect(new_batch.first).to eq({:u => 1})
+        end
+        table.filter!
+      end
+    end
+
+    context "when locked" do
+      before(:each) do
+        table.lock { |test| !!test[:lock] }
+      end
+
+      it "should not filter locked records" do
+        test_data = [{:u => 1, :lock => false},
+                     {:u => 2, :lock => false},
+                     {:u => 2, :lock => true}]
+        allow(table).to receive(:records_in_batches).and_yield(test_data)
+        allow(table).to receive(:delete_records) do |old_batch, new_batch|
+          expect(old_batch.size).to eq(3)
+          expect(new_batch.size).to eq(2)
+          expect(new_batch).
+            to eq([{:u => 1, :lock => false}, {:u => 2, :lock => true}])
+        end
+        table.filter!
+      end
+    end
+  end
+
+  context "when a sanitizer is specified" do
+    let(:database) {PgShrink::Database.new}
+    let(:table) { PgShrink::Table.new(database, :test_table) }
+    before(:each) do
+      table.sanitize do |test|
+        test[:u] = -test[:u]
+        test
+      end
+    end
+
+    it "should add sanitizer to sanitizers array" do
+      expect(table.sanitizers.size).to eq(1)
+    end
+
+    it "should alter values based on the block" do
+      expect(table.sanitizers.first.apply({:u => 1})).to eq({:u => -1})
+    end
+
+    context "when running sanitizers" do
+      it "returns an altered set of records" do
+        test_data = [{:u => 1}, {:u => 2}]
+        expect(table).to receive(:records_in_batches).and_yield(test_data)
+        expect(table).to receive(:update_records) do |old_batch, new_batch|
+          expect(old_batch).to eq(test_data)
+          expect(new_batch).to eq([{:u => -1}, {:u => -2}])
+        end
+        table.sanitize!
+      end
+    end
+  end
+
+  context "when a subtable filter is specified" do
+    let(:database) {PgShrink::Database.new}
+    let(:table) { PgShrink::Table.new(database, :test_table) }
+
+    before(:each) do
+      table.filter_subtable(:subtable)
+    end
+
+    it "yields back a table so additional manipulations can be made" do
+      table.filter_subtable(:subtable) do |f|
+        expect(f.class).to eq(PgShrink::Table)
+        expect(f.table_name).to eq(:subtable)
+      end
+    end
+
+    it "adds subtable_filter to subtable_filters array" do
+      expect(table.subtable_filters.size).to eq(1)
+    end
+
+    describe "when running filters" do
+      before(:each) do
+        table.filter_by do |test|
+          !!test[:u]
+        end
+      end
+
+      it "runs subtable filters with old and new batches" do
+        test_data = [{:u => true}, {:u => false}]
+        expect(table).to receive(:records_in_batches).and_yield(test_data)
+        expect(database).to receive(:delete_records)
+        expect(table).to receive(:filter_subtables) do |old_batch, new_batch|
+          expect(old_batch).to eq(test_data)
+          expect(new_batch).to eq([{:u => true}])
+        end
+        table.filter!
+      end
+    end
+  end
+  context "when a remove is specified" do
+    let(:database) {PgShrink::Database.new}
+    let(:table) { PgShrink::Table.new(database, :test_table) }
+    let(:test_data) {[{:u => 1}, {:u => 2}]}
+
+    before(:each) do
+      table.mark_for_removal!
+    end
+
+    it "should by default remove all" do
+      expect(table).to receive(:records_in_batches).and_yield(test_data)
+      expect(table).to receive(:delete_records) do |old_batch, new_batch|
+        expect(old_batch).to eq(test_data)
+        expect(new_batch).to eq([])
+      end
+      table.shrink!
+    end
+
+    it "should allow locking of records" do
+      table.lock do |u|
+        u[:u] == 1
+      end
+      expect(table).to receive(:records_in_batches).and_yield(test_data)
+      expect(table).to receive(:delete_records) do |old_batch, new_batch|
+        expect(old_batch).to eq(test_data)
+        expect(new_batch).to eq([{:u => 1}])
+      end
+      table.shrink!
+    end
+
+  end
+end