cluster-aware psycopg2 🚀
Tags: postgres, yugabytedb, python
Published: 2022-04-21 10:14:46

With a YugabyteDB database, you can connect your PostgreSQL application to any node in the cluster. Or you can choose a subset of nodes that are closer, limited to your availability zone, or region. The PostgreSQL drivers are not clustered-aware because they were built for a monolithic database where only the primary node can serve reads and writes.

You can manually balance your nodes, but, with a distributed SQL database like YugabyteDB, nodes can be added or removed transparently, to scale with the load or be resilient to failures. You can use a load balancer, like in a Kubernetes cluster. But to make it easy without a specific configuration, the YugabyteDB Smart Drivers detect the available nodes and balance the connections among them.

The Smart Driver for Java exists for some time. Here is the Python one, a fork from `psycopg2`.

## psycopg2

I'm using a virtual environment for this test, installing `pandas`for easy query and display, `sqlalchemy` to get a connection pool, and `psycopg2` to connect to YugabyteDB through the PostgreSQL protocol:
```bash
[ Franck ~]$ python -m venv  demo
[ Franck ~]$ .  demo/bin/activate
( demo) [ Franck ~]$
( demo) [ Franck ~]$ pip install --upgrade pip
...
Successfully installed pip-21.3.1

( demo) [ Franck ~]$ python -m pip install \
 sqlalchemy pandas psycopg2
...
( demo) [ Franck ~]$
```

There, here is my simple program to run with `python3`:

```python
import sqlalchemy
import threading
import pandas

pandas.set_option('display.max_colwidth', None)

# Connection pool to YugabyteDB database
yb=sqlalchemy.create_engine('\
postgresql+psycopg2://yugabyte:yugabyte@\
yb1.pachot.net:5433/yugabyte'\
,pool_size=1,max_overflow=5)

# function to grab a connection and show where I'm connected, 5 times
def mythread():
 for i in range(5):
  print(pandas.read_sql_query(f"""
   select pg_sleep(5), format( 'Thread %s connected to: %s'
    ,'{threading.current_thread().name}'
    , replace(current_setting('listen_addresses'),'0.0.0.0',host(inet_server_addr())::text)
    )""" , yb.connect()).to_string(index=False,header=False))

# start 9 threads
mythreads=[]
for i in range(9):
 t=threading.Thread(target=mythread)
 mythreads.append(t)
 t.start()

# wait and exit
for i in mythreads:
 t.join()

# end
```

This connects to my lab that is opened on the public internet (YugabyteDB is installed in multiple regions with VMs in the Oracle Cloud free tier). You can try this connection string and see the same if I'm not breaking my lab at that time. I provided the endpoint `yb1.pachot.net:5433` which resolves to 129.159.250.98 in the public internet, and bound to the 10.0.0.231 private IP in my VPC subnet. The program connects with multiple threads and displays the private IP with `inet_server_addr()`:

```
   Thread Thread-5 connected to: 10.0.0.231
   Thread Thread-1 connected to: 10.0.0.231
   Thread Thread-4 connected to: 10.0.0.231
   Thread Thread-2 connected to: 10.0.0.231
   Thread Thread-6 connected to: 10.0.0.231
   Thread Thread-3 connected to: 10.0.0.231
   Thread Thread-7 connected to: 10.0.0.231
   Thread Thread-8 connected to: 10.0.0.231
   Thread Thread-9 connected to: 10.0.0.231
   Thread Thread-1 connected to: 10.0.0.231
   Thread Thread-5 connected to: 10.0.0.231
   Thread Thread-4 connected to: 10.0.0.231
```

All threads are connected to the same endpoint. This is the default behavior of the `psycopg2` driver. It does the same as connecting directly to one node and using this connection:

```bash
[ Franck ~]$ psql -h yb1.pachot.net -p 5433 -c "create temporary table demo (me text); copy demo from program 'wget -qO- ifconfig.me' with (rows_per_transaction 0); select me, inet_server_addr() from demo;"

       me       | inet_server_addr
----------------+------------------
 129.159.250.98 | 10.0.0.231
```

## psycopg2-yugabytedb

I'm now installing the YugabyteDB Smart Driver, which overrides the `psycopg2` installation with this cluster-aware fork:
```bash
( demo) [ Franck ~]$ python -m pip install psycopg2-yugabytedb
Collecting psycopg2-yugabytedb
...
Successfully installed psycopg2-yugabytedb-2.9.3.post0
```

