so2db 0.1.0

data/lib/so2db/models.rb

@@ -0,0 +1,106 @@
+#--
+# Copyright (c) 2012 Chad Taylor
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy 
+# of this software and associated documentation files (the "Software"), to deal 
+# in the Software without restriction, including without limitation the rights 
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
+# copies of the Software, and to permit persons to whom the Software is 
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in 
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
+# SOFTWARE.
+
+require 'active_record'
+
+module SO2DB::Models
+
+  class Badge < ActiveRecord::Base
+    def self.exported_fields
+      return [ :id, :user_id, :name, :date ]
+    end
+  end
+
+  class Comment < ActiveRecord::Base
+    def self.exported_fields
+      return [ :id, :post_id, :score, :text, :creation_date, :user_id, 
+               :user_display_name ]
+    end
+  end
+
+  class Post < ActiveRecord::Base
+    def self.exported_fields
+      return [ :id, :post_type_id, :parent_id, :accepted_answer_id, 
+               :creation_date, :score, :view_count, :body, :owner_user_id, 
+               :last_editor_user_id, :last_editor_display_name, :last_edit_date,
+               :last_activity_date, :community_owned_date, :closed_date, :title,
+               :tags, :answer_count, :comment_count, :favorite_count, 
+               :owner_display_name ]
+    end
+  end
+
+  class PostHistory < ActiveRecord::Base
+    self.table_name = "post_history"
+
+    def self.exported_fields
+      return [ :id, :post_history_type_id, :post_id, :revision_guid, 
+               :creation_date, :user_id, :user_display_name, :comment, :text,
+               :close_reason_id ]
+    end
+  end
+
+  class User < ActiveRecord::Base
+    def self.exported_fields
+      return [ :id, :reputation, :creation_date, :display_name, :email_hash,
+               :last_access_date, :website_url, :location, :age, :about_me, 
+               :views, :up_votes, :down_votes ]
+    end
+  end
+
+  class Vote < ActiveRecord::Base
+    def self.exported_fields
+      return [ :id, :post_id, :vote_type_id, :creation_date, :user_id, 
+               :bounty_amount ]
+    end
+  end
+
+  class PostType < ActiveRecord::Base
+  end
+
+  class PostHistoryType < ActiveRecord::Base
+  end
+
+  class CloseReason < ActiveRecord::Base
+  end
+
+  class VoteType < ActiveRecord::Base
+  end
+
+  # Infrastructure.  Do not call this from your code.
+  class Lookup
+
+    @@map = { "badges" => :Badge, "comments" => :Comment, 
+              "posthistory" => :PostHistory, "posts" => :Post, "users" => :User,
+              "votes" => :Vote }
+
+    def self.find_class(file_name)
+      Object.const_get("SO2DB").const_get("Models")
+      .const_get(@@map[file_name].to_s)
+    end
+
+    def self.get_required_attrs(file_name)
+      raw = find_class(file_name).send :exported_fields
+      return raw.map {|f| f.to_s.camelize.sub(/Guid/, 'GUID')}
+    end
+
+  end
+
+end

data/lib/so2db.rb

