Write a Postgres proxy. Day 1. Getting familiar with the API.
Tags: python, sql, postgres, redisql
Published: 2019-05-17 14:32:16

Write a Postgres proxy. Day 1.
==============================

[RediSQL, SQL steroids for Redis.](https://redisql.com) Is a very fast in-memory SQL engine. Its main features are:

1.  Speed, up to 130,000 insert per second
2.  Familiarity, it support standard SQL, no weird dialects
3.  Simplicity, it is very easy to operate and to use with binding for any language.

Code on github: [RedBeardLab/rediSQL](https://github.com/RedBeardLab/rediSQL)

* * *

 

In this series of post we are writing a postgres proxy that accept connections made using the postgres (PG) protocol and forward them to [RediSQL](https://redisql.com).

[Motivation and introduction of the project are here.](https://redbeardlab.com/2019/04/20/write-a-postgresql-proxy-motivation/)

My hope for this project is to distill the knowledge I am getting from this work and help other that are interested in exploring the PG protocol.

Intro
-----

This post is about the first day of this project so it is mostly introduction of the references used during this work and a little bit of code.

In this day we quickly reach the stage where we are able to receive a query from `psql` (the CLI tool for PG).

We will start the post showing the PG references that are most useful and the Python references of the `asyncio` module we used.

Then we will explore very quickly the few lines of code that I ended up writing.

The last section will explore the error I made during this day, trivial errors but that where a big time sink anyway.  

The references
--------------

Before to start this work I questioned if I should implement this proxy for Postgres or for MySQL.

To choose I explored at the documentation of both projects and both are quite good. However, the documentation for PG looked simpler to follow and more linear and I just decide to go for PG.

The main documentation for this project is the [Chapter 50 of the PG documentation.](https://www.postgresql.org/docs/9.5/protocol.html)

In particular the following sections are of extreme interest:

*   The [Message Flow](https://www.postgresql.org/docs/9.5/protocol-flow.html) section explains what is the flow of messages between PG and the client. It helps in understanding what message we should expect from the client and what message we are required to send as a server.
*   The [Message Data Types](https://www.postgresql.org/docs/9.5/protocol-message-types.html) section simply explains how to read and interpretate the section “Message Formats”.
*   The [Message Formats](https://www.postgresql.org/docs/9.5/protocol-message-formats.html) section goes into the details and enumerate the format of each kind of message. As an example we discover that usually each message start with a single letter that identify the type of message (like `R` is used for authentication related messages or that queries start with `Q`), then 4 bytes (an Int32) indicate the length of the whole message and finally the body of the message itself.

Python “ASYNCIO”
-------------------

While I would like to merge this project in the main [RediSQL](http://redisql.com) rust codebase, I am a strong believer that starting the project in Python is a good idea. I will gain the knowledge necessary to successfully re-write the software in Rust while having already faced most of the implementation difficulties in a language that allows very fast iteration. Moreover, it will be just impossible to merge the Python code base into the RediSQL rust codebase, so I will just be forced to re-write it.

While I am not looking for performance I still opted to work with `asyncio`, mostly because it was a long time I didn’t do any big work in Python and I wanted to get a pulse of the available tools. Moreover I hoped that it would be closer to what I would find in Rust with Tokio, but it seems to me that the two models are not very similar.

On the Python side I keep referring to the [Callback Based API for](https://docs.python.org/3.5/library/asyncio-protocol.html#transports-and-protocols-callback-based-api) `[asyncio](https://docs.python.org/3.5/library/asyncio-protocol.html#transports-and-protocols-callback-based-api)`.

The API is very simple, you simply sub-classed the `asyncio.Protocol` class and implemented three callbacks:

1.  `connection_made` for when a new connection is created to the server.
2.  `data_received` for when a new packed of data arrives to the server.
3.  `connection_lost` for when we loose connection with the client.

As you can imagine all the logic is in the `data_received`callback, and it will be more complex than a standard web-server. Indeed HTTP is a stateless protocol, everything is simpler if the protocol is stateless, each request does not depends on the previous one.

The PG protocol is stateful, it means that we need to store and use information from previous messages. As an example, a client, before to send its queries, needs to send an handshake and to authenticate. This means that our server will have at least two state, an “initial” state where each connection start and a “ready” state where a connection end ups only if it completed the handshake and authenticate.

The code
--------

Finally [here the code of this first day of code.](https://github.com/RedBeardLab/pg-redis-proxy/blob/fe3408b02f7f1cd12061cae75de3a309a5c7d935/main.py) The code is mostly boilerplate copied from the Python documentation but it is already enough to accept a connection from `psql` and receive the first query.

To test our progresses we started the Python server and, at the same time we execute `psql` giving as input [a file with few SQL statements](https://github.com/RedBeardLab/pg-redis-proxy/blob/70d1c2de917648900289a36e142720237aed29b2/goal.sql) to execute.

`psql -f goal.sql -h localhost -p 8888`

### The workflow of the day

Other than boilerplate code, the interesting part of the code are the definition of the magic number that identifies the messages:

    SSLRequestCode = b'\x04\xd2\x16\x2f' # == hex(80877103)
    StartupMessageCode = b'\x00\x03\x00\x00' # == hex(196608)
    
    NoSSL = b'\x4E' # == 'N'
    
    AuthenticationOk = b'\x52\x00\x00\x00\x08\x00\x00\x00\x00'
    AuthenticationCleartextPassword = b'\x52\x00\x00\x00\x08\x00\x00\x00\x03'
    
    ReadyForQuery = b'\x5A\x00\x00\x00\x05\x49' # == Z0005I , the last I stand for Idle 

And the logic to reply to the client:

        def _reply(self, data):
            if self.state == "initial" and data[4:8] == SSLRequestCode:
                self.transport.write(NoSSL)
            elif self.state == "initial" and data[4:8] == StartupMessageCode:
                # we don't require a password
                self.transport.write(AuthenticationOk)
                # good to go for the first query!
                self.transport.write(ReadyForQuery)
            return

Let’s explore how we get to this few lines of code.

My discovering process
----------------------

Scanning quickly the documentation it could seems like the first message to expect is the `StartupMessage`, however, the first message sent by `psql` is the `SSLRequest` message, and this took quite a while to figure out.

The `SSLRequest` message is recognized because it contains the magic number 80877103 which we encode in the python code as `b'\x04\xd2\x16\x2f'`

Since we don’t yet support SSL we simply respond to the `SSLRequest` with `N`(encoded as `b'\x4E'`) to let know to the client that we are not going to use SSL. At this point, the client, can either drop the connection or decide to accept a non-encrypted connection and send the `StartupMessage` in plain text.

Also for the `StartupMessage` there is a magic number (196608) which we encoded as `b'\x00\x03\x00\x00'`.

Along with the magic number, the `StartupMessage` contains information like the user who is starting the connection, what database the user is trying to connect and other information. At the moment we ignore all those information.

After the `StartupMessage` the server requires authentication, in our case we don’t care about authentication just yet and we just send the `AuthenticationOk` message.

The next step is a little tricky.

We just send a message to the client, the `AuthenticationOk` message and so I would expect the client to send the server something back.

Wrong!

Now, the server need to be proactive and tell the client that it is ok to start sending queries. We need to send two messages, one after the other to the client.

Indeed you can see in the code that we immediately send the `ReadyForQuery` message.

At this point our time is over for this day, however we can clearly see from the log that the next message received by the server is the first query of our file!

Success!

Errors made during this day
---------------------------

During this coding section I wasted a lot of time because I didn’t read the documentation with enough care.

Indeed I was expecting the `StartupMessage` as first message and not the `SSLRequest`. I spend a lot of time trying to fit the `StartupMessage` into the `SSLRequest`, maybe I was reading the message with the wrong endianess? Maybe there was “garbage” from the protocol layer?

Nah! I am just reading the wrong message.

Another time sink was me reading the wrong column in the ASCII table trying to use the decimal, instead of the hexadecimal, encoding. All the messages start with a letter, in our cases we needed the `N` for rejecting the SSL, and the `R` for the `AuthenticationOk` message and finally the `Z` for the `ReadyForQuery` message. As an example the `N` is 78 in decimal and 4E in hexadecimal. I was trying to encode `N` as `b'\x78'` instead of `b'\x4E'`.

Conclusion
----------

I hope your enjoyed the post.

I will keep publishing about this topic on this blog, so if you are interested feel free to [follow me on twitter](https://twitter.com/siscia_) or subscribe to the mail list just below.

All post of this serie: [Writing a Postgres proxy.](https://redbeardlab.com/category/pg-redis-proxy/)