## load_balance=true

I'll run the same program with only one change, adding `?load_balance=true` to the connection URL

```python
yb=sqlalchemy.create_engine('\
postgresql+psycopg2://yugabyte:yugabyte@\
yb1.pachot.net:5433/yugabyte?load_balance=true'\
,pool_size=1,max_overflow=5)
```

This enables the additional cluster-aware logic. There is the output:

```
Couldn't connect to  129.151.253.58  adding to failed list
For cleanup purposes

   Thread Thread-15 connected to: 10.0.0.14
   Thread Thread-11 connected to: 10.0.0.200
   Thread Thread-12 connected to: 10.0.0.231
   Thread Thread-10 connected to: 10.0.2.100
   Thread Thread-13 connected to: 10.0.2.100
   Thread Thread-16 connected to: 10.0.0.14
   Thread Thread-17 connected to: 10.0.0.200
   Thread Thread-18 connected to: 10.0.0.231
   Thread Thread-15 connected to: 10.0.2.100
   Thread Thread-14 connected to: 10.0.2.100
   Thread Thread-11 connected to: 10.0.2.100
   Thread Thread-12 connected to: 10.0.0.14
   Thread Thread-10 connected to: 10.0.0.200
   Thread Thread-13 connected to: 10.0.0.231
   Thread Thread-16 connected to: 10.0.2.100
   Thread Thread-17 connected to: 10.0.2.100
   Thread Thread-18 connected to: 10.0.2.100
```
I can see many connections to different hosts. And a message that `129.151.253.58` is unavailable. The driver now tries all known nodes, and detects which ones are available (I have blocked the `129.151.253.58`, which is `yb7.pachot.net` and its host name `593ba5730dae`, but not opening the port).

Here is how my cluster looks like (and you see host names here):

![cluster](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9s6ujku8xqontptnk1k0.png)

Note that the private IP is the only easy one to display from a connection, with `inet_server_addr` but some of my nodes are on docker and show the same private IP. There other node in same region (Marseille) is `yb5.pachot.net`, `144.24.199.202` also showing `10.0.2.100` as private IP.

## yb_servers()

The smart driver is able to discover the other nodes because this is exposed by YugabyteDB with `yb_servers()`:

```
[ Franck ~]$ psql -h yb1.pachot.net -p 5433 -c "select * from yb_servers()"
      host      | port | num_connections | node_type | cloud  | region  |             zone              |    public_ip
----------------+------+-----------------+-----------+--------+---------+-------------------------------+-----------------
 65ab17172451   | 5433 |               0 | primary   | Oracle | Europe  | FR-Provence-Alpes-Cote-d'Azur | 144.24.199.202
 593ba5730dae   | 5433 |               0 | primary   | Oracle | Europe  | FR-Provence-Alpes-Cote-d'Azur | 129.151.253.58
 yb3.pachot.net | 5433 |               0 | primary   | Oracle | Europe  | CH-Zurich                     | 140.238.171.183
 yb2.pachot.net | 5433 |               0 | primary   | Oracle | Europe  | DE-Hesse                      | 150.230.147.84
 0424d4f96e7c   | 5433 |               0 | primary   | Oracle | Europe  | GB-England                    | 140.238.64.71
 774769fdf9c5   | 5433 |               0 | primary   | Oracle | America | US-Virginia                   | 132.145.164.186
 yb1.pachot.net | 5433 |               0 | primary   | Oracle | Europe  | DE-Hesse                      | 129.159.250.98
(7 rows)
```
 
This view shows the `host` name (with resolves to the private IP) and the `public ip`. When you connect, the host used in the connection string is resolved and compared to `public_ip` to detect that we connect from the public network. Or compared to the resolved `host` to detect that we connect in the private subnet. Then, the list of other addresses in the same (private or public) is read and the driver will load-balance new connections to them.

## topology_keys=cloud-region.zone,cloud-region.zone

In the above `yb_servers()` you can see that each host has a `cloud`, `region` and `zone` information. Those are defined when starting the `yb-tserver` with the placement flags.