@@ -0,0 +1,126 @@
+#--
+# Copyright (c) 2012 Chad Taylor
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy 
+# of this software and associated documentation files (the "Software"), to deal 
+# in the Software without restriction, including without limitation the rights 
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
+# copies of the Software, and to permit persons to whom the Software is 
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in 
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
+# SOFTWARE.
+
+module SO2DB
+
+  # Base class for StackOverflow data importers.  Drives database setup and
+  # data importing files from a directory.
+  #
+  # Implementations of this class must provide a method with the following 
+  # signature:
+  #
+  #   import_stream(formatter)
+  #
+  # This method may be private.  The purpose of this method is to actually
+  # perform the data import with data from the provided formatter.  The
+  # formatter is provided to support scenarios of streaming data to STDIN 
+  # (e.g., PostgreSQL's COPY command) as well as pushing data to a file before 
+  # import (e.g., for MySQL's mysqlimport utility).  It has type 
+  # SO2DB::Formatter.
+  #
+  # The importer uses ActiveRecord for table creation and Foreigner for creating
+  # table relationships. You are limited to the databases supported by these 
+  # libraries.  In addition, a 'uuid' method must be avaiable to the adapter
+  # provided to ActiveRecord.  (See so2pg for an example of an adapter extension
+  # that provides the method.)
+  #
+  # In addition, it provides two accessors for subclasses:
+  #
+  #   attr_reader :conn_opts
+  #   attr_accessor :delimiter
+  #
+  # The conn_opts property provides the ActiveRecord connection data (e.g., 
+  # :database, :host, etc.).  The delimiter property sets the delimiter used by
+  # the formatter.  The delimiter is \v (0xB) by default.
+  class Importer
+
+    # Initializes the importer.
+    #
+    # Arguments:
+    #   relations: (Boolean) Indicates whether database relationships should
+    #                        be created.
+    #   optionals: (Boolean) Indicates whether optional database tables and
+    #                        content should be created.
+    #   adapter:   (String)  The ActiveRecord adapter name (e.g., 'postgresql').
+    #   options:   (Hash)    The database connection options, as required by
+    #                        ActiveRecord for the provided adapter.
+    def initialize(relations = false, optionals = false, adapter = '', options = {})
+      @relations = relations
+      @optionals = optionals
+      @conn_opts = options.merge( { :adapter => adapter } )
+      @format_delimiter = 11.chr.to_s
+    end
+
+    # Creates the database tables and relationships, and imports the data in
+    # the files in the specified directory.  
+    #
+    # Arguments:
+    #   dir:  (String) The directory path containting the StackOverflow data
+    #                  dump XML files (e.g., badges.xml, posts.xml, etc.).
+    def import(dir)
+      setup
+      create_basics
+      import_data(dir)
+      create_relations if @relations
+      create_optionals if @optionals
+      create_optional_relations if @relations and @optionals
+    end
+
+    private
+
+    attr_reader :conn_opts
+    attr_accessor :format_delimiter
+
+    def setup
+      ActiveRecord::Base.establish_connection @conn_opts
+      Foreigner.load
+    end
+
+    def create_basics
+      SO2DB::CreateBasicTables.new.up
+    end
+
+    def import_data(dir)
+      files = Dir.entries(dir).delete_if { |x| !x.end_with? 'xml' }
+      files.each do |f|
+        f = Formatter.new(File.join(dir, f), @format_delimiter)
+        import_stream f
+      end
+    end
+
+    def create_relations
+      SO2DB::CreateRelationships.new.up
+    end
+
+    def create_optionals
+      SO2DB::CreateOptionals.new.up
+    end
+
+    def create_optional_relations
+      SO2DB::CreateOptionalRelationships.new.up
+    end
+
+  end
+end
+
+require 'so2db/formatter'
+require 'so2db/migrations'
+require 'so2db/models'

data/lib/so2pg.rb

