chicago-etl 0.0.9

data/lib/chicago/etl/table_builder.rb

@@ -0,0 +1,45 @@
+module Chicago
+  module ETL
+    # Builds ETL tables.
+    class TableBuilder
+      # Creates the necessary tables for the ETL process in the given
+      # database.
+      def self.build(db)
+        new(db).build
+      end
+
+      def initialize(db) # :nodoc:
+        @db = db
+      end
+
+      def build # :nodoc:
+        create_table :etl_batches do
+          primary_key :id, :type => :integer, :unsigned => true
+          timestamp   :started_at, :null => false, :default => :current_timestamp.sql_function
+          timestamp   :finished_at, :null => true, :default => nil
+          timestamp   :extracted_to, :null => true, :default => nil
+          enum        :state, :null => false, :elements => %w{Started Finished Error}, :default => "Started"
+        end
+
+        create_table :etl_task_invocations do
+          primary_key :id, :type => :integer, :unsigned => true
+          integer     :batch_id, :unsigned => true, :null => false
+          enum        :stage, :null => false, :elements => %w{Extract Transform Load}
+          varchar     :name, :null => false
+          timestamp   :started_at, :null => false, :default => :current_timestamp.sql_function
+          timestamp   :finished_at, :null => true, :default => nil
+          enum        :state, :null => false, :elements => %w{Created Started Finished Error}, :default => "Created"
+          smallint    :attempts, :null => false, :unsigned => true
+
+          index [:batch_id, :stage, :name], :unique => true
+        end
+      end
+
+      private
+
+      def create_table(table, &block)
+        @db.create_table(table, :engine => "innodb", &block) unless @db.tables.include?(table)
+      end
+    end
+  end
+end

data/lib/chicago/etl/task_invocation.rb

@@ -0,0 +1,32 @@
+module Chicago
+  module ETL
+
+    class TaskInvocation < Sequel::Model
+      set_dataset :etl_task_invocations
+      many_to_one :batch
+
+      # Executes a block of code.
+      #
+      # Sets the state to "Error" and re-raises any exception that the
+      # block of code raises.
+      def perform
+        raise RuntimeError.new("The task #{name} in batch #{batch_id} has already run") if finished?
+        update(:state => "Started", :attempts => attempts + 1)
+        begin
+          yield
+        rescue Exception => e
+          update(:state => "Error")
+          batch.error if batch
+          raise e
+        end
+        update(:state => "Finished", :finished_at => Time.now)
+      end
+
+      # Returns true if this task has finished running successfully.
+      def finished?
+        state == "Finished"
+      end
+    end
+
+  end
+end

data/lib/chicago/etl/tasks.rb

@@ -0,0 +1,34 @@
+require 'rake/tasklib'
+
+module Chicago
+  module ETL
+    # ETL Rake tasks for a Chicago project.
+    #
+    # To use, simply include:
+    #
+    #     Chicago::ETL::RakeTasks.new(db, schema)
+    #
+    # in your project's Rakefile.
+    #
+    # Provides the following tasks:
+    #
+    # +db:create_etl_tables+:: defines the tables used for ETL batches
+    #                          and the like
+    class RakeTasks < Rake::TaskLib
+      def initialize(db, schema)
+        @db = db
+        @schema = schema
+        define
+      end
+
+      def define
+        namespace :db do
+          desc "Creates the etl tables"
+          task :create_etl_tables do
+            Chicago::ETL::TableBuilder.build(@db)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/chicago/etl/transformations/add_insert_timestamp.rb

@@ -0,0 +1,16 @@
+module Chicago
+  module ETL
+    module Transformations
+      class AddInsertTimestamp
+        def initialize(timestamp=Time.now)
+          @insert_timestamp = timestamp.utc
+        end
+
+        def call(row, errors=[])
+          row[:_inserted_at] = @insert_timestamp
+          [row, errors]
+        end
+      end
+    end
+  end
+end

