25

Our system is running on Ubuntu, python 3.4, postgres 9.4.x and psycopg2.

We (will in the furture) split between dev, test and prod environments using schemas. I've create a convenience method for creating connections to our database. It uses json connection configuration files in order create the connection string. I want to configure the connection to use a particular schema for all following queries using the returned connection. I don't want my queries to have hardcoded schemas, because we should be able to easily switch between them depending on if we are in development, testing or production phase/environment.

Currently the convenience method looks like the following:

def connect(conn_config_file = 'Commons/config/conn_commons.json'):
    with open(conn_config_file) as config_file:    
        conn_config = json.load(config_file)

    conn = psycopg2.connect(
        "dbname='" + conn_config['dbname'] + "' " +
        "user='" + conn_config['user'] + "' " +
        "host='" + conn_config['host'] + "' " +
        "password='" + conn_config['password'] + "' " +
        "port=" + conn_config['port'] + " "
    )
    cur = conn.cursor()
    cur.execute("SET search_path TO " + conn_config['schema'])

    return conn

It works fine as long as you give it time to execute the set search_path query. Unfortunately, if I'm too fast with executing a following query a race condition happens where the search_path isn't set. I've tried to force the execution with doing a conn.commit() before the return conn, however, this resets the search_path to the default schema postgres so that it doesn't use, say, prod. Suggestions at the database or application layer is preferable, however, I know we probably could solve this at the OS level too, any suggestions in that direction are also welcomed.

An example json configuration file looks like the following:

{
    "dbname": "thedatabase",
    "user": "theuser",
    "host": "localhost",
    "password": "theusers_secret_password",
    "port": "6432",
    "schema": "prod"
}

Any suggestion is very appreciated.

André C. Andersen
  • 8,955
  • 3
  • 53
  • 79
  • Maybe wait a second before return? `import time; time.sleep(1)` – Tom-db Sep 28 '15 at 04:58
  • I can't believe I'm doing it but I'm going to go for your answer. If you'd like the credit you can post it as an answer and I'll accept it. – André C. Andersen Sep 30 '15 at 11:25
  • Do you mean, is it pretty inelegant? Yes, is only a workaround – Tom-db Sep 30 '15 at 11:58
  • 2
    Can you post some code that demonstrates the race condition? Given that you have a single connection and queries are executed sequentially, there shouldn't be any possibility for a race condition in the code you posted. If I copy/paste your code, hardcode the connection parameters for a local Postgres instance of mine, and then add a query `SHOW search_path`, the result of `SET` is visible. – Mark E. Haase Nov 30 '15 at 04:18
  • It's impossible to get a race condition in a single connection. As the name suggests, it needs at least two competitors. – klin Apr 19 '18 at 18:42
  • Ultimately, we ended up using different databases for different environments (e.g., prod vs test), and always explicitly give the schema in queries. – André C. Andersen Aug 19 '18 at 10:09

2 Answers2

34

I think a more elegant solution would be to set the search_path in options parameter of connect(), like so:

def connect(conn_config_file = 'Commons/config/conn_commons.json'):
    with open(conn_config_file) as config_file:    
        conn_config = json.load(config_file)

    schema = conn_config['schema']
    conn = psycopg2.connect(
        dbname=conn_config['dbname'],
        user=conn_config['user'],
        host=conn_config['host'],
        password=conn_config['password'],
        port=conn_config['port'],
        options=f'-c search_path={schema}',
    )
    return conn

Of course, you can use "options" as part of the connection string. But using keyword arguments prevents all the hassle with string concatenations.

I found this solution in this psycopg2 feature request. As for the "options" parameter itself, it's mentioned here.

butla
  • 695
  • 7
  • 15
  • 1
    This did not work for me on Ubuntu, but `f"--search_path={schema}"` did. – Nelewout Feb 09 '19 at 22:09
  • What error message did you get? I think psycopg delegates to Postgres' C library, which on ubuntu you get either by installing libpq-dev or postgresql (maybe with some suffix). But from what I see, libpq is also package in psycopg2's wheel. Can you try installing psycopg2 in a clean venv and say what you're getting? Is it the source distribution that you compile yourself, or is it the wheel? – butla Feb 10 '19 at 12:32
  • Minor addition (was useful for me): in case of `asyncpg` this should be passed through `server_settings={'search_path': schema)` to connection or pool creation function. – Anton Bryzgalov Nov 06 '21 at 21:03
15

I think a better idea is to have something like DatabaseCursor returning cursor you use to execute queries with "SET search_path..." instead of connection. Well I mean something like this:

class DatabaseCursor(object):

    def __init__(self, conn_config_file):
        with open(conn_config_file) as config_file:     
            self.conn_config = json.load(config_file)

    def __enter__(self):
        self.conn = psycopg2.connect(
            "dbname='" + self.conn_config['dbname'] + "' " + 
            "user='" + self.conn_config['user'] + "' " + 
            "host='" + self.conn_config['host'] + "' " + 
            "password='" + self.conn_config['password'] + "' " + 
            "port=" + self.conn_config['port'] + " " 
        )   
        self.cur = self.conn.cursor()
        self.cur.execute("SET search_path TO " + self.conn_config['schema'])

        return self.cur

    def __exit__(self, exc_type, exc_val, exc_tb):
        # some logic to commit/rollback
        self.conn.close()

and

with DatabaseCursor('Commons/config/conn_commons.json') as cur:
    cur.execute("...")
bersen
  • 370
  • 4
  • 8