jruby-pg 0.1-java

data/sample/warehouse_partitions.rb

@@ -0,0 +1,320 @@
+#!/usr/bin/env ruby
+# vim: set nosta noet ts=4 sw=4:
+#
+# Script to automatically move partitioned tables and their indexes
+# to a separate area on disk.
+#
+# Mahlon E. Smith <mahlon@martini.nu>
+#
+# Example use case:
+#
+#   - You've got a heavy insert table, such as syslog data.
+#   - This table has a partitioning trigger (or is manually partitioned)
+#     by date, to separate incoming stuff from archival/report stuff.
+#   - You have a tablespace on cheap or slower disk (maybe even
+#     ZFS compressed, or some such!)
+#
+# The only assumption this script makes is that your tables are dated, and
+# the tablespace they're moving into already exists.
+#
+# A full example, using the syslog idea from above, where each child
+# table is date partitioned by a convention of "syslog_YEAR-WEEKOFYEAR":
+#
+#    syslog             # <--- parent
+#    syslog_2012_06     # <--- inherited
+#    syslog_2012_07     # <--- inherited
+#    syslog_2012_08     # <--- inherited
+#    ...
+#
+# You'd run this script like so:
+#
+#    ./warehouse_partitions.rb -F syslog_%Y_%U
+#
+# Assuming this was week 12 of the year, tables syslog_2012_06 through
+# syslog_2012_11 would start sequentially migrating into the tablespace
+# called 'warehouse'.
+#
+
+
+begin
+	require 'date'
+	require 'ostruct'
+	require 'optparse'
+	require 'pathname'
+	require 'etc'
+	require 'pg'
+
+rescue LoadError # 1.8 support
+	unless Object.const_defined?( :Gem )
+		require 'rubygems'
+		retry
+	end
+	raise
+end
+
+
+### A tablespace migration class.
+###
+class PGWarehouse
+
+	def initialize( opts )
+		@opts = opts
+		@db = PG.connect(
+			:dbname   => opts.database,
+			:host     => opts.host,
+			:port     => opts.port,
+			:user     => opts.user,
+			:password => opts.pass,
+			:sslmode  => 'prefer'
+		)
+		@db.exec "SET search_path TO %s" % [ opts.schema ] if opts.schema
+
+		@relations = self.relations
+	end
+
+	attr_reader :db
+
+	######
+	public
+	######
+
+	### Perform the tablespace moves.
+	###
+	def migrate
+		if @relations.empty?
+			$stderr.puts 'No tables were found for warehousing.'
+			return
+		end
+
+		$stderr.puts "Found %d relation%s to move." % [ relations.length, relations.length == 1 ? '' : 's' ]
+		@relations.sort_by{|_,v| v[:name] }.each do |_, val|
+			$stderr.print "  - Moving table '%s' to '%s'... "  % [
+				val[:name], @opts.tablespace
+			]
+
+			if @opts.dryrun
+				$stderr.puts '(not really)'
+
+			else
+				age = self.timer do
+					db.exec "ALTER TABLE %s SET TABLESPACE %s;" % [
+						val[:name], @opts.tablespace
+					]
+				end
+				puts age
+			end
+
+			val[ :indexes ].each do |idx|
+				$stderr.print "      - Moving index '%s' to '%s'... "  % [
+					idx, @opts.tablespace
+				]
+				if @opts.dryrun
+					$stderr.puts '(not really)'
+
+				else
+					age = self.timer do
+						db.exec "ALTER INDEX %s SET TABLESPACE %s;" % [
+							idx, @opts.tablespace
+						]
+					end
+					puts age
+				end
+			end
+		end
+	end
+
+
+	#########
+	protected
+	#########
+
+	### Get OIDs and current tablespaces for everything under the
+	### specified schema.
+	###
+	def relations
+		return @relations if @relations
+		relations = {}
+
+		query =  %q{
+			SELECT c.oid AS oid,
+				c.relname AS name,
+				c.relkind AS kind,
+				t.spcname AS tspace
+			FROM pg_class AS c
+			LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
+			LEFT JOIN pg_tablespace t ON t.oid = c.reltablespace
+			WHERE c.relkind = 'r' }
+		query << "AND n.nspname='#{@opts.schema}'" if @opts.schema
+
+		# Get the relations list, along with each element's current tablespace.
+		#
+		self.db.exec( query ) do |res|
+			res.each do |row|
+				relations[ row['oid'] ] = {
+					:name       => row['name'],
+					:tablespace => row['tspace'],
+					:indexes    => [],
+					:parent     => nil
+				}
+			end
+		end
+
+		# Add table inheritence information.
+		#
+		db.exec 'SELECT inhrelid AS oid, inhparent AS parent FROM pg_inherits' do |res|
+			res.each do |row|
+				relations[ row['oid'] ][ :parent ] = row['parent']
+			end
+		end
+
+		# Remove tables that don't qualify for warehousing.
+		#
+		#   - Tables that are not children of a parent
+		#   - Tables that are already in the warehouse tablespace
+		#   - The currently active child (it's likely being written to!)
+		#   - Any table that can't be parsed into the specified format
+		#
+		relations.reject! do |oid, val|
+			begin
+				val[:parent].nil? ||
+				val[:tablespace] == @opts.tablespace ||
+				val[:name] == Time.now.strftime( @opts.format ) ||
+				! DateTime.strptime( val[:name], @opts.format )
+			rescue ArgumentError
+				true
+			end
+		end
+
+		query = %q{
+			SELECT c.oid AS oid,
+				i.indexname AS name
+			FROM pg_class AS c
+			INNER JOIN pg_indexes AS i
+				ON i.tablename = c.relname }
+		query << "AND i.schemaname='#{@opts.schema}'" if @opts.schema
+
+		# Attach index names to tables.
+		#
+		db.exec( query ) do |res|
+			res.each do |row|
+				relations[ row['oid'] ][ :indexes ] << row['name'] if relations[ row['oid'] ]
+			end
+		end
+
+		return relations
+	end
+
+
+	### Wrap arbitrary commands in a human readable timer.
+	###
+	def timer
+		start = Time.now
+		yield
+		age = Time.now - start
+
+		diff = age
+		secs = diff % 60
+		diff = ( diff - secs ) / 60
+		mins = diff % 60
+		diff = ( diff - mins ) / 60
+		hour = diff % 24
+
+		return "%02d:%02d:%02d" % [ hour, mins, secs ]
+	end
+end
+
+
+### Parse command line arguments.  Return a struct of global options.
+###
+def parse_args( args )
+	options          = OpenStruct.new
+	options.database = Etc.getpwuid( Process.uid ).name
+	options.host     = '127.0.0.1'
+	options.port     = 5432
+	options.user     = Etc.getpwuid( Process.uid ).name
+	options.sslmode  = 'prefer'
+	options.tablespace = 'warehouse'
+
+	opts = OptionParser.new do |opts|
+		opts.banner = "Usage: #{$0} [options]"
+
+		opts.separator ''
+		opts.separator 'Connection options:'
+
+		opts.on( '-d', '--database DBNAME',
+				"specify the database to connect to (default: \"#{options.database}\")" ) do |db|
+			options.database = db
+		end
+
+		opts.on( '-h', '--host HOSTNAME', 'database server host' ) do |host|
+			options.host = host
+		end
+
+		opts.on( '-p', '--port PORT', Integer,
+				"database server port (default: \"#{options.port}\")" ) do |port|
+			options.port = port
+		end
+
+		opts.on( '-n', '--schema SCHEMA', String,
+				"operate on the named schema only (default: none)" ) do |schema|
+			options.schema = schema
+		end
+
+		opts.on( '-T', '--tablespace SPACE', String,
+				"move old tables to this tablespace (default: \"#{options.tablespace}\")" ) do |tb|
+			options.tablespace = tb
+		end
+
+		opts.on( '-F', '--tableformat FORMAT', String,
+				"The naming format (strftime) for the inherited tables (default: none)" ) do |format|
+			options.format = format
+		end
+
+		opts.on( '-U', '--user NAME',
+				"database user name (default: \"#{options.user}\")" ) do |user|
+			options.user = user
+		end
+
+		opts.on( '-W', 'force password prompt' ) do |pw|
+			print 'Password: '
+			begin
+				system 'stty -echo'
+				options.pass = gets.chomp
+			ensure
+				system 'stty echo'
+				puts
+			end
+		end
+
+		opts.separator ''
+		opts.separator 'Other options:'
+
+		opts.on_tail( '--dry-run', "don't actually do anything" ) do
+			options.dryrun = true
+		end
+
+		opts.on_tail( '--help', 'show this help, then exit' ) do
+			$stderr.puts opts
+			exit
+		end
+
+		opts.on_tail( '--version', 'output version information, then exit' ) do
+			puts Stats::VERSION
+			exit
+		end
+	end
+
+	opts.parse!( args )
+	return options
+end
+
+
+if __FILE__ == $0
+	opts = parse_args( ARGV )
+	raise ArgumentError, "A naming format (-F) is required." unless opts.format
+
+	$stdout.sync = true
+	PGWarehouse.new( opts ).migrate
+end
+
+

