Merge pull request #559 from larskanis/reset-with-dns

Renew address resolution in conn.reset

ext/pg_connection.c

@@ -563,6 +563,27 @@ pgconn_sync_reset( VALUE self )
 	return self;
 }
 
+static VALUE
+pgconn_reset_start2( VALUE self, VALUE conninfo )
+{
+	t_pg_connection *this = pg_get_connection( self );
+
+	/* Close old connection */
+	pgconn_close_socket_io( self );
+	PQfinish( this->pgconn );
+
+	/* Start new connection */
+	this->pgconn = gvl_PQconnectStart( StringValueCStr(conninfo) );
+
+	if( this->pgconn == NULL )
+		rb_raise(rb_ePGerror, "PQconnectStart() unable to allocate PGconn structure");
+
+	if ( PQstatus(this->pgconn) == CONNECTION_BAD )
+		pg_raise_conn_error( rb_eConnectionBad, self, "%s", PQerrorMessage(this->pgconn));
+
+	return Qnil;
+}
+
 /*
  * call-seq:
  *    conn.reset_start() -> nil
@@ -4468,6 +4489,7 @@ init_pg_connection(void)
 	rb_define_method(rb_cPGconn, "finished?", pgconn_finished_p, 0);
 	rb_define_method(rb_cPGconn, "sync_reset", pgconn_sync_reset, 0);
 	rb_define_method(rb_cPGconn, "reset_start", pgconn_reset_start, 0);
+	rb_define_private_method(rb_cPGconn, "reset_start2", pgconn_reset_start2, 1);
 	rb_define_method(rb_cPGconn, "reset_poll", pgconn_reset_poll, 0);
 	rb_define_alias(rb_cPGconn, "close", "finish");
 

lib/pg/connection.rb

@@ -565,7 +565,12 @@ def encrypt_password( password, username, algorithm=nil )
 	# Resets the backend connection. This method closes the
 	# backend connection and tries to re-connect.
 	def reset
-		reset_start
+		iopts = conninfo_hash.compact
+		if iopts[:host] && !iopts[:host].empty? && PG.library_version >= 100000
+			iopts = self.class.send(:resolve_hosts, iopts)
+		end
+		conninfo = self.class.parse_connect_args( iopts );
+		reset_start2(conninfo)
 		async_connect_or_reset(:reset_poll)
 		self
 	end
@@ -773,6 +778,40 @@ def new(*args)
 		alias setdb new
 		alias setdblogin new
 
+		# Resolve DNS in Ruby to avoid blocking state while connecting.
+		# Multiple comma-separated values are generated, if the hostname resolves to both IPv4 and IPv6 addresses.
+		# This requires PostgreSQL-10+, so no DNS resolving is done on earlier versions.
+		private def resolve_hosts(iopts)
+			ihosts = iopts[:host].split(",", -1)
+			iports = iopts[:port].split(",", -1)
+			iports = [nil] if iports.size == 0
+			iports = iports * ihosts.size if iports.size == 1
+			raise PG::ConnectionBad, "could not match #{iports.size} port numbers to #{ihosts.size} hosts" if iports.size != ihosts.size
+
+			dests = ihosts.each_with_index.flat_map do |mhost, idx|
+				unless host_is_named_pipe?(mhost)
+					if Fiber.respond_to?(:scheduler) &&
+								Fiber.scheduler &&
+								RUBY_VERSION < '3.1.'
+
+						# Use a second thread to avoid blocking of the scheduler.
+						# `TCPSocket.gethostbyname` isn't fiber aware before ruby-3.1.
+						hostaddrs = Thread.new{ Addrinfo.getaddrinfo(mhost, nil, nil, :STREAM).map(&:ip_address) rescue [''] }.value
+					else
+						hostaddrs = Addrinfo.getaddrinfo(mhost, nil, nil, :STREAM).map(&:ip_address) rescue ['']
+					end
+				else
+					# No hostname to resolve (UnixSocket)
+					hostaddrs = [nil]
+				end
+				hostaddrs.map { |hostaddr| [hostaddr, mhost, iports[idx]] }
+			end
+			iopts.merge(
+				hostaddr: dests.map{|d| d[0] }.join(","),
+				host: dests.map{|d| d[1] }.join(","),
+				port: dests.map{|d| d[2] }.join(","))
+		end
+
 		private def connect_to_hosts(*args)
 			option_string = parse_connect_args(*args)
 			iopts = PG::Connection.conninfo_parse(option_string).each_with_object({}){|h, o| o[h[:keyword].to_sym] = h[:val] if h[:val] }
@@ -782,37 +821,7 @@ def new(*args)
 				# hostaddr is provided -> no need to resolve hostnames
 
 			elsif iopts[:host] && !iopts[:host].empty? && PG.library_version >= 100000
-				# Resolve DNS in Ruby to avoid blocking state while connecting.
-				# Multiple comma-separated values are generated, if the hostname resolves to both IPv4 and IPv6 addresses.
-				# This requires PostgreSQL-10+, so no DNS resolving is done on earlier versions.
-				ihosts = iopts[:host].split(",", -1)
-				iports = iopts[:port].split(",", -1)
-				iports = [nil] if iports.size == 0
-				iports = iports * ihosts.size if iports.size == 1
-				raise PG::ConnectionBad, "could not match #{iports.size} port numbers to #{ihosts.size} hosts" if iports.size != ihosts.size
-
-				dests = ihosts.each_with_index.flat_map do |mhost, idx|
-					unless host_is_named_pipe?(mhost)
-						if Fiber.respond_to?(:scheduler) &&
-									Fiber.scheduler &&
-									RUBY_VERSION < '3.1.'
-
-							# Use a second thread to avoid blocking of the scheduler.
-							# `TCPSocket.gethostbyname` isn't fiber aware before ruby-3.1.
-							hostaddrs = Thread.new{ Addrinfo.getaddrinfo(mhost, nil, nil, :STREAM).map(&:ip_address) rescue [''] }.value
-						else
-							hostaddrs = Addrinfo.getaddrinfo(mhost, nil, nil, :STREAM).map(&:ip_address) rescue ['']
-						end
-					else
-						# No hostname to resolve (UnixSocket)
-						hostaddrs = [nil]
-					end
-					hostaddrs.map { |hostaddr| [hostaddr, mhost, iports[idx]] }
-				end
-				iopts.merge!(
-					hostaddr: dests.map{|d| d[0] }.join(","),
-					host: dests.map{|d| d[1] }.join(","),
-					port: dests.map{|d| d[2] }.join(","))
+				iopts = resolve_hosts(iopts)
 			else
 				# No host given
 			end

spec/helpers.rb

@@ -675,6 +675,11 @@ def with_env_vars(**kwargs)
 			ENV.update(old_values)
 		end
 	end
+
+	# Append or change 'rubypg_test' host entry in /etc/hosts to a given IP address
+	def set_etc_hosts(hostaddr)
+		system "sudo --non-interactive sed -i '/.* rubypg_test$/{h;s/.*/#{hostaddr} rubypg_test/};${x;/^$/{s//#{hostaddr} rubypg_test/;H};x}' /etc/hosts" or skip("unable to change /etc/hosts file")
+	end
 end
 
 

spec/pg/connection_spec.rb

@@ -1690,6 +1690,23 @@
 		conn.close
 	end
 
+	it "refreshs DNS address while conn.reset", :without_transaction, :ipv6 do
+		set_etc_hosts "::1"
+		conn = described_class.connect( "postgres://rubypg_test/test" )
+		conn.exec("select 1")
+
+		set_etc_hosts "127.0.0.1"
+		conn.reset
+		conn.exec("select 1")
+
+		set_etc_hosts "::2"
+		expect do
+			conn.reset
+			conn.exec("select 1")
+		end.to raise_error(PG::Error)
+	end
+
+
 	it "closes the IO fetched from #socket_io when the connection is closed", :without_transaction do
 		conn = PG.connect( @conninfo )
 		io = conn.socket_io