data/lib/chicago/etl/transformations/uk_post_code.rb

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+module Chicago
+  module ETL
+    module Transformations
+      # Cleans and reformats UK-based postcodes.
+      #
+      # Transformations are based on observed errors in data entry, so
+      # shift key slips (i.e. typing '!' where '1' was meant) are
+      # corrected, as are use of numbers where letters were intended
+      # i.e. (0X -> OX for Oxfordshire postcodes).
+      #
+      # Leaves BFPO postcodes alone.
+      class UkPostCode
+        # Creates a new post code transformation.
+        #
+        # @param Symbol column_name - the name of the column
+        #   containing the post code.
+        #
+        # @param Proc filter_block - an optional block, which takes a
+        #   row. If the block returns false, the transformation will
+        #   not be run. This can be useful to only run the
+        #   transformation on UK addresses, based a country field in
+        #   the row for example.
+        def initialize(column_name, &filter_block)
+          @column_name = column_name
+          @filter_block = filter_block
+        end
+
+        def call(row, errors=[])
+          return [row, errors] if @filter_block && !@filter_block.call(row)
+
+          row[@column_name] = UkPostCodeField.new.
+            normalize(row[@column_name])[:post_code]
+
+          [row, errors]
+        end
+      end
+    end
+  end
+end

data/lib/chicago/etl/transformations/uk_post_code_field.rb

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+
+module Chicago
+  module ETL
+    module Transformations
+      class UkPostCodeField
+        MATCH = /\A([A-Z][A-Z]?[0-9][0-9A-Z]?)(?:([0-9][A-Z]{2}))?\Z/
+        
+        # Returns cleaned, formatted data about a UK Post Code.
+        #
+        # Example:
+        #
+        #     UkPostCodeField.new.normalize(" SW !2 4 GH")
+        #     # => { :post_code => "SW12 4GH",
+        #            :outcode => "SW12",
+        #            :incode => "4GH" }
+        #
+        # Partial postcodes will be returned without the incode. BFPO
+        # postcodes are supported, but have no incode or
+        # outcode. Postcodes that do not follow the format will be
+        # returned as is, with an invalid key set.
+        def normalize(raw_post_code)
+          reformat(clean(raw_post_code)) ||
+            {:post_code => raw_post_code, :invalid => true}
+        end
+
+        private
+
+        def clean(raw_post_code)
+          raw_post_code.
+            strip.
+            upcase.
+            tr('!"$%^&*()', '124567890').
+            gsub("£", "3").
+            sub(/^0([XL])/, 'O\1').
+            sub(/^([PSCY])0/, '\1O')
+        end
+
+        def reformat(post_code)
+          if post_code[0..3] == "BFPO"
+            { :post_code => post_code.sub(/BFPO\s*/, "BFPO ") }
+          else
+            reformat_standard_post_code(post_code)
+          end
+        end
+        
+        def reformat_standard_post_code(post_code)
+          match = post_code.gsub(/\s+/,'').match(MATCH)
+
+          unless match.nil?
+            { :outcode => match[1],
+              :incode => match[2],
+              :post_code => [match[1], match[2]].join(' ').strip }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/chicago/etl.rb

@@ -0,0 +1,35 @@
+require 'sequel'
+
+require 'chicago/etl/counter'
+require 'chicago/etl/key_builder'
+require 'chicago/etl/sink'
+require 'chicago/etl/mysql_load_file_value_transformer'
+require 'chicago/etl/buffering_insert_writer'
+require 'chicago/etl/mysql_dumpfile'
+
+require 'chicago/etl/load_dataset_builder'
+
+# Sequel Extensions
+require 'chicago/etl/sequel/filter_to_etl_batch'
+require 'chicago/etl/sequel/load_data_infile'
+require 'chicago/etl/sequel/dependant_tables'
+
+# Screens
+require 'chicago/etl/screens/column_screen'
+require 'chicago/etl/screens/composite_screen'
+require 'chicago/etl/screens/missing_value'
+require 'chicago/etl/screens/invalid_element'
+require 'chicago/etl/screens/out_of_bounds'
+
+# Transformations
+require 'chicago/etl/transformations/add_insert_timestamp'
+require 'chicago/etl/transformations/uk_post_code'
+require 'chicago/etl/transformations/uk_post_code_field'
+
+module Chicago
+  module ETL
+    autoload :TableBuilder,   'chicago/etl/table_builder.rb'
+    autoload :Batch,          'chicago/etl/batch.rb'
+    autoload :TaskInvocation, 'chicago/etl/task_invocation.rb'
+  end
+end

