nfcollector 3.2.0

data/lib/nfcollector/mapping/indexer.rb

@@ -0,0 +1,21 @@
+module Nfcollector
+  module Mapping
+    class Indexer
+      def initialize(account_id)
+        @users      = {}
+        @account_id = account_id
+      end
+
+      def perform(mapped_row, indicies)
+        username, _ = mapped_row.values_at(*indicies)
+        # May seem clunky but hashes are much faster than Sets
+        # (plus we can store other things here later)
+        @users[username] = nil
+      end
+
+      def commit!
+        Configuration.user_committer.commit(@account_id, @users)
+      end
+    end
+  end
+end

data/lib/nfcollector/mapping/mapped_row.rb

@@ -0,0 +1,21 @@
+module Nfcollector
+  module Mapping
+    class MappedRow < Array
+      attr_writer :category_ids
+      attr_accessor :id
+
+      def initialize(*args)
+        super()
+        concat(args)
+      end
+
+      def category_ids
+        @category_ids || []
+      end
+
+      def inspect
+        Array([@id, super]).compact.join(":")
+      end
+    end
+  end
+end

data/lib/nfcollector/mapping/output.rb

@@ -0,0 +1,59 @@
+module Nfcollector
+  module Mapping
+    # TODO: Consider merging this and DefaultOutput
+    class Output
+      class << self
+        attr_accessor :categories_processor
+
+        def outputs
+          @output ||= []
+        end
+
+        def output(name, options = {})
+          self.outputs << new(name, options)
+        end
+
+        def categories(options = {})
+          self.categories_processor = CategoriesProcessor.new(options)
+        end
+
+        def build_transpiler(input_definition, column_options = {})
+          Transpiler.new(input_definition, column_options).tap do |transpiler|
+            outputs.each do |output|
+              transpiler.add_column(output)
+            end
+            transpiler.process_categories_with(self.categories_processor)
+          end
+        end
+      end
+
+      attr_reader :name
+      attr_reader :inputs
+      attr_reader :process_with
+
+      def initialize(name, options = {})
+        @name         = name.to_sym
+        @inputs       = Array(options.fetch(:inputs, @name))
+        @process_with = options.fetch(:process_with, @name)
+        @pkey         = options[:primary_key]
+        @optional     = options[:optional]
+      end
+
+      def optional?
+        @optional == true
+      end
+
+      def required?
+        !optional?
+      end
+
+      def primary_key?
+        @pkey == true
+      end
+
+      def method_missing(method, arg, options = {})
+        return arg
+      end
+    end
+  end
+end

data/lib/nfcollector/mapping/transpiler.rb

@@ -0,0 +1,92 @@
+
+module Nfcollector
+  module Mapping
+    class AfterRowHook
+      # Represents a hook that is called after each row has been transpiled
+      # @param [Object,Class] processor an object or class that responds to the perform method
+      # @param [Array<Integer>] an array of indicies from which the processing arguments are taken in the transpiled row
+      def initialize(processor, indicies)
+        @processor = processor
+        @indicies  = indicies
+      end
+
+      def perform(mapped_row)
+        @processor.perform(mapped_row, @indicies)
+      end
+    end
+
+    class Transpiler
+      attr_reader :outputs
+
+      def initialize(input_definition, column_options = {})
+        @input_definition   = input_definition
+        @outputs            = []
+        @after_row_hooks    = []
+        @column_options     = column_options
+      end
+
+      def headers
+        @outputs.map(&:name)
+      end
+
+      def add_column(output)
+        # TODO: Make this output.should_transpile?(@input_definition)
+        if output.inputs.blank? || @input_definition.has_index_for?(output.inputs) || output.required?
+          column_transpiler = ColumnTranspiler.new(output, @column_options)
+          column_transpiler.build!(@input_definition)
+          @outputs << column_transpiler
+          if output.primary_key?
+            @primary_key_index = @outputs.size - 1
+          end
+        end
+      end
+
+      def process_categories_with(categories_processor)
+        @categories_processor = categories_processor
+        if @input_definition.has_index_for?(:category) || @categories_processor.required?
+          @category_index = @input_definition.column_index(:category)
+        end
+      end
+
+      def after_row(processor, options = {})
+        required_columns = options[:using]
+        indicies = Array(required_columns).map { |column| index_of(column) }
+        @after_row_hooks << AfterRowHook.new(processor, indicies)
+      end
+
+      # @param [Symbol] name of the column
+      # @return the index of the column name in the output
+      def index_of(column)
+        @outputs.index { |output| output.name == column }
+      end
+
+      def transpile(row)
+        MappedRow.new.tap do |out|
+          @outputs.each do |ct|
+            out << ct.go(row)
+          end
+          out.category_ids = process_categories(row)
+          out.id           = primary_key(out)
+          process_after_row_hooks(out)
+        end
+      end
+
+      def process_after_row_hooks(mapped_row)
+        @after_row_hooks.each do |hook|
+          hook.perform(mapped_row)
+        end
+      end
+
+      def process_categories(row)
+        if @categories_processor.present? && @category_index.present?
+          @categories_processor.perform(row[@category_index])
+        end
+      end
+
+      def primary_key(output_row)
+        return nil if @primary_key_index.blank?
+        output_row[@primary_key_index]
+      end
+    end
+  end
+end

