theman 0.0.5 → 0.1.0

data/.rspec

@@ -0,0 +1 @@
+--colour

data/README.rdoc

@@ -2,39 +2,49 @@
 
 The man getting you down?
 
-FasterCSV is great but when you get to 100MB files it takes too long and 
-you may only be looking for certain records that match some criteria, 
-enter Theman.
+Theman lets you import lots of data into postgres very fast.
 
 == Installation
 
   gem install 'theman'
 
-== Basic Usage
+== Basic usage
 
-  my_agent = ::Theman::Agency.new 'pretty.csv'
-  my_agent.instance.count
+Say we have a csv file called <tt>sample.csv</tt> with 220 rows:
 
-Say if you wanted to analyse some data the man has given you, you could make the 
-table persist by doing the following in an IRB session:
+  conn  = PGconn.open(:dbname => 'test')
+  
+  agent = Theman::Agency.new(conn, 'sample.csv')
+  agent.create!
 
-  require 'active_record'
-  require 'theman'
+  res = conn.exec("SELECT count(*) FROM #{agent.table_name}")
+  res.getvalue(0,0)
 
-  ActiveRecord::Base.establish_connection :database => $USER, :host => "localhost", :adapter => "postgresql"
+  => 220
 
-  my_agent = ::Theman::Agency.new('pretty.csv', ActiveRecord::Base, :temporary => false) {|a| a.date :date_col }
+== Basic usage with Active Record and a simple object
+  
+  conn  = ActiveRecord::Base.connection.raw_connection
+  
+  agent = Theman::Agency.new(conn, 'sample.csv')
+  agent.create!
 
-After you have figured out what the hell is going on with the data, you can get the man of 
-your back (for now).
+  model = Theman::Object(agent.table_name, ActiveRecord::Base)
+  model.count
 
-== Advanced Usage
-    
-  my_agent = ::Theman::Agency.new 'ugly.csv', ActiveRecord::Base, :primary_key => true do |s|
-    s.nulls /"N"/, /"UNKNOWN"/, /""/
-    s.seds "-n -e :a -e '1,15!{P;N;D;};N;ba'"
-    s.delimiter "|"
-    s.table do |t|
+  => 220
+
+== Advanced usage with Active Record and an existing model
+
+Theman will call the +create!+ method if you pass in a block.
+
+  conn  = ActiveRecord::Base.connection.raw_connection
+
+  agent = Theman::Agency.new conn, 'ugly.csv' |ag|
+    ag.nulls /"N"/, /"UNKNOWN"/, /""/
+    ag.seds "-n -e :a -e '1,15!{P;N;D;};N;ba'"
+    ag.delimiter "|"
+    ag.table do |t|
       t.string :name, :limit => 50
       t.date :date
       t.integer :ext_id
@@ -42,8 +52,11 @@ your back (for now).
       t.boolean :exited
     end
   end
-
-  my_agent.instance.where(:exited => true).count
+  
+  MyModel.table_name = agent.table_name
+  MyModel.where(:exited => true).count
+  
+  => 220
 
 In the above example we omitted the last 15 rows, made some things null, 
 added a primary key and changed some column data types to something else.
@@ -52,27 +65,32 @@ If you do not provide a table block your columns will be VARCHAR(255); you
 can cherry pick the columns that you want to change the data types for.
 
 The temp table has no id column by default, but you can add one using the options 
-hash or calling add_primary_key, this will add the agents_pkey column.
+hash or calling +add_primary_key+, this will add the agents_pkey column.
 
-If you want to call this procedural style just don't pass in the path to the file 
-and Theman will not create a table, in this case you will need to call everything 
-explicitly:
+WARNING: if you have user input in your sed commands, don't.
 
-  smith = ::Theman::Agency.new
-  smith.stream 'real_ugly.csv'
-  smith.create_table
-  smith.pipe_it
+== Drop on commit
 
-WARNING: if you have user input in your sed commands, don't.
+If you want to use <tt>ON COMMIT DROP</tt> you will need to pass in 
+<tt>:on_commit => :drop</tt> into options and do everthing inside a transacton.
+
+    agent = Theman::Agency.new conn, csv, :on_commit => :drop
+
+    agent.transaction do
+      agent.create!
+      # do stuff
+    end
+
+== No headers
 
-NOTE: When passing in extra sed commands each one will be given it's own subprocess, 
-this makes it much faster if you have more than one core.
+If you data does not have headers pass into options <tt>:headers => false</tt>, but 
+each column must be specified or the import will fail.
 
 == Dates
 
 Ah dates, the joy! Use datestyle to tell Theman to then tell PostgreSQL:
 
