mosql 0.1.0

data/lib/mosql/log.rb

@@ -0,0 +1,7 @@
+module MoSQL
+  module Logging
+    def log
+      @@logger ||= Log4r::Logger.new("Stripe::MoSQL")
+    end
+  end
+end

data/lib/mosql/schema.rb

@@ -0,0 +1,149 @@
+module MoSQL
+  class SchemaError < StandardError; end;
+
+  class Schema
+    include MoSQL::Logging
+
+    def to_ordered_hash(lst)
+      hash = BSON::OrderedHash.new
+      lst.each do |ent|
+        raise "Invalid ordered hash entry #{ent.inspect}" unless ent.is_a?(Hash) && ent.keys.length == 1
+        field, type = ent.first
+        hash[field] = type
+      end
+      hash
+    end
+
+    def parse_spec(spec)
+      out = spec.dup
+      out[:columns] = to_ordered_hash(spec[:columns])
+      out
+    end
+
+    def initialize(map)
+      @map = {}
+      map.each do |dbname, db|
+        @map[dbname] ||= {}
+        db.each do |cname, spec|
+          @map[dbname][cname] = parse_spec(spec)
+        end
+      end
+    end
+
+    def create_schema(db, clobber=false)
+      @map.values.map(&:values).flatten.each do |collection|
+        meta = collection[:meta]
+        log.info("Creating table '#{meta[:table]}'...")
+        db.send(clobber ? :create_table! : :create_table?, meta[:table]) do
+          collection[:columns].each do |field, type|
+            column field, type
+          end
+          if meta[:extra_props]
+            column '_extra_props', 'TEXT'
+          end
+          primary_key [:_id]
+        end
+      end
+    end
+
+    def find_ns(ns)
+      db, collection = ns.split(".")
+      schema = (@map[db] || {})[collection]
+      if schema.nil?
+        log.debug("No mapping for ns: #{ns}")
+        return nil
+      end
+      schema
+    end
+
+    def find_ns!(ns)
+      schema = find_ns(ns)
+      raise SchemaError.new("No mapping for namespace: #{ns}") if schema.nil?
+      schema
+    end
+
+    def transform(ns, obj, schema=nil)
+      schema ||= find_ns!(ns)
+
+      obj = obj.dup
+      row = []
+      schema[:columns].each do |name, type|
+        v = obj.delete(name)
+        case v
+        when BSON::Binary, BSON::ObjectId
+          v = v.to_s
+        end
+        row << v
+      end
+
+      if schema[:meta][:extra_props]
+        # Kludgily delete binary blobs from _extra_props -- they may
+        # contain invalid UTF-8, which to_json will not properly encode.
+        obj.each do |k,v|
+          obj.delete(k) if v.is_a?(BSON::Binary)
+        end
+        row << obj.to_json
+      end
+
+      log.debug { "Transformed: #{row.inspect}" }
+
+      row
+    end
+
+    def all_columns(schema)
+      cols = schema[:columns].keys
+      if schema[:meta][:extra_props]
+        cols << "_extra_props"
+      end
+      cols
+    end
+
+    def copy_data(db, ns, objs)
+      schema = find_ns!(ns)
+      data = objs.map { |o| transform_to_copy(ns, o, schema) }.join("\n")
+      db.synchronize do |pg|
+        sql = "COPY \"#{schema[:meta][:table]}\" " +
+          "(#{all_columns(schema).map {|c| "\"#{c}\""}.join(",")}) FROM STDIN"
+        pg.execute(sql)
+        objs.each do |o|
+          pg.put_copy_data(transform_to_copy(ns, o, schema) + "\n")
+        end
+        pg.put_copy_end
+        begin
+          pg.get_result.check
+        rescue PGError => e
+          db.send(:raise_error, e)
+        end
+      end
+    end
+
+    def quote_copy(val)
+      case val
+      when nil
+        "\\N"
+      when true
+        't'
+      when false
+        'f'
+      else
+        val.to_s.gsub(/([\\\t\n\r])/, '\\\\\\1')
+      end
+    end
+
+    def transform_to_copy(ns, row, schema=nil)
+      row.map { |c| quote_copy(c) }.join("\t")
+    end
+
+    def table_for_ns(ns)
+      find_ns!(ns)[:meta][:table]
+    end
+
+    def all_mongo_dbs
+      @map.keys
+    end
+
+    def collections_for_mongo_db(db)
+      (@map[db]||{}).keys
+    end
+  end
+end

data/lib/mosql/sql.rb

