janko 0.0.2

data/config/environment.rb

@@ -0,0 +1,7 @@
+require "active_record"
+require "pg"
+require "logger"
+
+ActiveRecord::Base.logger = Logger.new("tmp/test.log")
+config = YAML::load(IO.read("config/database.yml"))
+ActiveRecord::Base.establish_connection(config["development"])

data/janko.gemspec

@@ -0,0 +1,32 @@
+# coding: utf-8
+
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+
+require 'janko/version'
+
+Gem::Specification.new do |spec|
+    spec.name = "janko"
+    spec.version = Janko::VERSION
+    spec.authors = ["Don Werve"]
+    spec.email = ["don@werve.net"]
+    spec.summary = %q{High-performance import, merge, and upsert for PostgreSQL.}
+    spec.description = %q{Because sometimes you just need to feed PostgreSQL a lot of data.}
+    spec.homepage = "https://github.com/matadon/janko"
+    spec.license = "Apache-2.0"
+    spec.files = `git ls-files -z`.split("\x0")
+    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', '~> 0.17', '> 0.17'
+    spec.add_runtime_dependency 'agrippa', '~> 0.0.1', '>= 0.0.1'
+    spec.add_development_dependency "bundler", "~> 1.6"
+    spec.add_development_dependency "rake", "~> 10.0"
+    spec.add_development_dependency 'rspec', '~> 3.0', '>= 3.0.0'
+    spec.add_development_dependency 'guard', '~> 2.8'
+    spec.add_development_dependency 'guard-rspec', '~> 4.3'
+    spec.add_development_dependency 'ruby_gntp', '~> 0'
+    spec.add_development_dependency 'simplecov', '~> 0'
+    spec.add_development_dependency 'pry', '~> 0'
+    spec.add_development_dependency 'activerecord', '~> 4.1', '> 4.1'
+end

data/lib/janko.rb

@@ -0,0 +1,5 @@
+require "janko/version"
+
+module Janko
+  # Your code goes here...
+end

data/lib/janko/bulk_merge.rb