data/spec/data/expected_trace.out

@@ -0,0 +1,26 @@
+To backend> Msg Q
+To backend> "SELECT 1 AS one"
+To backend> Msg complete, length 21
+From backend> T
+From backend (#4)> 28
+From backend (#2)> 1
+From backend> "one"
+From backend (#4)> 0
+From backend (#2)> 0
+From backend (#4)> 23
+From backend (#2)> 4
+From backend (#4)> -1
+From backend (#2)> 0
+From backend> D
+From backend (#4)> 11
+From backend (#2)> 1
+From backend (#4)> 1
+From backend (1)> 1
+From backend> C
+From backend (#4)> 11
+From backend> "SELECT"
+From backend> Z
+From backend (#4)> 5
+From backend> Z
+From backend (#4)> 5
+From backend> T

data/spec/lib/helpers.rb

@@ -0,0 +1,350 @@
+#!/usr/bin/env ruby
+
+require 'pathname'
+require 'rspec'
+require 'shellwords'
+require 'pg'
+
+TEST_DIRECTORY = Pathname.getwd + "tmp_test_specs"
+
+class Object
+  def jruby?
+    RUBY_PLATFORM =~ /java/
+  end
+end
+
+module PG::TestingHelpers
+
+
+	# Set some ANSI escape code constants (Shamelessly stolen from Perl's
+	# Term::ANSIColor by Russ Allbery <rra@stanford.edu> and Zenin <zenin@best.com>
+	ANSI_ATTRIBUTES = {
+		'clear'      => 0,
+		'reset'      => 0,
+		'bold'       => 1,
+		'dark'       => 2,
+		'underline'  => 4,
+		'underscore' => 4,
+		'blink'      => 5,
+		'reverse'    => 7,
+		'concealed'  => 8,
+
+		'black'      => 30,   'on_black'   => 40,
+		'red'        => 31,   'on_red'     => 41,
+		'green'      => 32,   'on_green'   => 42,
+		'yellow'     => 33,   'on_yellow'  => 43,
+		'blue'       => 34,   'on_blue'    => 44,
+		'magenta'    => 35,   'on_magenta' => 45,
+		'cyan'       => 36,   'on_cyan'    => 46,
+		'white'      => 37,   'on_white'   => 47
+	}
+
+
+	###############
+	module_function
+	###############
+
+	### Create a string that contains the ANSI codes specified and return it
+	def ansi_code( *attributes )
+		attributes.flatten!
+		attributes.collect! {|at| at.to_s }
+
+		return '' unless /(?:vt10[03]|xterm(?:-color)?|linux|screen)/i =~ ENV['TERM']
+		attributes = ANSI_ATTRIBUTES.values_at( *attributes ).compact.join(';')
+
+		# $stderr.puts "  attr is: %p" % [attributes]
+		if attributes.empty?
+			return ''
+		else
+			return "\e[%sm" % attributes
+		end
+	end
+
+
+	### Colorize the given +string+ with the specified +attributes+ and return it, handling
+	### line-endings, color reset, etc.
+	def colorize( *args )
+		string = ''
+
+		if block_given?
+			string = yield
+		else
+			string = args.shift
+		end
+
+		ending = string[/(\s)$/] || ''
+		string = string.rstrip
+
+		return ansi_code( args.flatten ) + string + ansi_code( 'reset' ) + ending
+	end
+
+
+	### Output a message with highlighting.
+	def message( *msg )
+		$stderr.puts( colorize(:bold) { msg.flatten.join(' ') } )
+	end
+
+
+	### Output a logging message if $VERBOSE is true
+	def trace( *msg )
+		return unless $VERBOSE
+		output = colorize( msg.flatten.join(' '), 'yellow' )
+		$stderr.puts( output )
+	end
+
+
+	### Return the specified args as a string, quoting any that have a space.
+	def quotelist( *args )
+		return args.flatten.collect {|part| part.to_s =~ /\s/ ? part.to_s.inspect : part.to_s }
+	end
+
+
+	### Run the specified command +cmd+ with system(), failing if the execution
+	### fails.
+	def run( *cmd )
+		cmd.flatten!
+
+		if cmd.length > 1
+			trace( quotelist(*cmd) )
+		else
+			trace( cmd )
+		end
+
+		system( *cmd )
+		raise "Command failed: [%s]" % [cmd.join(' ')] unless $?.success?
+	end
+
+
+	### Run the specified command +cmd+ after redirecting stdout and stderr to the specified
+	### +logpath+, failing if the execution fails.
+	def log_and_run( logpath, *cmd )
+		cmd.flatten!
+
+		if cmd.length > 1
+			trace( quotelist(*cmd) )
+		else
+			trace( cmd )
+		end
+
+		# Eliminate the noise of creating/tearing down the database by
+		# redirecting STDERR/STDOUT to a logfile if the Ruby interpreter
+		# supports fork()
+		logfh = File.open( logpath, File::WRONLY|File::CREAT|File::APPEND )
+
+		if jruby?
+      # FIXME: for some reason redirection in the system method don't
+      # work, so I ended up with this lengthy hack
+      logpath ||= "/dev/null"
+      f = File.open logpath, 'a'
+      old_stdout = $stdout
+      old_stderr = $stderr
+      $stdout = $stderr = f
+			system( *cmd )
+      $stdout = old_stdout
+      $stderr = old_stderr
+		else
+		  begin
+			  pid = fork
+		  rescue NotImplementedError
+			  logfh.close
+      end
+			if pid
+				logfh.close
+			else
+				$stdout.reopen( logfh )
+				$stderr.reopen( $stdout )
+				$stderr.puts( ">>> " + cmd.shelljoin )
+				exec( *cmd )
+				$stderr.puts "After the exec()?!??!"
+				exit!
+			end
+
+			Process.wait( pid )
+		end
+
+    unless $?.success?
+      system("cat #{logpath}")
+      raise "Command failed: [%s]" % [cmd.join(' ')]
+    end
+	end
+
+  def certs_directory
+    File.expand_path "#{__FILE__}/../../../certs"
+  end
+
+  alias :old_log_and_run :log_and_run
+	def log_and_run( logpath, *cmd )
+    return old_log_and_run logpath, cmd unless cmd.first =~ /initdb/
+
+    $stderr.puts "Setting up pg_hba.conf with some dummy usernames and passwords"
+    tmp = old_log_and_run logpath, cmd
+    File.open("#{@test_pgdata}/pg_hba.conf", "w") do |f|
+      f.puts "# TYPE    DATABASE        USER            ADDRESS         METHOD"
+      f.puts "  hostssl all             ssl             samehost        password"
+      f.puts "  host    all             ssl             samehost        reject"
+      f.puts "  host    all             password        samehost        password"
+      f.puts "  host    all             encrypt         samehost        md5"
+      f.puts "  host    all             all             samehost        trust"
+      f.puts "  local   all             all                             trust"
+    end
+    File.open("#{@test_pgdata}/postgresql.conf", "a") do |f|
+      f.puts "ssl = yes"
+    end
+    FileUtils.cp "#{certs_directory}/server.key", @test_pgdata
+    File.chmod 0400, "#{@test_pgdata}/server.key"
+    FileUtils.cp "#{certs_directory}/server.crt", @test_pgdata
+    tmp
+  end
+
+	### Check the current directory for directories that look like they're
+	### testing directories from previous tests, and tell any postgres instances
+	### running in them to shut down.
+	def stop_existing_postmasters
+		# tmp_test_0.22329534700318
+		pat = Pathname.getwd + 'tmp_test_*'
+		Pathname.glob( pat.to_s ).each do |testdir|
+			datadir = testdir + 'data'
+			pidfile = datadir + 'postmaster.pid'
+			if pidfile.exist? && pid = pidfile.read.chomp.to_i
+				$stderr.puts "pidfile (%p) exists: %d" % [ pidfile, pid ]
+				begin
+					Process.kill( 0, pid )
+				rescue Errno::ESRCH
+					$stderr.puts "No postmaster running for %s" % [ datadir ]
+					# Process isn't alive, so don't try to stop it
+				else
+					$stderr.puts "Stopping lingering database at PID %d" % [ pid ]
+					run pg_bin_path('pg_ctl'), '-D', datadir.to_s, '-m', 'fast', 'stop'
+				end
+			else
+				$stderr.puts "No pidfile (%p)" % [ pidfile ]
+			end
+		end
+	end
+
+
+  def pg_bin_path cmd
+    begin
+      bin_dir = `pg_config --bindir`.strip
+      "#{bin_dir}/#{cmd}"
+    rescue
+      cmd
+    end
+  end
+
+	### Set up a PostgreSQL database instance for testing.
+	def setup_testing_db( description )
+		require 'pg'
+		stop_existing_postmasters()
+
+		puts "Setting up test database for #{description}"
+		@test_pgdata = TEST_DIRECTORY + 'data'
+		@test_pgdata.mkpath
+
+		@port = 54321
+		ENV['PGPORT'] = @port.to_s
+		ENV['PGHOST'] = 'localhost'
+		@conninfo = "host=127.0.0.1 port=#{@port} dbname=test"
+
+		@logfile = TEST_DIRECTORY + 'setup.log'
+		trace "Command output logged to #{@logfile}"
+
+		begin
+			unless (@test_pgdata+"postgresql.conf").exist?
+				FileUtils.rm_rf( @test_pgdata, :verbose => $DEBUG )
+				$stderr.puts "Running initdb"
+				log_and_run @logfile, pg_bin_path('initdb'), '-E', 'UTF8', '--no-locale', '-D', @test_pgdata.to_s
+			end
+
+			$stderr.puts "Starting postgres"
+			log_and_run @logfile, pg_bin_path('pg_ctl'), 'start', '-l', @logfile.to_s, '-w', '-o', "-k #{TEST_DIRECTORY.to_s.dump}",
+				'-D', @test_pgdata.to_s
+			sleep 2
+
+			$stderr.puts "Creating the test DB"
+			log_and_run @logfile, pg_bin_path('psql'), '-e', '-c', 'DROP DATABASE IF EXISTS test', 'postgres'
+			log_and_run @logfile, pg_bin_path('createdb'), '-e', 'test'
+
+		rescue => err
+			$stderr.puts "%p during test setup: %s" % [ err.class, err.message ]
+			$stderr.puts "See #{@logfile} for details."
+			$stderr.puts *err.backtrace if $DEBUG
+			fail
+		end
+
+		conn = PG.connect( @conninfo )
+		conn.set_notice_processor do |message|
+			$stderr.puts( description + ':' + message ) if $DEBUG
+		end
+
+		return conn
+	end
+
+
+	def teardown_testing_db( conn )
+		puts "Tearing down test database"
+
+		if conn
+			check_for_lingering_connections( conn )
+			conn.finish
+		end
+
+		log_and_run @logfile, pg_bin_path('pg_ctl'), 'stop', '-m', 'fast', '-D', @test_pgdata.to_s
+	end
+
+
+	def check_for_lingering_connections( conn )
+		conn.exec( "SELECT * FROM pg_stat_activity" ) do |res|
+			conns = res.find_all {|row| row['pid'].to_i != conn.backend_pid }
+			unless conns.empty?
+				puts "Lingering connections remain:"
+				conns.each do |row|
+					puts "  [%d] {%s} %s -- %s" % row.values_at( 'pid', 'state', 'application_name', 'query' )
+				end
+			end
+		end
+	end
+
+	def connection_string_should_contain_application_name(conn_args, app_name)
+		conn_name = conn_args.match(/application_name='(.*)'/)[1]
+		conn_name.should include(app_name[0..10])
+		conn_name.should include(app_name[-10..-1])
+		conn_name.length.should <= 64
+	end
+
+	# Ensure the connection is in a clean execution status.
+	def verify_clean_exec_status
+		@conn.send_query( "VALUES (1)" )
+		@conn.get_last_result.values.should == [["1"]]
+	end
+end
+
+
+RSpec.configure do |config|
+	ruby_version_vec = RUBY_VERSION.split('.').map {|c| c.to_i }.pack( "N*" )
+
+	config.include( PG::TestingHelpers )
+	config.treat_symbols_as_metadata_keys_with_true_values = true
+
+	config.mock_with :rspec
+	config.filter_run_excluding :ruby_19 if ruby_version_vec <= [1,9,1].pack( "N*" )
+	if RUBY_PLATFORM =~ /mingw|mswin/
+		config.filter_run_excluding :unix
+	else
+		config.filter_run_excluding :windows
+	end
+	config.filter_run_excluding :socket_io unless
+		PG::Connection.instance_methods.map( &:to_sym ).include?( :socket_io )
+
+	config.filter_run_excluding :postgresql_90 unless
+		PG::Connection.instance_methods.map( &:to_sym ).include?( :escape_literal )
+
+	if !PG.respond_to?( :library_version )
+		config.filter_run_excluding( :postgresql_91, :postgresql_92, :postgresql_93 )
+	elsif PG.library_version < 90200
+		config.filter_run_excluding( :postgresql_92, :postgresql_93 )
+	elsif PG.library_version < 90300
+		config.filter_run_excluding( :postgresql_93 )
+	end
+end
+