-  my_agent = ::Theman::Agency.new 'uber_foie_gras.csv' do |schmit|
+  agent = Theman::Agency.new conn, 'uber_foie_gras.csv' do |schmit|
     schmit.datestyle 'European'
     schmit.table do |t|
       t.date :start_date
@@ -83,34 +101,36 @@ Ah dates, the joy! Use datestyle to tell Theman to then tell PostgreSQL:
 Refer to PostgreSQL docs for more info in the mean time here is some 
 copy and paste action:
 
-=== ISO
+<b>ISO</b>
 
 Use ISO 8601-style dates and times (YYYY-MM-DD HH:MM:SS). This is the default. 
 
-=== SQL
+<b>SQL</b>
 
 Use Oracle/Ingres-style dates and times. 
 
-=== PostgreSQL
+<b>PostgreSQL</b>
 
 Use traditional PostgreSQL format. 
 
-=== German
+<b>German</b>
 
 dd.mm.yyyy
 
-=== European
+<b>European</b>
 
 dd/mm/yyyy
 
-=== US
+<b>US</b>
 
 mm/dd/yyyy
 
-== Troubles
+== My table is empty?
 
-Table empty? the man (the real life one) has given you crappy data and 
-PostgresSQL has silently dissed it.
+PostgreSQL <tt>COPY</tt> requires that the data be well formed, any rows that 
+are different to what is expected by the table and the whole import will fail. 
+If you are importing very large files and the import fails space on disc will still
+be used untill <tt>VACUUM</tt>.
 
 == Copyright
 

data/lib/theman/agency/columns.rb

@@ -0,0 +1,95 @@
+module Theman
+  class Agency
+    class Columns
+      attr_accessor :column
+      attr_reader :connection
+
+      def initialize(conn)
+        @connection = conn
+        @columns    = []
+      end
+
+      def to_sql #:nodoc
+        @columns.map{|column| column_to_sql(*column)}.join(', ')
+      end
+
+      %w( string text integer float decimal datetime timestamp time date binary boolean ).each do |type|
+        class_eval <<-EOV, __FILE__, __LINE__ + 1
+          def #{type}(name, *args)
+            column(name, '#{type}', *args)
+          end
+        EOV
+      end
+      
+      def symbolize(name) #:nodoc
+        name.is_a?(Symbol) ? name : name.gsub(/ /,"_").gsub(/\W/, "").downcase.to_sym
+      end
+
+      def column(name, type, *args) #:nodoc
+        sym_col = symbolize(name)
+        @columns.each_with_index do |column, index|
+          if column[0] == sym_col
+            @columns[index] = [sym_col, type, *args]
+            return
+          end
+        end
+        @columns << [sym_col, type, *args]
+      end
+      
+      def column_to_sql(name, type, options = {}) #:nodoc
+        sql = [quote_column_name(name)]
+        case type
+        when 'integer'
+          if options[:limit]
+            case options[:limit]
+            when 1, 2;
+              sql << 'smallint'
+            when 3, 4;
+              sql << 'integer'
+            when 5..8;
+              sql << 'bigint'
+            else
+              raise ArgumentError, "No integer type has byte size #{limit}."
+            end
+          else
+            sql << 'integer'
+          end
+        when 'decimal'
+          sql << 'double precision'
+        when 'float'
+          sql << 'double precision'
+        when 'string'
+          if options[:limit]
+            sql << "character varying(#{options[:limit]})"
+          else
+            sql << 'character varying(255)'
+          end
+        when 'binary'
+          sql << 'oid'
+        when 'time'
+          sql << 'time without time zone'
+        when 'datetime'
+          sql << 'timestamp without time zone'
+        when 'timestamp'
+          sql << 'timestamp without time zone'
+        else
+          sql << type
+        end
+
+        if options[:null] ==  false
+          sql << 'NOT NULL'
+        end
+
+        if options[:default]
+          sql << "DEFAULT #{options[:default]}"
+        end
+
+        sql.join(' ')
+      end
+      
+      def quote_column_name(name) #:nodoc
+        @connection.quote_ident(name.to_s)
+      end
+    end
+  end
+end

data/lib/theman/agency/table.rb

@@ -0,0 +1,24 @@
+module Theman
+  class Agency
+    class Table
+      def initialize(name, columns, temporary = nil, on_commit = nil)
+        @name       = name
+        @columns    = columns
+        @temporary  = temporary
+        @on_commit  = on_commit
+      end
+
+      def to_sql(sql = []) #:nodoc
+        sql << ["CREATE"]
+        sql << "TEMPORARY TABLE" unless @temporary == false
+        sql << @name
+        sql << "(#{@columns})"
+        unless @on_commit.nil?
+          sql << "ON COMMIT"
+          sql << @on_commit.to_s.upcase.gsub(/_/," ")
+        end
+        sql.join(" ")
+      end
+    end
+  end
+end

