DEV Community

Franck Pachot
Franck Pachot

Posted on

cluster-aware psycopg2 🚀

#postgres #yugabytedb #python

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 pandasfor easy query and display, sqlalchemy to get a connection pool, and psycopg2 to connect to YugabyteDB through the PostgreSQL protocol:

[ 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 ~]$
Enter fullscreen mode Exit fullscreen mode

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

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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:

[ 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
Enter fullscreen mode Exit fullscreen mode

psycopg2-yugabytedb

I'm now installing the YugabyteDB Smart Driver, which overrides the psycopg2 installation with this cluster-aware fork:

( demo) [ Franck ~]$ python -m pip install psycopg2-yugabytedb
Collecting psycopg2-yugabytedb
...
Successfully installed psycopg2-yugabytedb-2.9.3.post0
Enter fullscreen mode Exit fullscreen mode

load_balance=true

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

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

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
Enter fullscreen mode Exit fullscreen mode

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

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)
Enter fullscreen mode Exit fullscreen mode

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:

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)
Enter fullscreen mode Exit fullscreen mode

In my lab, I've set them from ipinfo.io to the geographical location of the IP address:

--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
Enter fullscreen mode Exit fullscreen mode

This has set the following topology:
cloud.zone.region

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:

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)
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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:

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)
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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/

Top comments (0)

Read next

pauq profile image

Scraping New Telegram Channels

pauq -

ortizdavid profile image

Como Tratar Diferentes Tipos de Logs em Go?

Ortiz de Arcanjo Antonio David -

jmegnidro profile image

Distance de Levenshtein : Le Guide Ultime pour Mesurer la Similarité Textuelle

Dominique Megnidro -

smuniraj profile image

Convert an Excel dataset into a SQL insert statement

Saravanan Muniraj -