@@ -0,0 +1,59 @@
+module MoSQL
+  class SQLAdapter
+    include MoSQL::Logging
+
+    attr_reader :db
+
+    def initialize(schema, uri, pgschema=nil)
+      @schema = schema
+      connect_db(uri, pgschema)
+    end
+
+    def connect_db(uri, pgschema)
+      @db = Sequel.connect(uri, :after_connect => proc do |conn|
+                             if pgschema
+                               begin
+                                 conn.execute("CREATE SCHEMA \"#{pgschema}\"")
+                               rescue PG::Error
+                               end
+                               conn.execute("SET search_path TO \"#{pgschema}\"")
+                             end
+                           end)
+    end
+
+    def table_for_ns(ns)
+      @db[@schema.table_for_ns(ns).intern]
+    end
+
+    def upsert_ns(ns, obj)
+      h = {}
+      cols = @schema.all_columns(@schema.find_ns(ns))
+      row  = @schema.transform(ns, obj)
+      cols.zip(row).each { |k,v| h[k] = v }
+      upsert(table_for_ns(ns), h)
+    end
+
+    def upsert(table, item)
+      begin
+        upsert!(table, item)
+      rescue Sequel::DatabaseError => e
+        wrapped = e.wrapped_exception
+        if wrapped.result
+          log.warn("Ignoring row (_id=#{item['_id']}): #{e}")
+        else
+          raise e
+        end
+      end
+    end
+
+    def upsert!(table, item)
+      begin
+        table.insert(item)
+      rescue Sequel::DatabaseError => e
+        raise e unless e.message =~ /duplicate key value violates unique constraint/
+        table.where(:_id => item['_id']).update(item)
+      end
+    end
+  end
+end
+

data/lib/mosql/tailer.rb

@@ -0,0 +1,36 @@
+module MoSQL
+  class Tailer < Mongoriver::AbstractPersistentTailer
+    def self.create_table(db, tablename)
+      db.create_table?(tablename) do
+        column :service,   'TEXT'
+        column :timestamp, 'INTEGER'
+        primary_key [:service]
+      end
+      db[tablename.to_sym]
+    end
+
+    def initialize(backends, type, table, opts)
+      super(backends, type, opts)
+      @table   = table
+      @service = opts[:service] || "mosql"
+    end
+
+    def read_timestamp
+      row = @table.where(:service => @service).select([:timestamp]).first
+      if row
+        BSON::Timestamp.new(row[:timestamp], 0)
+      else
+        BSON::Timestamp.new(0, 0)
+      end
+    end
+
+    def write_timestamp(ts)
+      begin
+        @table.insert({:service => @service, :timestamp => ts.seconds})
+      rescue Sequel::DatabaseError => e
+        raise unless e.message =~ /duplicate key value violates unique constraint/
+        @table.where(:service => @service).update(:timestamp => ts.seconds)
+      end
+    end
+  end
+end

data/lib/mosql/version.rb

@@ -0,0 +1,3 @@
+module MoSQL
+  VERSION = "0.1.0"
+end

data/mosql.gemspec

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+$:.unshift(File.expand_path("lib", File.dirname(__FILE__)))
+require 'mosql/version'
+
+Gem::Specification.new do |gem|
+  gem.authors       = ["Nelson Elhage"]
+  gem.email         = ["nelhage@stripe.com"]
+  gem.description   = %q{A library for streaming MongoDB to SQL}
+  gem.summary       = %q{MongoDB -> SQL streaming bridge}
+  gem.homepage      = "https://github.com/stripe/mosql"
+
+  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.name          = "mosql"
+  gem.require_paths = ["lib"]
+  gem.version       = MoSQL::VERSION
+
+  %w[sequel pg mongo bson_ext rake log4r json
+     mongoriver].each { |dep| gem.add_runtime_dependency(dep) }
+
+  gem.add_development_dependency "minitest"
+  gem.add_development_dependency "mocha"
+end

data/test/_lib.rb

@@ -0,0 +1,18 @@
+require 'rubygems'
+require 'bundler/setup'
+
+require 'minitest/autorun'
+require 'minitest/spec'
+require 'mocha'
+
+$:.unshift(File.expand_path(File.join(File.dirname(__FILE__), '../lib')))
+
+require 'mosql'
+
+module MoSQL
+  class Test < ::MiniTest::Spec
+    def setup
+      # Put any stubs here that you want to apply globally
+    end
+  end
+end

data/test/functional/_lib.rb