data/spec/db_connections.yml.dist

@@ -0,0 +1,4 @@
+adapter: mysql
+username: root
+socket: /tmp/mysql.sock
+database: chicago_test_db

data/spec/etl/batch_spec.rb

@@ -0,0 +1,86 @@
+require 'spec_helper'
+
+describe Chicago::ETL::Batch do
+  before :each do
+    TEST_DB.drop_table(*(TEST_DB.tables))
+    ETL::TableBuilder.build(TEST_DB)
+    ETL::Batch.db = TEST_DB
+    Chicago.project_root = File.expand_path(File.join(File.dirname(__FILE__), ".."))
+    tmpdir = File.expand_path(File.join(File.dirname(__FILE__), "..", "tmp"))
+    FileUtils.rm_r(tmpdir) if File.exists?(tmpdir)
+  end
+
+  it "should return a new batch when instance is called and there are no outstanding batches in error" do
+    ETL::Batch.instance.should be_new
+  end
+
+  it "should set the start timestamp of the batch to now when created" do
+    ETL::Batch.instance.start.started_at.to_i.should == Time.now.to_i
+  end
+
+  it "should have a state of 'Started' when started" do
+    ETL::Batch.instance.start.state.should == "Started"
+  end
+
+  it "should have a default extracted_to datetime of midnight (this morning)" do
+    now = Time.now
+    ETL::Batch.instance.start.extracted_to.should == Time.local(now.year, now.month, now.day, 0,0,0)
+  end
+
+  it "should be able to specify an extract to date" do
+    now = Date.today - 1
+    ETL::Batch.instance.start(now).extracted_to.should == Time.local(now.year, now.month, now.day, 0,0,0)
+  end
+
+  it "should create a directory tmp/batches/1 under the project root when created" do
+    ETL::Batch.instance.start
+    File.should be_directory(Chicago.project_root + "/tmp/batches/1")
+  end
+
+  it "should return the batch directory path from #dir" do
+    ETL::Batch.instance.start.dir.should == Chicago.project_root + "/tmp/batches/1"
+  end
+
+  it "should set the finished_at timestamp when #finish is called" do
+    batch = ETL::Batch.instance.start
+    batch.finish
+    batch.finished_at.should_not be_nil
+    batch.state.should == "Finished"
+  end
+
+  it "should return true from #error? if in the error state" do
+    batch = ETL::Batch.instance.start
+    batch.error
+    batch.should be_in_error
+  end
+
+  it "should not return a new batch if the last batch was not finished" do
+    batch = ETL::Batch.instance.start
+    ETL::Batch.instance == batch
+  end
+
+  it "should not return a new batch if the last batch ended in error" do
+    batch = ETL::Batch.instance.start
+    batch.error
+    ETL::Batch.instance.should == batch
+  end
+
+  it "should create a log in tmp/batches/1/log" do
+    ETL::Batch.instance.start
+    File.read(Chicago.project_root + "/tmp/batches/1/log").
+      should include("Started ETL batch 1.")
+  end
+
+  it "should perform a task only once" do
+    batch = ETL::Batch.instance.start
+    i = 0
+    2.times { batch.perform_task("Transform", "Test") { i += 1} }
+    i.should == 1
+    batch.task_invocations_dataset.filter(:stage => "Transform", :name => "Test").count.should == 1
+  end
+
+  it "should not complain when given a symbol as the stage name" do
+    batch = ETL::Batch.instance.start
+    lambda { batch.perform_task(:transform, "Test") {} }.should_not raise_error(Sequel::DatabaseError)
+  end
+end

data/spec/etl/counter_spec.rb