data/lib/nfcollector/nfcollector_exception.rb

@@ -0,0 +1,4 @@
+
+module Nfcollector
+  class NfcollectorException < RuntimeError; end
+end

data/lib/nfcollector/partition.rb

@@ -0,0 +1,76 @@
+
+module Nfcollector
+  class Partition
+    attr_reader :rows
+    attr_reader :partition_id
+    attr_reader :headers
+    attr_reader :account_id
+
+    # Instantiates a partitioner that will generate partitions
+    # of this type
+    #
+    # @param [Integer,String] account_id the account_id for which we are processing data
+    # @param [Integer] partition_key_index the index in each row of the partition key
+    # @param [Array] headers an array of headers for writing the data to file
+    #
+    def self.get_partitioner(account_id, partition_key_index, headers)
+      Partitioner.new(self, account_id, partition_key_index).tap do |part|
+        part.set_headers!(headers)
+      end
+    end
+
+    # The ActiveRecord class used to back this partition type
+    class_attribute :ar_klass
+
+    # Create a new partition with the given values
+    # @param [Array] values
+    #
+    def initialize(value, account_id, headers)
+      @value      = value
+      @account_id = account_id
+      @headers    = headers
+      @rows       = []
+    end
+
+    def partition_id
+      self.class.partition_id(@value)
+    end
+
+    # Adds a row to the partition
+    # @param [MappedRow] row
+    #
+    def add_row(row)
+      # TODO: Is there a cleaner way to do this??
+      @rows << [ @account_id ] + row
+    end
+
+    # Return the table name that should be used for the data
+    # in this partition
+    # @param [Integer,String] account_id
+    #
+    def table_name
+      return @table_name if @table_name
+      partition = ar_klass.partitions.find_or_create_for(keys)
+      @table_name = partition.nil? ? ar_klass.table_name : partition.name
+    end
+
+    def file_name
+      @file_name ||= File.expand_path("#{table_name}_#{randstr}.copy", Configuration.output_dir)
+    end
+
+    # @return [Hash] a set of keys used to generate the table for this partition
+    #
+    def keys
+      raise NotImplemented
+    end
+
+    private
+    def randstr
+      "".tap do |str|
+        8.times do
+          str << ((rand * 25).to_i + 97).chr
+        end
+      end
+    end
+  end
+end

data/lib/nfcollector/partitioner.rb

@@ -0,0 +1,37 @@
+module Nfcollector
+  class Partitioner
+    attr_reader :column_indexes
+    attr_reader :data
+
+    # TODO: use 2.1 named parameters when we can guarantee Ruby 2.1
+    def initialize(partition_klass, account_id, column_index)
+      @partition_klass = partition_klass
+      @account_id      = account_id
+      @column_index    = column_index
+      @data            = {}
+    end
+
+    def set_headers!(headers)
+      @headers = headers
+    end
+
+    def add_row(row)
+      p_value = row_value(row)
+      p_id = @partition_klass.partition_id(p_value)
+      @data[p_id] ||= @partition_klass.new(p_value, @account_id, @headers)
+      @data[p_id].add_row(row)
+    end
+
+    def row_value(row)
+      row[@column_index]
+    end
+
+    def commit!(account_id)
+      @data.map do |(key,partition)|
+        CopyFileWriter.open(partition) do |copy|
+          copy.commit!
+        end
+      end
+    end
+  end
+end

data/lib/nfcollector/payload_processor.rb