@@ -0,0 +1,65 @@
+require File.join(File.dirname(__FILE__), "../_lib")
+
+module MoSQL
+  class Test::Functional < MoSQL::Test
+    attr_reader :sequel, :mongo
+
+    def sql_test_uri
+      ENV['MONGOSQL_TEST_SQL'] || 'postgres:///test'
+    end
+    def mongo_test_uri
+      ENV['MONGOSQL_TEST_MONGO'] || 'mongodb://localhost'
+    end
+    def mongo_test_dbname
+      ENV['MONGOSQL_TEST_MONGO_DB'] || 'test'
+    end
+
+    def connect_sql
+      begin
+        conn = Sequel.connect(sql_test_uri)
+        conn.test_connection
+        conn
+      rescue Sequel::DatabaseConnectionError
+        $stderr.puts <<EOF
+
+*********************************************************************
+Unable to connect to PostgreSQL database at `#{sql_test_uri}'. Either
+configure a PostgresSQL database running locally without
+authentication with a 'test' database, or set \$MONGOSQL_TEST_SQL in
+the environment.
+*********************************************************************
+
+EOF
+        exit(1)
+      end
+    end
+
+    def connect_mongo
+      begin
+        Mongo::Connection.from_uri(mongo_test_uri)
+      rescue Mongo::ConnectionFailure, Mongo::ConnectionError
+        $stderr.puts <<EOF
+
+*********************************************************************
+Unable to connect to MongoDB database at `#{mongo_test_uri}'. Either
+configure a MongoDB database running on localhost without
+authentication with a 'test' database, or set \$MONGOSQL_TEST_MONGO in
+the environment.
+*********************************************************************
+
+EOF
+        exit(1)
+      end
+    end
+
+    def mongo_db
+      mongo.db(mongo_test_dbname)
+    end
+
+    def setup
+      @sequel = connect_sql
+      @mongo  = connect_mongo
+      super
+    end
+  end
+end

data/test/functional/functional.rb

@@ -0,0 +1,7 @@
+require File.join(File.dirname(__FILE__), '_lib.rb')
+
+class MoSQL::Test::Functional::SchemaTest < MoSQL::Test::Functional
+  it 'runs tests' do
+    assert(true)
+  end
+end

data/test/functional/schema.rb

@@ -0,0 +1,66 @@
+require File.join(File.dirname(__FILE__), '_lib.rb')
+
+class MoSQL::Test::Functional::SchemaTest < MoSQL::Test::Functional
+  TEST_MAP = <<EOF
+---
+db:
+  collection:
+    :meta:
+      :table: sqltable
+    :columns:
+      - _id: TEXT
+      - var: INTEGER
+  with_extra_props:
+    :meta:
+      :table: sqltable2
+      :extra_props: true
+    :columns:
+      - _id: INTEGER
+EOF
+
+  before do
+    @map = MoSQL::Schema.new(YAML.load(TEST_MAP))
+
+    @sequel.drop_table?(:sqltable)
+    @sequel.drop_table?(:sqltable2)
+    @map.create_schema(@sequel)
+  end
+
+  def table; @sequel[:sqltable]; end
+  def table2; @sequel[:sqltable2]; end
+
+  it 'Creates the tables with the right columns' do
+    assert_equal(Set.new([:_id, :var]),
+                 Set.new(table.columns))
+    assert_equal(Set.new([:_id, :_extra_props]),
+                 Set.new(table2.columns))
+  end
+
+  it 'Can COPY data' do
+    objects = [
+               {'_id' => "a", 'var' => 0},
+               {'_id' => "b", 'var' => 1},
+               {'_id' => "c"},
+               {'_id' => "d", 'other_var' => "hello"}
+              ]
+    @map.copy_data(@sequel, 'db.collection', objects.map { |o| @map.transform('db.collection', o) } )
+    assert_equal(4, table.count)
+    rows = table.select.sort_by { |r| r[:_id] }
+    assert_equal(%w[a b c d], rows.map { |r| r[:_id] })
+    assert_equal(nil, rows[2][:var])
+    assert_equal(nil, rows[3][:var])
+  end
+
+  it 'Can COPY BSON::ObjectIDs' do
+    o = {'_id' => BSON::ObjectId.new, 'var' => 0}
+    @map.copy_data(@sequel, 'db.collection', [ @map.transform('db.collection', o)] )
+    assert_equal(o['_id'].to_s, table.select.first[:_id])
+  end
+
+  it 'Can transform BSON::ObjectIDs' do
+    o = {'_id' => BSON::ObjectId.new, 'var' => 0}
+    row = @map.transform('db.collection', o)
+    table.insert(row)
+    assert_equal(o['_id'].to_s, table.select.first[:_id])
+  end
+end