@@ -0,0 +1,179 @@
+#--
+# Copyright (c) 2012 Chad Taylor
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy 
+# of this software and associated documentation files (the "Software"), to deal 
+# in the Software without restriction, including without limitation the rights 
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
+# copies of the Software, and to permit persons to whom the Software is 
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in 
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
+# SOFTWARE.
+
+require 'optparse'
+require 'so2db'
+
+module ActiveRecord::ConnectionAdapters
+  class PostgreSQLAdapter < AbstractAdapter
+
+    # Extends the adapter to include support for a uuid type.  This is required
+    # by the importer (see SO2DB::Importer for more information).  For 
+    # PostgreSQL, simply use the native 'uuid' type (for MySQL, use something
+    # a bit more contrived, like CHAR(16)).
+    def uuid
+      return 'uuid', {}
+    end
+  end
+end
+
+# Imports the StackOverflow data into a PostgreSQL data.
+class PgImporter < SO2DB::Importer
+
+  # (See SO2DB::Importer.initialize documentation)
+  def initialize(relations = false, optionals = false, options = {})
+    super(relations, optionals, "postgresql", options)
+  end
+
+  def self.import_from_argv(argv)
+    # Parse the command-line options
+    cmd_opts = PgOptionsParser.parse(ARGV)
+
+    # If all validation passed, then execute the import!
+    if cmd_opts
+      start = Time.now
+      pg = PgImporter.new(cmd_opts.has_key?(:relationships),
+                          cmd_opts.has_key?(:optionals),
+                          cmd_opts)
+      pg.import(cmd_opts[:dir])
+      puts "Import completed in #{Time.now - start}s"
+    end
+  end
+
+  private
+
+  # Imports the data from the formatter into the PostgreSQL database.
+  #
+  # Note that what follows is just one way to implement the importer.  You
+  # could just as easily push the formatted data into a file and then ask
+  # the database to suck that file in.
+  def import_stream(formatter)
+    puts "Importing file #{formatter.file_name}..."
+    start = Time.now
+
+    sql = build_sql(formatter.value_str)
+    cmd = build_cmd(sql)
+    execute_cmd(cmd, formatter)
+
+    puts "   -> #{Time.now - start}s"
+  end
+
+  # Builds the SQL command used for bulk loading the tables.
+  def build_sql(value_str)
+    "COPY #{value_str} FROM STDIN WITH (FORMAT csv, DELIMITER E'\x0B')"
+  end
+
+  # Builds the import command with the given SQL command and the global
+  # connection options.
+  #   
+  # Example:
+  #   >> sql = "COPY ..."
+  #   >> puts build_cmd(sql)
+  #   => psql -d test -h localhost -c "COPY ..."
+  def build_cmd(sql)
+    # Only exists within the context of this script (not exported), so this
+    # does not degrade security posture after the script has completed
+    ENV['PGPASSWORD'] = conn_opts[:password] if conn_opts.has_key? :password
+
+    cmd = "psql"
+    cmd << " -d #{conn_opts[:database]}" if conn_opts.has_key? :database
+    cmd << " -h #{conn_opts[:host]}" if conn_opts.has_key? :host
+    cmd << " -U #{conn_opts[:username]}" if conn_opts.has_key? :username
+    cmd << " -p #{conn_opts[:port]}" if conn_opts.has_key? :port
+    cmd << " -c \"#{sql}\""
+
+    return cmd
+  end
+
+  # Executes the provided shell command and pumps the data from the formatter
+  # to it over stdin.
+  def execute_cmd(cmd, formatter)
+    IO.popen(cmd, 'r+') do |s|
+      formatter.format(s)
+      s.close_write
+    end
+  end
+
+end
+
+class PgOptionsParser
+
+  # Parses the command-line arguments into a Hash object.  Note that the members
+  # of the Hash have the same name as the ActiveRecord parameters (e.g., :host,
+  # :database, etc.).  This Hash will actually be passed to ActiveRecord for
+  # consumption.
+  def self.parse(args)
+    options = {}
+
+    opts = OptionParser.new do |opts|
+      opts.banner = <<-EOB
+Imports a StackOverflow data dump into a PostgreSQL database.
+Usage: so2pg [options]
+      EOB
+
+      opts.on("-H", "--host HOST", "The database host") do |host|
+        options[:host] = host
+      end
+
+      opts.on("-d", "--database DBNAME", "The name of the database (REQUIRED)") do |dbname|
+        options[:database] = dbname
+      end
+
+      opts.on("-D", "--directory DIRECTORY", "The data directory path (REQUIRED)") do |dir|
+        options[:dir] = dir
+      end
+
+      opts.on("-u", "--user USER", "The user name") do |user|
+        options[:username] = user
+      end
+
+      opts.on("-p", "--password PASSWORD", "The user's password") do |password|
+        options[:password] = password
+      end
+
+      opts.on("-P", "--port PORT_NUMBER", "The port number") do |port|
+        options[:port] = port
+      end
+
+      opts.on("-O", "--include-optionals", "Includes optional tables") do
+        options[:optionals] = true
+      end
+
+      opts.on("-R", "--include-relationships", "Includes table relationships") do
+        options[:relationships] = true
+      end
+
+      opts.on("-h", "--help", "Show this help screen") do |help|
+        options[:help] = true
+      end
+
+    end
+
+    opts.parse!(args)
+    if(options[:help] or !options.has_key? :dir or !options.has_key? :database)
+      puts opts.help
+      nil
+    else
+      options
+    end
+  end
+
+end

data/test/test_formatter.rb

