DanaDanger-data_transport 0.1.1

data/lib/data_transport.rb

@@ -0,0 +1,30 @@
+require "data_transport/map"
+require "data_transport/data_store"
+
+module DataTransport
+  def self.default_batch_size
+    1000
+  end
+  
+  def self.map(input, output, options = {}, &block)
+    # Extract options.
+    ignore_duplicates = options.delete(:ignore_duplicates)
+    unless options.empty?
+      raise(ArgumentError, "unrecognized options: `#{options.join("', `")}'")
+    end
+    # If ignore_duplicates is true, make sure the output is a MySQL database.
+    if ignore_duplicates
+      unless output.is_a?(DataStore::ActiveRecord) && output.klass.connection.is_a?(::ActiveRecord::ConnectionAdapters::MysqlAdapter)
+        raise ArgumentError, "ignore_duplicates can only be used with an ActiveRecord data store connected to a MySQL database"
+      end
+    end
+    # Run the transport.
+    output.reset
+    output.ignore_duplicates = true if ignore_duplicates
+    map = DataTransport::Map.new(&block)
+    input.each_record do |record|
+      output.write_record(map.map(record))
+    end
+    output.finalize
+  end
+end

data/lib/data_transport/data_store.rb

@@ -0,0 +1,31 @@
+require "data_transport/data_store/file"
+require "data_transport/data_store/active_record"
+
+module DataTransport
+  class DataStore
+    def count
+      raise NotImplementedError
+    end
+    
+    def each_record(batch_size = 1000)
+      raise NotImplementedError
+    end
+    
+    def write_record(record)
+      raise NotImplementedError
+    end
+    
+    def finalize
+      # Do nothing by default.
+    end
+    
+  protected
+    
+    def check_batch_size(batch_size)
+      batch_size ||= DataTransport.default_batch_size
+      raise(TypeError, "batch size must be an integer") unless batch_size.is_a?(Integer)
+      raise(RangeError, "batch size must be greater than zero") if batch_size < 1
+      batch_size
+    end
+  end
+end

data/lib/data_transport/data_store/active_record.rb

@@ -0,0 +1,146 @@
+module DataTransport
+  class DataStore
+    class ActiveRecord < DataStore
+      attr_accessor :ignore_duplicates
+      
+      def initialize(options = {})
+        super()
+        # Extract options.
+        @class      = options.delete(:class)
+        @connection = options.delete(:connection)
+        @table_name = options.delete(:table_name)
+        @conditions = options.delete(:conditions)
+        @truncate   = options.delete(:truncate)
+        # Make sure a class or connection and table name was provided.
+        if @class.nil? && (@connection.nil? || @table_name.nil?)
+          raise(ArgumentError, "missing required option `class', or `connection' and `table_name'")
+        end
+        raise(TypeError, "class must be a class") if @class && !@class.is_a?(Class)
+        # If connection specs were provided instead of a class, make an
+        # anonymous ActiveRecord subclass.
+        unless @class
+          @class = Class.new(::ActiveRecord::Base)
+          @class.set_table_name @table_name
+          @class.establish_connection @connection
+        end
+        # Make sure the class descends from ActiveRecord::Base.
+        klass = @class.superclass
+        is_active_record = false
+        while klass
+          if klass == ::ActiveRecord::Base
+            is_active_record = true
+            break
+          end
+          klass = klass.superclass
+        end
+        raise(TypeError, "class must descend from ActiveRecord::Base") unless is_active_record
+        # Check for unknown options.
+        unless options.empty?
+          raise(ArgumentError, "unrecognized options: `#{options.join("', `")}'")
+        end
+        # Ask the database how much data it can handle in one query. This only
+        # works on MySQL.
+        begin
+          rows = @class.connection.select_all("SHOW VARIABLES LIKE 'max_allowed_packet'")
+          @max_allowed_packet = rows.first["Value"].to_i - 512
+        rescue
+        end
+        # Fetch column information
+        @columns = {}
+        @class.columns.each {|c| @columns[c.name.to_sym] = c}
+      end
+      
+      def klass
+        @class
+      end
+      
+      def count
+        @class.count(:conditions => @conditions)
+      end
+      
+      def each_record(batch_size = nil)
+        batch_size = check_batch_size(batch_size)
+        
+        conn = @class.connection
+        column_names = conn.columns(@class.table_name).collect {|c| c.name}
+        
+        offset = 0
+        record = {}
+        base_query = "SELECT * FROM #{conn.quote_table_name(@class.table_name)}"
+        @class.send(:add_conditions!, base_query, @conditions) unless @conditions.nil?
+        while true
+          sql = base_query.dup
+          conn.add_limit_offset!(sql, :limit => batch_size, :offset => offset)
+          offset += batch_size
+          rows = conn.select_rows(sql)
+          break if rows.empty?
+          rows.each do |row|
+            record.clear
+            column_names.each_with_index do |column_name, i|
+              column_name = column_name.to_sym
+              record[column_name] = @columns[column_name].type_cast(row[i])
+            end
+            yield record
+          end
+        end
+      end
+      
+      def write_record(record)
+        conn = @class.connection
+        # If no SQL has been produced yet, start an INSERT statement.
+        @sql_buffer ||= start_insert_sql(record)
+        # Convert the record into a string of quoted values.
+        values = []
+        record.each {|k, v| values << conn.quote(v, @columns[k])}
+        values = "(#{values.join ","}),"
+        # Write the record.
+        if @max_allowed_packet.nil?
+          # We have no information on the database's maximum allowed packet
+          # size, so it's safest to write the record immediately.
+          @sql_buffer << values
+          finalize
+        elsif @sql_buffer.length + record.length > @max_allowed_packet
+          # Appending this record to the SQL buffer will exceed the maximum
+          # allowed packet size. Send the buffer to the database and start a
+          # new statement with this record.
+          finalize
+          @sql_buffer = start_insert_sql
+          @sql_buffer << values
+        else
+          # This record will not cause the SQL buffer to exceed the maximum
+          # allowed packet size. Append it to the SQL buffer.
+          @sql_buffer << values
+        end
+      end
+      
+      def finalize
+        if @truncate
+          conn = @class.connection
+          begin
+            conn.execute("TRUNCATE TABLE #{conn.quote_table_name(@class.table_name)}")
+          rescue
+            @class.delete_all
+          end
+          @truncate = false
+        end
+        if @sql_buffer[-1,1] == ","
+          @sql_buffer.chop!
+          @class.connection.execute(@sql_buffer)
+        end
+      end
+      
+      def reset
+        self.ignore_duplicates = false
+        @sql_buffer = nil
+      end
+      
+    private
+      
+      def start_insert_sql(record)
+        "INSERT #{ignore_duplicates ? "IGNORE " : " "}INTO " +
+          "#{@class.connection.quote_table_name(@class.table_name)} " +
+          "(#{record.keys.join ","}) VALUES "
+      end
+    end
+  end
+end

data/lib/data_transport/data_store/file.rb

@@ -0,0 +1,117 @@
+module DataTransport
+  class DataStore
+    class File < DataStore
+      attr_reader :mode
+      
+      def initialize(options = {})
+        super()
+        
+        @header    = options.delete(:header)
+        @delimiter = options.delete(:delimiter) || "\t"
+        @path      = options.delete(:path)
+        @null      = options.delete(:null) || ""
+        @keys      = options.delete(:keys)
+        
+        raise(ArgumentError, "missing required option `path'") if @path.nil?
+        unless options.empty?
+          raise(ArgumentError, "unrecognized options: `#{options.join("', `")}'")
+        end
+      end
+      
+      def count
+        return @count if @count
+        self.mode = :input
+        line_count = 0
+        rewind_and_restore do
+          io.readline if @header
+          until io.eof?
+            io.gets
+            line_count += 1
+          end
+        end
+        @count = line_count
+      end
+      
+      def each_record(batch_size = nil)
+        self.mode = :input
+                
+        batch_size = check_batch_size(batch_size)
+        
+        io.rewind
+        io.readline if @header
+        until io.eof?
+          line = io.gets || break
+          line.chomp!
+          values = line.split(/#{@delimiter}/)
+          if keys.length != values.length
+            raise RuntimeError, "wrong number of fields (#{values.length} for #{keys.length})"
+          end
+          record = {}
+          keys.length.times {|i| record[keys[i]] = values[i]}
+          yield record
+        end
+      end
+      
+      def write_record(record)
+        self.mode = :output
+        # If no key order was ever specified, make one up.
+        @keys ||= record.keys.sort {|a,b| a.to_s <=> b.to_s}
+        # Write the header if this is the first record.
+        if @header && io.pos == 0
+          io.puts(keys.join(@delimiter))
+        end
+        # Write the values in a predictable order.
+        values = keys.collect do |k|
+          record[k].nil?? @null : record[k]
+        end
+        io.puts(values.join(@delimiter))
+      end
+      
+    private
+      
+      def mode=(new_mode)
+        if !@mode.nil? && @mode != new_mode
+          raise RuntimeError, "can't switch mode from #{@mode} to #{new_mode}"
+        end
+        unless [:input, :output].include?(new_mode)
+          raise ArgumentError, "unknown mode `#{new_mode}'"
+        end
+        @mode = new_mode
+      end
+      
+      def io
+        return @io if @io
+        if mode == :output
+          @io = ::File.open(@path, "w")
+          @io.rewind
+          @io
+        else
+          @io = ::File.open(@path, "r")
+        end
+      end
+      
+      def keys
+        return @keys if @keys
+        return [] if mode == :output
+        line = rewind_and_restore { io.readline }
+        line.chomp!
+        fields = line.split(/#{@delimiter}/)
+        if @header
+          @keys = fields.collect! {|hdr| hdr.downcase.to_sym}
+        else
+          @keys = (0..(fields.length - 1)).to_a.collect! do |i|
+            sprintf("column%02d", i).to_sym
+          end
+        end
+      end
+      
+      def rewind_and_restore
+        pos = io.pos
+        io.rewind
+        result = yield
+        io.seek(pos)
+        result
+      end
+    end
+  end
+end

data/lib/data_transport/map.rb

@@ -0,0 +1,21 @@
+require "data_transport/record/source"
+require "data_transport/record/destination"
+
+module DataTransport
+  class Map
+    attr_reader :source, :destination
+    
+    def initialize(&block)
+      @block = block
+      @source = DataTransport::Record::Source.new
+      @destination = DataTransport::Record::Destination.new
+    end
+    
+    def map(record)
+      @source.record = record
+      @destination.reset!
+      @block.call(@source, @destination)
+      @destination.record
+    end
+  end
+end

data/lib/data_transport/record/destination.rb

@@ -0,0 +1,28 @@
+module DataTransport
+  module Record
+    class Destination
+      attr_reader :record
+    
+      def initialize
+        @record = {}
+      end
+    
+      def reset!
+        @record.clear
+      end
+    
+      def method_missing(name, *args)
+        name_s = name.to_s
+        if name_s[-1,1] == "="
+          unless args.length == 1
+            raise ArgumentError, "wrong number of arguments (#{args.length} for 1)"
+          end
+          name_s.chop!
+          @record[name_s.to_sym] = args.first
+        else
+          super
+        end
+      end
+    end
+  end
+end

data/lib/data_transport/record/source.rb

@@ -0,0 +1,24 @@
+module DataTransport
+  module Record
+    class Source
+      def record=(record)
+        @record = record
+      end
+      
+      def id
+        method_missing :id
+      end
+      
+      def method_missing(name, *args)
+        if @record.has_key?(name)
+          unless args.empty?
+            raise ArgumentError, "wrong number of arguments (#{args.length} for 0)"
+          end
+          @record[name]
+        else
+          super
+        end
+      end
+    end
+  end
+end

metadata

@@ -0,0 +1,59 @@
+--- !ruby/object:Gem::Specification 
+name: DanaDanger-data_transport
+version: !ruby/object:Gem::Version 
+  version: 0.1.1
+platform: ruby
+authors: 
+- Dana Danger
+autorequire: 
+bindir: bin
+cert_chain: []
+
+date: 2009-04-08 00:00:00 -07:00
+default_executable: 
+dependencies: []
+
+description: 
+email: 
+executables: []
+
+extensions: []
+
+extra_rdoc_files: []
+
+files: 
+- lib/data_transport.rb
+- lib/data_transport/map.rb
+- lib/data_transport/data_store.rb
+- lib/data_transport/data_store/active_record.rb
+- lib/data_transport/data_store/file.rb
+- lib/data_transport/record/destination.rb
+- lib/data_transport/record/source.rb
+has_rdoc: false
+homepage: http://github.com/DanaDanger/data_transport
+post_install_message: 
+rdoc_options: []
+
+require_paths: 
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement 
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      version: "0"
+  version: 
+required_rubygems_version: !ruby/object:Gem::Requirement 
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      version: "0"
+  version: 
+requirements: []
+
+rubyforge_project: 
+rubygems_version: 1.2.0
+signing_key: 
+specification_version: 2
+summary: A gem for importing and exporting large quantities of data.
+test_files: []
+