@@ -0,0 +1,57 @@
+require "janko/merge_result"
+require "janko/import"
+require "janko/upsert"
+
+module Janko
+    class BulkMerge
+        def initialize(options = {})
+            @options = options
+            @target = "merge_#{SecureRandom.hex(8)}"
+            @upsert = Upsert.new(options.merge(from_table: @target))
+            @importer = Import.new(strategy: Janko::CopyImporter,
+                table: @target, connection: connection,
+                columns: options[:columns])
+        end
+
+        def start
+            create_copy_target
+            @importer.start
+            self
+        end
+
+        def push(*values)
+            @importer.push(*values)
+            self
+        end
+
+        def stop
+            @importer.stop
+            @upsert.process.cleanup
+            drop_copy_target
+            self
+        end
+
+        def result
+            @upsert.result
+        end
+
+        def connection
+            @options[:connection]
+        end
+
+        private
+
+        def create_copy_target
+            connection.exec(<<-END)
+                CREATE TEMP TABLE #{@target} WITHOUT OIDS ON COMMIT DROP
+                AS (SELECT * FROM #{@options[:table]}) WITH NO DATA;
+            END
+            self
+        end
+
+        def drop_copy_target
+            connection.exec("DROP TABLE #{@target}")
+            self
+        end
+    end
+end

data/lib/janko/column_list.rb

@@ -0,0 +1,193 @@
+require "janko/tagged_column"
+require "janko/constants"
+require "agrippa/mutable_hash"
+
+module Janko
+    class ColumnList
+        include Enumerable
+
+        include Agrippa::MutableHash
+
+        def ColumnList.build(source)
+            source.is_a?(ColumnList) ? source : ColumnList.new.add(*source)
+        end
+
+        def default_state
+            { columns: {} }
+        end
+
+        def builder
+            self
+        end
+
+        def add(*names)
+            each_column(names) { |name| add_column(name) }
+        end
+
+        def remove(*names)
+            each_column(names) { |name| remove_column(name) }
+        end
+
+        def tag(tag, *names)
+            each { |_, column| column.untag(tag) }
+            each_column(names) { |name| add_column(name).tag(tag) }
+        end
+
+        def untag(tag, *names)
+            each_column(names) { |_, column| column and column.untag(tag) }
+        end
+
+        def tagged(tag = nil)
+            filter_columns do |name, column|
+                (tag.nil? and column.tagged?) or column.has_tag?(tag)
+            end
+        end
+
+        def not_tagged(tag)
+            filter_columns { |name, column| not column.has_tag?(tag) }
+        end
+
+        def none_tagged?(tag)
+            @state[:columns].none? { |_, column| column.has_tag?(tag) }
+        end
+
+        def alter(*names)
+            each_column(names) do |name, column|
+                raise("Unknown column: #{name}") unless column
+                yield(column)
+            end
+        end
+
+        def set(state)
+            chain(state)
+        end
+
+        def pack(values)
+            pack_hash(stringify_keys(values))
+        end
+
+        def to_list
+            map_and_join { |_, column| column.quoted }
+        end
+
+        def to_conditions(left = nil, right = nil)
+            map_and_join(" AND ") { |_, c| c.to_condition(left, right) }
+        end
+
+        def to_setters(left = nil, right = nil)
+            map_and_join { |_, c| c.to_setter(left, right) }
+        end
+
+        def to_list_with_defaults
+            map_and_join { |_, c| c.to_value }
+        end
+
+        def to_binds
+            map_and_join_with_index { |_, c, i| c.to_bind(i + 1) }
+        end
+
+        def to_typecast_binds
+            map_and_join_with_index { |_, c, i| c.to_typecast_bind(i + 1) }
+        end
+
+        def each(&block)
+            @state[:columns].each(&block)
+        end
+
+        def map(&block)
+            @state[:columns].map(&block)
+        end
+
+        def map_and_join(separator = ",", &block)
+            map(&block).join(separator)
+        end
+
+        def map_and_join_with_index(separator = ",")
+            output = each_with_index.map { |pair, index| yield(*pair, index) }
+            output.join(separator)
+        end
+
+        def empty?
+            @state[:columns].empty?
+        end
+
+        def columns
+            @state[:columns].keys
+        end
+
+        def connection
+            @state[:parent].connection
+        end
+
+        def table
+            @state[:parent].table
+        end
+
+        def inspect
+            children = @state[:columns].map { |name, column|
+                "#{name}(#{column.inspect})" }
+            "#<#{self.class}:0x#{self.__id__.to_s(16)} #{children.join(" ")}>"
+        end
+
+        def length
+            @length ||= @state[:columns].length
+        end
+
+        def pack_hash(values)
+            result = columns.map { |column| values.delete(column) }
+            return(result) if values.empty?
+            raise(ArgumentError, "Unknown columns: #{values.keys.join(" ")}")
+        end
+
+        private
+
+        def add_column(name)
+            @state[:columns][name.to_s] ||= TaggedColumn.new(name: name,
+                parent: self)
+        end
+
+        def remove_column(name)
+            @state[:columns].delete(name.to_s)
+            self
+        end
+
+        def filter_columns
+            result = {}
+            each { |n, c| result[n] = c if yield(n, c) }
+            self.class.new(@state.merge(columns: result))
+        end
+
+        def stringify_keys(hash)
+            output = {}
+            hash.each { |k, v| output[k.to_s] = v }
+            output
+        end
+
+        def matching_names(names)
+            result = []
+            names.flatten.each do |name|
+                if(name == Janko::ALL)
+                    result.concat(connection.column_list(table))
+                elsif(name == Janko::DEFAULT)
+                    result.concat(connection.column_list(table) - [ "id" ])
+                elsif(name.is_a?(Hash) and name.has_key?(:except))
+                    all_columns = connection.column_list(table)
+                    exceptions = [ name[:except] ].flatten.map(&:to_s)
+                    result.concat(all_columns - exceptions)
+                elsif(name.nil? or name == "")
+                    raise("Blank or nil column names are not allowed.")
+                else
+                    result.push(name.to_s)
+                end
+            end
+            result
+        end
+ 
+        def each_column(names)
+            matching_names(names).each do |name|
+                yield(name, @state[:columns][name.to_s])
+            end
+            self
+        end
+    end
+end

data/lib/janko/connection.rb

@@ -0,0 +1,107 @@
+require "agrippa/maybe"
+require "agrippa/delegation"
+
+module Janko
+    class Connection
+        include Agrippa::Delegation
+
+        def Connection.build(backend)
+            return(default) if backend.nil?
+            return(backend) if backend.is_a?(Connection)
+            new(backend)
+        end
+
+        def Connection.default
+            if Kernel.const_defined?("ActiveRecord::Base")
+                new(Kernel.const_get("ActiveRecord::Base"))
+            else
+                raise("No default connection available.")
+            end
+        end
+
+        def Connection.cache_catalog
+            @catalog ||= yield
+        end
+
+        def Connection.reset_cached_catalog
+            @catalog = nil
+            self
+        end
+
+        attr_reader :backend
+
+        delegate *%w(exec prepare exec_prepared async_exec put_copy_data
+            put_copy_end get_last_result), to: "backend"
+
+        def initialize(backend)
+            @backend = extract_raw_connection(backend)
+        end
+
+        def in_transaction?
+            backend.transaction_status > 0
+        end
+
+        def failed?
+            backend.transaction_status >= 3
+        end
+
+        # http://dba.stackexchange.com/questions/22362/
+        # http://www.postgresql.org/docs/9.3/static/catalog-pg-attribute.html
+        # http://www.postgresql.org/docs/9.3/static/catalog-pg-attrdef.html
+        def catalog
+            Connection.cache_catalog do
+                result = backend.exec(<<-END)
+                    SELECT relname, attname, typname, pg_get_expr(adbin, 0)
+                    FROM pg_class
+                    LEFT JOIN pg_namespace ON (
+                        pg_class.relnamespace = pg_namespace.oid)
+                    LEFT JOIN pg_attribute ON (
+                        pg_class.oid = pg_attribute.attrelid)
+                    LEFT JOIN pg_attrdef ON (
+                        pg_attribute.attrelid = pg_attrdef.adrelid
+                        AND pg_attribute.attnum = pg_attrdef.adnum)
+                    LEFT JOIN pg_type ON (
+                        pg_attribute.atttypid = pg_type.oid)
+                    WHERE pg_class.relkind IN ('r','')
+                        AND pg_namespace.nspname
+                            NOT IN ('pg_catalog', 'pg_toast')
+                        AND pg_table_is_visible(pg_class.oid)
+                        AND attnum > 0
+                        AND NOT attisdropped;
+                END
+
+                output = {}
+                result.each_row do |row|
+                    output[row[0]] ||= {}
+                    output[row[0]][row[1]] = { type: row[2], default: row[3] }
+                end
+                output
+            end
+        end
+
+        def column_list(table)
+            catalog[table].keys
+        end
+
+        def column_type(table, column)
+            catalog[table][column][:type]
+        end
+
+        def column_default(table, column)
+            catalog[table][column][:default]
+        end
+
+        private
+
+        def maybe(*args)
+            Agrippa::Maybe.new(*args)
+        end
+
+        def extract_raw_connection(backend)
+            return(backend) if backend.is_a?(PG::Connection)
+            maybe(backend).raw_connection._ \
+                or maybe(backend).connection.raw_connection._ \
+                or raise("Unable to extract a connection from: #{backend}")
+        end
+    end
+end

data/lib/janko/constants.rb

@@ -0,0 +1,13 @@
+require "janko/flag"
+
+module Janko
+    module Constants
+        NULL = nil
+
+        DEFAULT = Flag.new(2 ** 0)
+
+        ALL = Flag.new(2 ** 1)
+
+        KEEP = Flag.new(2 ** 2)
+    end
+end

data/lib/janko/copy_importer.rb

@@ -0,0 +1,36 @@
+require "csv"
+require "agrippa/mutable"
+
+module Janko
+    class CopyImporter
+        include Agrippa::Mutable
+
+        state_reader :connection, :table, :columns
+
+        def start
+            connection.async_exec(sprintf("COPY %s(%s) FROM STDOUT CSV",
+                table, columns.to_list))
+            self
+        end
+
+        def push(values)
+            begin
+                line = CSV.generate_line(columns.pack(values))
+                connection.put_copy_data(line)
+            rescue
+                stop
+                raise
+            end
+            self
+        end
+
+        def stop
+            connection.put_copy_end
+            result = connection.get_last_result
+            return(self) if (result.result_status == PG::PGRES_COMMAND_OK)
+            return(self) if (result.result_status == PG::PGRES_COPY_IN)
+            raise(PG::Error, result.error_message)
+            self
+        end
+    end
+end