@@ -0,0 +1,120 @@
+require 'test/unit'
+require 'mocha'
+require 'so2db'
+
+class FormatterTest < Test::Unit::TestCase
+  include Rake::DSL
+
+  def setup
+    @formatter = SO2DB::Formatter.new("/tmp/badges.xml")
+  end
+
+  def test_initializer_with_default_args
+    f = SO2DB::Formatter.new
+    assert_equal '', f.instance_variable_get(:@path)
+    assert_equal 11.chr.to_s, f.instance_variable_get(:@delimiter)
+  end
+
+  def test_initializer_with_provided_args
+    path = '/my/test/path'
+    delimiter = 12.chr.to_s
+    f = SO2DB::Formatter.new(path, delimiter)
+
+    assert_equal path, f.instance_variable_get(:@path)
+    assert_equal delimiter, f.instance_variable_get(:@delimiter)
+  end
+
+  def test_format
+    file = "file"
+    outstream = "outstream"
+    attrs = [ :a, :b, :c ]
+    SO2DB::Models::Lookup.expects(:get_required_attrs).with("badges").once.returns(attrs)
+    File.expects(:open).with("/tmp/badges.xml").once.returns(file)
+    @formatter.expects(:format_from_stream).with(file, attrs, outstream).once.returns("x")
+
+    result = @formatter.format(outstream)
+    assert_equal "x", result
+  end
+
+  def test_format_from_stream
+    x = <<-eoxml
+<data>
+<fake-row id="3"/>
+<row Id="1" UserId="2" Name="Autobiographer" Date="2010-07-20T19:07:22.990" />
+</data>
+    eoxml
+
+    r, w = IO.pipe
+
+    arr = [ "Id", "UserId", "Name", "Date", "Missing" ]
+    @formatter.send(:format_from_stream, x, arr, w)
+
+    values = [ '2010-07-20T19:07:22.990', '1', '', 'Autobiographer', '2' ]
+    expected = values.join(11.chr.to_s) << "\n"
+    actual = r.gets
+
+    assert_equal expected, actual
+  end
+
+  def test_file_name
+    assert_equal "badges.xml", @formatter.file_name
+  end
+
+  def test_value_str
+    assert_equal "badges(date,id,name,user_id)", @formatter.value_str
+  end
+
+  def create_node_stub(name, node_type)
+    obj = mock()
+    obj.stubs(:name).returns(name)
+    obj.stubs(:node_type).returns(node_type)
+
+    obj
+  end
+
+  def test_element_start_with_good_values
+    node = create_node_stub("row", Nokogiri::XML::Reader::TYPE_ELEMENT)
+    assert @formatter.send(:element_start?, node)
+  end
+
+  def test_element_start_with_invalid_type
+    node = create_node_stub("row", Nokogiri::XML::Reader::TYPE_END_ELEMENT)
+    assert_equal false, @formatter.send(:element_start?, node)
+  end
+
+  def test_element_start_with_invalid_name
+    node = create_node_stub("badges", Nokogiri::XML::Reader::TYPE_ELEMENT)
+    assert_equal false, @formatter.send(:element_start?, node)
+  end
+
+  def test_format_node
+    node = mock()
+    node.stubs(:attribute).with("Id").returns("1")
+    node.stubs(:attribute).with("Name").returns("Anony Mous")
+
+    @formatter.expects(:scrub).with("1").once.returns("1")
+    @formatter.expects(:scrub).with("Anony Mous").once.returns("Anony Mous")
+
+    result = @formatter.send(:format_node, node, [ "Id", "Name" ])
+
+    assert_equal "1\vAnony Mous", result
+  end
+
+  def test_format_node_with_missing_attribute
+    node = mock()
+    node.stubs(:attribute).with("Id").returns("1")
+    node.stubs(:attribute).with("Name").returns(nil)
+
+    @formatter.expects(:scrub).once.with("1").returns("1")
+
+    result = @formatter.send(:format_node, node, [ "Id", "Name" ])
+
+    assert_equal "1\v", result
+  end
+
+  def test_scrub
+    assert_equal '&lt;asdffdsa&gt;', @formatter.send(:scrub, "<asdf\nfdsa\r>")
+  end
+
+end
+

data/test/test_models.rb

@@ -0,0 +1,20 @@
+require 'test/unit'
+require 'so2db'
+
+class LookupTest < Test::Unit::TestCase
+  include Rake::DSL
+
+  def test_lookup_badges
+    assert_equal SO2DB::Models::Badge, SO2DB::Models::Lookup::find_class("badges")
+  end
+
+  def test_required_attrs_badges
+    attrs = SO2DB::Models::Lookup::get_required_attrs("badges")
+    assert_equal [ "Id", "UserId", "Name", "Date" ], attrs
+  end
+
+  def test_guid_capitalization
+    attrs = SO2DB::Models::Lookup::get_required_attrs("posthistory")
+    assert_block { attrs.include? 'RevisionGUID' }
+  end
+end

data/test/test_so2db.rb