This can be automatically gathered from the cloud provider metadata and I could have used this to set the `oracle.com` cloud provider here, its Region and Availability Domain:
```bash
placement_cloud: $(curl -s http://169.254.169.254/opc/v1/instance/regionInfo/realmDomainComponent)
placement_region: $(curl -s http://169.254.169.254/opc/v1/instance/region)
placement_zone: $(curl -s http://169.254.169.254/opc/v1/instance/ociAdName)
```

 In my lab, I've set them from `ipinfo.io` to the geographical location of the IP address:
```bash
--placement_cloud  '$(curl -s ipinfo.io | jq -r .org | cut -f2 -d " "    | iconv -f utf-8 -t ascii//TRANSLIT | tr " ,." -  )' \
--placement_region '$(curl -s ipinfo.io | jq -r .timezone | cut -d/ -f1  | iconv -f utf-8 -t ascii//TRANSLIT | tr " ,." 0  )' \
--placement_zone   '$(curl -s ipinfo.io | jq -r .country+"\" \""+.region | iconv -f utf-8 -t ascii//TRANSLIT | tr " ,." -  )' \
--use_private_ip=zone 
```
This has set the following topology:
![cloud.zone.region](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/l68ee9oy3svulj8dluc9.png)
 

This `cloud`, `region` and `zone` information can be used to connect to a subset of the cluster. Typically, you run your application servers in multiple Availability Zones, and it makes sense that each one connects to the same zone. Because, in case of AZ failure, both go down at the same time and the others are not impacted. In order to do that, we can add a `topology_keys`:

```python
yb=sqlalchemy.create_engine('\
postgresql+psycopg2://yugabyte:yugabyte@\
yb1.pachot.net:5433/yugabyte?load_balance=true&topology_keys=Oracle.Europe.FR-Provence-Alpes-Cote-d\'Azur'\
,pool_size=1,max_overflow=5)
```
Here is the output:
```
( demo) [ Franck ~]$ python /home/opc/demo/tmp/test-yb-smart-driver.py
Couldnt connect to  129.151.253.58  adding to failed list
For cleanup purposes
   Thread Thread-2 connected to: 10.0.0.231
   Thread Thread-3 connected to: 10.0.0.231
   Thread Thread-6 connected to: 10.0.2.100
   Thread Thread-5 connected to: 10.0.2.100
   Thread Thread-7 connected to: 10.0.2.100
   Thread Thread-2 connected to: 10.0.0.231
   Thread Thread-8 connected to: 10.0.0.231
   Thread Thread-9 connected to: 10.0.2.100
   Thread Thread-3 connected to: 10.0.2.100
   Thread Thread-1 connected to: 10.0.0.231
```

`10.0.2.100` is one of my host in Marseille. The other is the one with public IP `129.151.253.58` for which the ports are not open. The initial endpoint, even when not in the topology zone defined, is still used in the pool. This is different from the JDBC Smart Driver, but remember that they are in beta version. I'll update this post after checking if it is expected.

The `topology_keys` can be a comma-separated list (so now you understand my `tr " ,."` to avoid commas, dots and spaces in the placement names), like this:

```python
yb=sqlalchemy.create_engine('\
postgresql+psycopg2://yugabyte:yugabyte@\
yb1.pachot.net:5433/yugabyte\
?load_balance=true\
&topology_keys=Oracle.Europe.CH-Zurich,Oracle.Europe.DE-Hesse,Oracle.Europe.FR-Provence-Alpes-Cote-d\'Azur'\
,pool_size=1,max_overflow=5)
```

I've run my program with this connection string, though `sort | uniq -c`:
```
Couldnt connect to  129.151.253.58  adding to failed list
For cleanup purposes
     10    Thread Thread-1 connected to: 10.0.0.231
     10    Thread Thread-2 connected to: 10.0.0.200
     10    Thread Thread-3 connected to: 10.0.0.200
     10    Thread Thread-4 connected to: 10.0.0.14
     10    Thread Thread-5 connected to: 10.0.0.14
     10    Thread Thread-6 connected to: 10.0.2.100
```
This shows that my connections were distributed to multiple hosts.

The project is on https://github.com/yugabyte/psycopg2
Your feedback is welcome, here, twitter, https://twitter.com/FranckPachot/status/1517085003901030400?s=20&t=dWHVsJ1scfdH0GDRDJvxLw or the Yugabyte community channels https://www.yugabyte.com/community/