data/lib/theman/agency.rb

@@ -1,129 +1,147 @@
 module Theman
   class Agency
-    attr_reader :instance, :column_names
-
-    def initialize(stream = nil, parent = ::ActiveRecord::Base, options = {})
-      @options = options
-      @stream  = stream
-
-      agent_id = sprintf "agent%010d", rand(100000000)
-      @column_names = {}
-      @instance = Class.new(parent) do
-        instance_eval <<-EOV, __FILE__, __LINE__ + 1
-          set_table_name "#{agent_id}"
-          def table_name
-            "#{agent_id}"
-          end
-          def inspect
-            "Agent (#{agent_id})"
-          end
-        EOV
-      end
-
-      yield self if block_given?
-      return unless stream
-      create_table
-      pipe_it
-      if @options[:primary_key]
-        add_primary_key
+    attr_reader :columns, :table_name, :connection
+    # create a new agent object - if a block is passed create! is called
+    #
+    # ==== Parameters
+    # * +conn+ - A database connection from the <tt>PGconn</tt> class 
+    #   or <tt>ActiveRecord::Base.connection.raw_connection</tt> which 
+    #   is the same class.
+    # * +stream+ - path to the data file.
+    # * +options+ - Additional options are <tt>:temporary</tt>, 
+    #   <tt>:on_commit</tt> and <tt>:headers</tt>
+    #
+    # ==== Examples
+    #   # Update all customers with the given attributes
+    #   conn  = PGconn.open(:dbname => 'test')
+    #   agent = Theman::Agency.new(conn, 'sample.csv')
+    #   agent.create!
+    #   res = conn.exec("SELECT count(*) FROM #{agent.table_name}")
+    #   res.getvalue(0,0)
+    def initialize(conn, stream, options = {}, &block)
+      @stream       = stream
+      @connection   = conn
+      @options      = options
+
+      @table_name         = sprintf "agent%010d", rand(100000000)
+      @columns            = Columns.new(conn)
+      @stream_columns_set = false
+
+      if block_given?
+        yield self
+        create!
       end
     end
-
-    def table
-      yield self if block_given?
+    
+    # create a transaction block for use with :on_commit => :drop
+    def transaction(&block)
+      connection.exec "BEGIN;"
+      yield
+      connection.exec "COMMIT;"
     end
 
-    # columnn data type methods
-    %w( string text integer float decimal datetime timestamp time date binary boolean ).each do |column_type|
-      class_eval <<-EOV, __FILE__, __LINE__ + 1
-        def #{column_type}(column_name, *args)
-          column(column_name, '#{column_type}', *args)
-        end
-      EOV
+    def create_stream_columns #:nodoc
+      @stream_columns_set = true
+      headers.split(delimiter_regexp).each do |column|
+        @columns.string column
+      end
     end
 
-    # overides the default string type column
-    def column(column_name, column_type, *args)
-      @column_names.merge! column_name.to_sym => [column_name, column_type, *args]
+    def headers #:nodoc
+      File.open(@stream, "r"){ |infile| infile.gets }
     end
-
+    
+    # create default columns from stream and replace selected
+    # columns with custom data types from block
+    def table(&block)
+      create_stream_columns unless @options[:headers] == false
+      yield @columns
+    end
+    
+    # the location of the data to be sent to Postgres via STDIN (requires a header row)
     def stream(arg)
       @stream = arg
     end
-
+    
+    # datestyle of date columns
     def datestyle(arg)
       @datestyle = arg
     end
-
+    
+    # values in stream to replace with NULL
     def nulls(*args)
       @nulls = args
     end
 
+    # custom seds to parse stream with
     def seds(*args)
       @seds = args
     end
 
+    # delimter used in stream - comma is the default
     def delimiter(arg)
       @delimiter = arg
     end
     
-    def symbolize(name)
-      name.gsub(/ /,"_").gsub(/\W/, "").downcase.to_sym
-    end
-    
-    def psql_copy(psql = [])
-      psql << "COPY #{@instance.table_name} FROM STDIN WITH"
+    def psql_copy(psql = []) #:nodoc
+      psql << "COPY #{table_name} FROM STDIN WITH"
       psql << "DELIMITER '#{@delimiter}'" unless @delimiter.nil?
-      psql << "CSV HEADER"
+      psql << "CSV"
+      psql << "HEADER" unless @options[:headers] == false
       psql
     end
 
-    def psql_command(psql = [])
+    def psql_command(psql = []) #:nodoc
       psql << "SET DATESTYLE TO #{@datestyle}" unless @datestyle.nil?
       psql << psql_copy.join(" ")
       psql
     end
 