@@ -0,0 +1,17 @@
+require 'test/unit'
+require 'mocha'
+require 'so2db'
+
+class ImporterTest < Test::Unit::TestCase
+
+  def test_import_data
+    importer = SO2DB::Importer.new
+    Dir.expects(:entries).once.with('/tmp').returns([ 'test.bak', 'test.xml' ])
+    SO2DB::Formatter.expects(:new).once
+              .with('/tmp/test.xml', 11.chr.to_s).returns('formatter')
+    importer.expects(:import_stream).once.with('formatter')
+
+    importer.send(:import_data, '/tmp')
+  end
+
+end

data/test/test_so2pg.rb

@@ -0,0 +1,120 @@
+require 'test/unit'
+require 'mocha'
+require 'so2pg'
+
+class PgImporterTest < Test::Unit::TestCase
+
+  def setup
+    @importer = PgImporter.new(true, true, { :database => "dbname", :dir => "dir" })
+  end
+
+  def test_import_stream
+    formatter = mock()
+    formatter.stubs(:file_name).returns("file_name")
+    formatter.expects(:value_str).once.returns('badges(id,name)')
+
+    @importer.expects(:build_sql).once.with('badges(id,name)').returns("COPY...")
+    @importer.expects(:build_cmd).once.with("COPY...").returns("cmd")
+    @importer.stubs(:execute_cmd).returns('')
+    @importer.expects(:execute_cmd).once.with("cmd", formatter).returns('')
+
+    $stdout.stubs(:puts).returns('')
+
+    @importer.send(:import_stream, formatter)
+  end
+
+  def test_build_sql
+    exp = "COPY badges(id,name) FROM STDIN WITH (FORMAT csv, DELIMITER E'\x0B')"
+    assert_equal exp, @importer.send(:build_sql, "badges(id,name)")
+  end
+
+  def test_build_cmd
+    expected = "psql -d dbname -c \"COPY...\""
+    actual = @importer.send(:build_cmd, "COPY...")
+
+    assert_equal expected, actual
+  end
+
+  def test_build_cmd_sets_env_password
+    importer = PgImporter.new(true, true, { :database => "dbname",
+                                            :dir => "dir",
+                                            :password => "asdf1234" })
+
+    importer.send(:build_cmd, "COPY...")
+    assert_equal "asdf1234", ENV['PGPASSWORD']
+  end
+
+  def test_execute_cmd
+    strm = mock()
+    formatter = mock()
+
+    IO.expects(:popen).once.with("cmd", "r+").yields(strm)
+    formatter.expects(:format).once.with(strm)
+    strm.expects(:close_write).once
+
+    @importer.send(:execute_cmd, "cmd", formatter)
+  end
+
+end
+
+class PgOptionsParserTest < Test::Unit::TestCase
+
+  def test_all_options
+    host = 'localhost'
+    database = 'database_name'
+    directory = 'data_directory'
+    user = 'anony'
+    password = 'mous'
+    port = '1234'
+
+    cmd = [ '-H', host, '-d', database, '-D', directory, '-u', user, 
+            '-p', password, '-P', port, '-O', '-R' ]
+
+    options = PgOptionsParser.parse(cmd)
+
+    assert_equal host, options[:host]
+    assert_equal database, options[:database]
+    assert_equal directory, options[:dir]
+    assert_equal user, options[:username]
+    assert_equal password, options[:password]
+    assert_equal port, options[:port]
+    assert options[:optionals]
+    assert options[:relationships]
+  end
+
+  def assert_help_displayed(cmd)
+    $stdout.expects(:puts).returns('')
+    options = PgOptionsParser.parse(cmd)
+
+    assert_nil options
+  end
+
+  def test_options_without_database
+    assert_help_displayed [ '-D', 'data_dir' ]
+  end
+
+  def test_options_without_data_dir
+    assert_help_displayed [ '-d', 'database_name' ]
+  end
+
+  def test_options_with_help
+    assert_help_displayed [ '-D', 'data_dir', '-d', 'database_name', '-h' ]
+  end
+
+  def test_options_without_optionals
+    cmd = [ '-D', 'data_dir', '-d', 'database_name', '-R' ]
+    options = PgOptionsParser.parse(cmd)
+
+    assert options[:relationships]
+    assert !options[:optionals]
+  end
+
+  def test_options_without_relationships
+    cmd = [ '-D', 'data_dir', '-d', 'database_name', '-O' ]
+    options = PgOptionsParser.parse(cmd)
+
+    assert !options[:relationships]
+    assert options[:optionals]
+  end
+
+end