@@ -0,0 +1,44 @@
+require 'spec_helper'
+
+describe Chicago::ETL::Counter do
+  it "returns the next available key" do
+    key = described_class.new(3)
+    key.next.should == 4
+    key.next.should == 5
+  end
+
+  it "can have the initial key set via a block" do
+    counter = described_class.new { 1 + 1 }
+    counter.next.should == 3
+  end
+
+  it "defaults the counter to 0 if the block returns nil" do
+    counter = described_class.new { nil }
+    counter.next.should == 1
+  end
+
+  it "prefers the block to the argument for setting initial state" do
+    counter = described_class.new(5) { 2 }
+    counter.next.should == 3
+  end
+
+  it "can be constructed with no argument, implying 0" do
+    described_class.new.next.should == 1
+  end
+
+  it "updates keys in a thread-safe fashion" do
+    key = described_class.new
+
+    # These seem to need to be a fairly large number of times to see
+    # errors
+    [Thread.new { 100_000.times {|i| key.next } },
+     Thread.new { 100_000.times {|i| key.next } },
+     Thread.new { 100_000.times {|i| key.next } }].each(&:join)
+
+    key.next.should == 300_001
+  end
+
+  it "has a current value" do
+    described_class.new.current.should == 0
+  end
+end

data/spec/etl/etl_batch_id_dataset_filter.rb

@@ -0,0 +1,29 @@
+module Chicago
+  module ETL
+    class EtlBatchIdDatasetFilter
+      def initialize(etl_batch_id)
+        @etl_batch_id = etl_batch_id
+      end
+
+      # Returns a new dataset, filtered by all tables where the etl
+      # batch id matches.
+      def filter(dataset)
+        dataset.filter(conditions(filterable_tables(dataset)))
+      end
+
+      private
+
+      def filterable_tables(dataset)
+        dataset.dependant_tables.select {|t|
+          dataset.db.schema(t).map(&:first).include?(:etl_batch_id)
+        }
+      end
+
+      def conditions(tables)
+        tables.
+          map {|t| {:etl_batch_id.qualify(t) => @etl_batch_id} }.
+          inject {|a,b| a | b}
+      end
+    end
+  end
+end

data/spec/etl/key_builder_spec.rb

