Checking Postgres availability with pure Ruby

A quick introduction to Postgres wire protocol

Recently I had to write a simple script, which single job was to wait until Postgres becomes available. This is really important when you need to run queries on your database as part of some automatic workflow, but you don’t have any guarantees that database service will be ready in time. My concrete use cases were:

Running tests on CI server using Docker Compose. I needed to wait for Postgres before starting test script. Otherwise I could see random test failures due to database connection errors

Automatically running database seeds on Heroku after creating new apps. I needed to wait until Postgres addon was fully provisioned and attached to the app

The obvious way to implement such script is to either use psql or Postgres library for your language of choice. But because I played a little bit with Postgres protocol in the past, I’ve decided that it will be good idea to write it in pure Ruby and avoid installing psql inside Docker container.

This article describes small portion of Postgres protocol and how to write Ruby program that communicates with Postgres using this protocol.

TL;DR If you are interested just in the script, you can jump straight to the repository. If you want to learn about the protocol or some Ruby tricks, read on.

Protocol basics

Postgres protocol is quite extensive with over 50 message types and various modes. But for the purpose of this article, you need to know only few of them. Take a look at diagram below that describes message flow when establishing new connection to the server:

Technically the format is binary, but most of the data is sent as ASCII. And almost all message types have the same format, comprised of three different parts:

Character Tag (1 byte) - identifies message type

Length (4 bytes) - length of the message including length field and excluding character tag

Payload - the actual body of the message, depending on message type

Startup message is the exception here. It does not contain the character tag:

Name / Value pairs - parameters send to the server, most important are user and database, where user is required parameter

Zero byte - required at the end of startup message

Going back to the diagram, here is the flow again, explained step by step:

Clients sends Startup Message with required parameters

If authentication is required, server responds with Authentication Request message. This message varies based on server configuration and selected authentication method. If authentication is not required, server immediately responds with Authentication OK message and the flow continues from step 5.

Let’s walk through it step by step, starting with build_startup_message method.

It’s main job is to construct array of values and translate it into byte stream that can be send to the Postgres server. The easiest way to do this in Ruby is to use Array#pack method. It takes an array and a format string as parameters and returns an output string containing all values encoded according to the specified format string. This method takes care of correct padding, big / little endian encoding, terminating strings with zero bytes etc. You can read more about it here and here

The method starts with some variable definitions:

message = [0, 196608]
message_size = 4 + 4
pack_string = "L>L>"

The first two parts of startup message are it’s length and protocol version. At this point we don’t know yet what message size will be, so we need to put 0 as the first element. Protocol version is 3.0. Two most significant bytes represent major version, and two least significant bytes represent 0. This gives value 196608.

message_size stores current message size. We initialize it with value of 8, because message size and protocol version are already added to the message body. Writing 4 + 4 instead of just 8 better indicates that this is 4 bytes for message length and another 4 bytes for protocol version.

pack_string stores current format string that will be passed to Array#pack method. Final format string will depend on passed arguments to the build_startup_message method, so we should store it in a variable. And it will be easy to extend this method with support for other parameters.

We iterate over the params array and add each value to the message body. In the same time we need to increase the message size and update format string. Notice that message size is increased by one extra byte. This is for zero byte that terminates each value.

message << 0
message_size += 1
pack_string << "C"

We need to add zero byte at the end of the message and of course update size and format string.

message[0] = message_size

At this point the whole message body is ready, so we can set the correct message length inside the message body.

message.pack(pack_string)

And finally we encode message body as a string. For example, if user is postgres and database name is my_database, the final message will look like this:

This is main event loop. At the beginning we read 1-byte character tag from the server, followed by message length and actual message payload. Depending on received character tag (message type) we can executed appropriate code.

case char_tag
when "E"
# Received **ErrorResponse**
break
# hidden for brevity
end

This matches error response code. In such case we break the loop and the method will return false.

This piece matches various authentiation requests. First 4-bytes of payload specify what type of request it is. Value 0 means that authentication was successful and we don’t have to do anything more.

Value 3 means that cleartext password is required, so we are sending Password Message. This packet is encoded in the same way as startup message, using Array.pack method. But the format is simpler. First element is the character tag 112 (lowercase p). Then message length and password itself. Message length is always password length + 5 bytes. 4 bytes for message length field and 1 byte for zero byte at the end of the message.

Value 5 means that MD5-encoded password is required. In this case server also sends salt value in the message payload. We are using user name, password and received salt to calculate the correct password challange and then we are sending it to the server. The message format is the same as with cleartext password. The only difference is how the password is encoded.

Conclusion

I hope you liked the article and learnt something new. I realize that knowing Postgres protocol is not a skill that you will use everyday, but it has some practical usages. In the future I plan to write simple Postgres proxy that will be able to simulate various database problems in tests such as timeouts, terminating connections, cancelled requests etc. Stay tuned!