@@ -0,0 +1,46 @@
+module Nfcollector
+  class PayloadProcessor
+
+    def initialize(account_id, tz, attributes_string)
+      @account_id = account_id
+      @tz         = tz
+      input_definition = Attributes.parse(attributes_string)
+      # TODO: Will need to specify the Output class to use later
+      @transpiler = Mapping::DefaultOutput.build_transpiler(input_definition, tz: @tz)
+      @categoriser = Categoriser.new(account_id)
+      @indexer     = Mapping::Indexer.new(account_id)
+      # TODO: Do the category lookups inside the categoriser and pass it the categories string
+      # (Only do this if the input definition has the categories - actuall ALL the required inputs)
+      @transpiler.after_row(@categoriser, using: [ :domain, :host, :created_at ])
+      @transpiler.after_row(@indexer, using: [ :username ])
+    end
+
+    def process_payload(io)
+      # TODO: Look at streaming with GzipReader#read_partial later (not the lowest hanging fruit right now)
+      gz = Zlib::GzipReader.new(io)
+      strio = StringIO.new(gz.read)
+      gz.close
+
+      weblog_partitioner = WeblogPartition.get_partitioner(@account_id, @transpiler.index_of(:created_at), @transpiler.headers)
+
+      # Improvements for the future
+      # - Switched from excelsior to CSV because it was segfaulting under the latest Ruby
+      # - Ideally we would use ccsv here but it cannot read from an IO, only a file
+      # We could look at using sidekiq with Pure Ruby
+      # or have a job unzip to a second file first (not very memory efficient, and error prone)
+      # Or look at rewriting the collector in Go and use http://www.goworker.org/ (Ideal but lots of work)
+      #
+      csv = CSV.new(strio)
+      while !csv.eof?
+        mapped_row = @transpiler.transpile(csv.readline) # TODO: rename the transpile method
+        weblog_partitioner.add_row(mapped_row)
+      end
+
+      # TODO: These could write on the fly in time
+      cfws = weblog_partitioner.commit!(@account_id)
+      @categoriser.commit!
+      @indexer.commit!
+      return cfws
+    end
+  end
+end

data/lib/nfcollector/sequence_generator.rb

@@ -0,0 +1,11 @@
+module Nfcollector
+  class SequenceGenerator
+    def initialize
+      @start = 0 # TODO: This will have to be smarter (needs to keep a track of the last used value somewhere)
+    end
+
+    def next
+      @start += 1
+    end
+  end
+end

data/lib/nfcollector/version.rb

@@ -0,0 +1,3 @@
+module Nfcollector
+  VERSION = "3.2.0"
+end

data/lib/nfcollector/weblog_partition.rb

@@ -0,0 +1,26 @@
+
+module Nfcollector
+  class WeblogPartition < Partition
+
+    # Calculates the ID for this partition
+    # Think of it as a statistical 'bin'
+    #
+    # In this case we take the date from a timestamp
+    #
+    def self.partition_id(value)
+      Date.new(value.year, value.month, value.mday)
+    end
+
+    def keys
+      { account_id: account_id, created_at: range_from_date }
+    end
+
+    def range_from_date
+      date       = partition_id
+      next_date  = partition_id.next
+      start_time = Time.utc(date.year, date.month, date.mday)
+      end_time   = Time.utc(next_date.year, next_date.month, next_date.mday)
+      start_time...end_time
+    end
+  end
+end

data/nfcollector.gemspec

@@ -0,0 +1,30 @@
+# -*- encoding: utf-8 -*-
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'nfcollector/version'
+
+Gem::Specification.new do |gem|
+  gem.name          = "nfcollector"
+  gem.version       = Nfcollector::VERSION
+  gem.authors       = ["Dan Draper"]
+  gem.email         = ["daniel@codefire.com"]
+  gem.description   = %q{Data collection system for NetFox}
+  gem.summary       = %q{Data Collection}
+  gem.homepage      = ""
+
+  gem.files         = `git ls-files`.split($/)
+  gem.executables   = gem.files.grep(%r{^bin/}).map{ |f| File.basename(f) }
+  gem.test_files    = gem.files.grep(%r{^(test|spec|features)/})
+  gem.require_paths = ["lib"]
+
+  gem.add_development_dependency 'test-unit'
+  gem.add_development_dependency 'rspec'
+  gem.add_development_dependency 'shoulda'
+  gem.add_development_dependency 'timecop'
+  gem.add_development_dependency 'differ'
+  gem.add_development_dependency 'byebug'
+  gem.add_development_dependency 'activerecord'
+
+  gem.add_runtime_dependency 'domainatrix', '>= 0.0.10'
+  gem.add_runtime_dependency 'activesupport', '>= 3.2'
+end

data/spec/attribute_validator_spec.rb

@@ -0,0 +1,23 @@
+require 'spec_helper'
+
+describe Nfcollector::AttributeValidator do
+  subject(:validator) { Nfcollector::AttributeValidator.new }
+
+  specify "raise if a duplicate attribute is provided" do
+    expect {
+      subject.validate(">a,Rr,Rh,Rp,>a,>A,t,Un,Hs,mt,<s,Cs")
+    }.to raise_error(Nfcollector::DuplicateAttribute)
+  end
+
+  specify "raise if an unknown attribute is provided" do
+    expect {
+      subject.validate(">a,Rr,Rh,Rp,>A,t,Un,Hs,mt,<s,Cs,GG")
+    }.to raise_error(Nfcollector::UnknownAttribute, "Unknown Attribute GG")
+  end
+
+  specify "raise if a required attribute is NOT provided" do
+    expect {
+      subject.validate(">a,Rr,Rh,Rp,>A,Un,Hs,mt,<s,Cs")
+    }.to raise_error(Nfcollector::MissingRequiredAttribute, "Missing attribute 't'")
+  end
+end