@@ -0,0 +1,190 @@
+require 'spec_helper'
+
+describe Chicago::ETL::KeyBuilder do
+  before :all do
+    @schema = Chicago::StarSchema.new
+    @schema.define_dimension(:user) do
+      columns { integer :original_id }
+    end
+
+    @schema.define_dimension(:address) do
+      columns do
+        string :line1
+        string :post_code
+      end
+
+      natural_key :line1, :post_code
+    end
+
+    @schema.define_dimension(:random) do
+      columns do
+        string :foo
+      end
+    end
+
+    @schema.define_fact(:addresses) do
+      dimensions :user, :address
+      natural_key :user, :address
+    end
+  end
+  
+  before :each do
+    @db = stub(:staging_database).as_null_object
+    @db.stub(:[]).and_return(stub(:max => nil, :select_hash => {}))
+    @writer = stub(:writer).as_null_object
+    Chicago::ETL::BufferingInsertWriter.stub(:new).and_return(@writer)
+  end
+
+  describe "for identifiable dimensions" do
+    before :each do
+      @dimension = @schema.dimension(:user)
+    end
+
+    it "returns an incrementing key, given a row" do
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 2).should == 1
+      builder.key(:original_id => 3).should == 2
+    end
+
+    it "returns the same key for the same record" do
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 2).should == 1
+      builder.key(:original_id => 2).should == 1
+    end
+
+    it "takes into account the current maximum key in the database" do
+      @db.stub(:[]).with(:keys_dimension_user).and_return(stub(:max => 2, :select_hash => {}))
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 1).should == 3
+    end
+
+    it "returns previously created keys" do
+      dataset = stub(:dataset, :max => 1, :select_hash => {40 => 1})
+      @db.stub(:[]).with(:keys_dimension_user).and_return(dataset)
+
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 30).should == 2
+      builder.key(:original_id => 40).should == 1
+    end
+
+    it "raises an error when original_id isn't present in the row" do
+      builder = described_class.for_table(@dimension, @db)
+      expect { builder.key(:foo => :bar) }.to raise_error(Chicago::ETL::KeyError)
+    end
+
+    it "flushes new keys to a key table" do
+      pending
+      dataset = stub(:dataset, :max => 1, :select_hash => {40 => 1})
+      dataset.stub(:insert_replace => dataset)
+      @db.stub(:[]).with(:keys_dimension_user).and_return(dataset)
+
+      dataset.should_receive(:multi_insert).
+        with([{:original_id => 30, :dimension_id => 2}])
+                                              
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 30)
+      builder.key(:original_id => 40)
+      builder.flush
+    end
+
+    it "flushes new keys only once" do
+      pending
+      dataset = stub(:dataset, :max => 1, :select_hash => {40 => 1})
+      dataset.stub(:insert_replace => dataset)
+      @db.stub(:[]).with(:keys_dimension_user).and_return(dataset)
+
+      dataset.should_receive(:multi_insert).
+        with([{:original_id => 30, :dimension_id => 2}])
+      dataset.should_receive(:multi_insert).with([])
+                                              
+      builder = described_class.for_table(@dimension, @db)
+      builder.key(:original_id => 30)
+      builder.key(:original_id => 40)
+      builder.flush
+      builder.flush
+    end
+
+    it "replaces old mappings with new values" do
+      pending
+      dataset = stub(:dataset, :max => 1, :select_hash => {40 => 1}, :multi_insert => nil)
+      @db.stub(:[]).with(:keys_dimension_user).and_return(dataset)
+
+      dataset.should_receive(:insert_replace).and_return(dataset)
+      described_class.for_table(@dimension, @db).flush
+    end
+  end
+
+  describe "for non-identifiable dimensions with natural keys" do
+    before :each do
+      @builder = described_class.for_table(@schema.dimension(:address), @db)
+    end
+
+    it "returns an incrementing key, given a row" do
+      @builder.key(:line1 => "some street", :post_code => "TW3 X45").
+        should == 1
+      @builder.key(:line1 => "some road", :post_code => "TW3 X45").
+        should == 2
+    end
+
+    it "returns the same incrementing key, ignoring case" do
+      @builder.key(:line1 => "some street", :post_code => "TW3 X45").
+        should == 1
+      @builder.key(:line1 => "some STREET", :post_code => "TW3 X45").
+        should == 1
+    end
+
+    it "can override default hash preparation" do
+      @builder.hash_preparation = lambda {|c| c }
+
+      @builder.key(:line1 => "some street", :post_code => "TW3 X45").
+        should == 1
+      @builder.key(:line1 => "some STREET", :post_code => "TW3 X45").
+        should == 2
+    end
+
+    it "inserts the hash as a binary literal" do
+      # Yuck. Don't like the implementation test, but mock
+      # expectations fail here for some reason, maybe because of the
+      # Sequel::LiteralString?
+      @builder.key_for_insert(@builder.original_key(:line1 => "some street", :post_code => "TW3 X45")).should == "0x817860F2417EB83D81FEA9D82E6B213A".lit
+    end
+
+    it "selects the Hex version of the binary column for the cache" do
+      dataset = stub(:dataset, :max => 1).as_null_object
+      @db.stub(:[]).with(:keys_dimension_address).and_return(dataset)
+      @builder = described_class.for_table(@schema.dimension(:address), @db)
+
+      dataset.should_receive(:select_hash).with(:hex.sql_function(:original_id).as(:original_id), :dimension_id).and_return({})
+      
+      @builder.key(:line1 => "foo")
+    end
+
+    it "uses all columns as the natural key if one isn't defined" do
+      described_class.
+        for_table(@schema.dimension(:random), @db).
+        original_key(:foo => "bar").
+        should == "3D75EEC709B70A350E143492192A1736"
+    end
+  end
+
+  describe "for facts" do
+    before :each do 
+      @builder = described_class.for_table(@schema.fact(:addresses), @db)
+    end
+
+    it "increments the id, regardless of row equality" do
+      @builder.key({}).should == 1
+      @builder.key({}).should == 2
+    end
+
+    it "increments from the last id stored id in the fact table" do
+      @db.stub(:[]).with(:facts_addresses).and_return(stub(:max => 100, :select_hash => {}))
+      @builder = described_class.for_table(@schema.fact(:addresses), @db)
+      @builder.key({}).should == 101
+    end
+
+    it "supports the flush interface as a no-op" do
+      lambda { @builder.flush }.should_not raise_error
+    end
+  end
+end