-    def sed_command(sed = [])
+    def sed_command(sed = []) #:nodoc
       sed << nulls_to_sed unless @nulls.nil?
       sed << @seds unless @seds.nil?
       sed
     end
 
-    def nulls_to_sed
+    def nulls_to_sed #:nodoc
       @nulls.map do |regex|
         "-e 's/#{regex.source}//g'"
       end
     end
 
-    # creates a delimiter regular expresion
-    def delimiter_regexp
-      Regexp.new(@delimiter.nil? ? "," : "\\#{@delimiter}")
+    def delimiter_regexp #:nodoc
+      @delimiter_regexp ||= Regexp.new(@delimiter.nil? ? "," : "\\#{@delimiter}")
     end
-
-    # read the first line from the stream to create a table with
-    def create_table
-      f = File.open(@stream, 'r')
-      options = {:id => false}
-      options.merge!(:temporary => true) if @options[:temporary].nil?
-      instance.connection.create_table(instance.table_name, options) do |t|
-        f.each_line do |line|
-          line.split(delimiter_regexp).each do |col|
-            column_name = symbolize(col)
-            if custom = @column_names.fetch(column_name, nil)
-              t.column(*custom) 
-            else
-              t.string column_name
-            end
-          end
-          break
-        end
+    
+    # Postgress COPY command using STDIN
+    # - reads chunks of 8192 bytes to save memory
+    # System command for IO subprocesses are piped to 
+    # take advantage of multi cores
+    def create!
+      unless @stream_columns_set || @options[:headers] == false
+        create_stream_columns
       end
+      connection.exec Table.new(table_name, @columns.to_sql, @options[:temporary], @options[:on_commit]).to_sql
+      pipe_it
     end
     
-    # system command for IO subprocesses, commands are piped to 
-    # take advantage of multi cores
-    def system_command
+    # adds a serial column called agents_pkey and sets as primary key
+    def add_primary_key!
+      connection.exec "ALTER TABLE #{table_name} ADD COLUMN agents_pkey serial PRIMARY KEY;"
+    end
+    
+    # analyzes the table for efficent query contstruction on tables larger than ~1000 tuples
+    def analyze!
+      connection.exec "ANALYZE #{table_name};"
+    end
+
+    # explicitly drop table
+    def drop!
+      connection.exec "DROP TABLE #{table_name};"
+      @table_name = nil
+    end
+    
+    def system_command #:nodoc
       unless sed_command.empty?
         "cat #{@stream} | sed #{sed_command.join(" | sed ")}" 
       else
@@ -131,27 +149,17 @@ module Theman
       end
     end
 
-    # addition of a primary key after the data has been piped to 
-    # the table
-    def add_primary_key
-      instance.connection.raw_connection.query "ALTER TABLE #{instance.table_name} ADD COLUMN agents_pkey serial PRIMARY KEY;"
-    end
-
-    # use postgress COPY command using STDIN with CSV HEADER
-    # reads chunks of 8192 bytes to save memory
-    def pipe_it(l = "")
-      raise "table does not exist" unless instance.table_exists?
-      raw = instance.connection.raw_connection
-      raw.query psql_command.join("; ")
+    def pipe_it(l = "") #:nodoc
+      connection.exec psql_command.join("; ")
       f = IO.popen(system_command)
       begin
         while f.read(8192, l)
-          raw.put_copy_data l
+          connection.put_copy_data l
         end
       rescue EOFError
         f.close
       end
-      raw.put_copy_end
+      connection.put_copy_end
     end
   end
 end

data/lib/theman/object.rb

@@ -0,0 +1,22 @@
+module Theman
+  class Object
+    def self.new(table_name, parent = ::Object, conn = nil)
+      Class.new(parent) do
+        unless conn.nil?
+          @@connection = conn
+        end
+        instance_eval <<-EOV, __FILE__, __LINE__ + 1
+          set_table_name "#{table_name}"
+          
+          def table_name
+            "#{table_name}"
+          end
+          
+          def inspect
+            "Agent (#{table_name})"
+          end
+        EOV
+      end
+    end
+  end
+end

data/lib/theman/version.rb

@@ -1,3 +1,3 @@
 module Theman
-  VERSION = "0.0.5"
+  VERSION = "0.1.0"
 end

data/lib/theman.rb

@@ -1,4 +1,11 @@
+require 'active_record'
+require 'active_record/connection_adapters/postgresql_adapter'
+
 require 'theman/agency'
+require 'theman/agency/columns'
+require 'theman/agency/table'
+
+require 'theman/object'
 
 module Theman
 end