CODEWITHSUNDEEP
pythonpostgresqlpsycopg2
dranobob
dranobob

Reputation: 836

How to create a postgres connection that uses context manager?

I was trying to create a wrapper class that used enter/exit to allow for a context manager to close a postgres db connection.

class ClosingConnection:
def __init__(self, schema_name: str) -> None:
    """
    :param schema_name: the db schema (i.e. the tenant)
    """
    super().__init__()
    self.schema_name = schema_name

def __enter__(self):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

def __exit__(self, exc_type, exc_val, exc_tb):
    self.conn.close()

Usage

def get_db(schema_name: str):
    return ClosingConnection(schema_name)

with get_db("test") as db:
    cursor = db.cursor()
    cursor.execute("some sql")
    rv = cursor.fetchall()
    cursor.close()

This works fine to retrieve data, but if I do an upsert, only the first upsert is performed. The remaining ones rv returns correctly, but the actual database is never updated.

In contrast if I get rid of the closing class, this works fine

Usage

def get_db(schema_name: str):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

conn = get_db("test") 
try:      
    cursor = db.cursor()
    cursor.execute("some sql")
    rv = cursor.fetchall()
    cursor.close()
finally:
    conn.close()

For curiosity here is the upsert command I'm using (I've also tried the older with upsert version)

INSERT INTO settings_account (setting_name, setting_value)
    VALUES (%(setting_name)s, %(setting_value)s)
    ON CONFLICT (setting_name)
    DO
        UPDATE SET setting_value=EXCLUDED.setting_value
    RETURNING *

Why is the wrapper class causing the db writes to fail? Is the idea of wrapper class flawed?

Upvotes: 2

Views: 1703

Answers (1)

dranobob
dranobob

Reputation: 836

Thanks the comment above, i was able to get this working with:

    class ClosingConnection:
def __init__(self, schema_name: str) -> None:
    """
    :param schema_name: the db schema (i.e. the tenant)
    """
    super().__init__()
    self.schema_name = schema_name

def __enter__(self):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

def __exit__(self, exc_type, exc_val, exc_tb):
    self.conn.commit()
    self.conn.close()

def get_db(schema_name: str):
    return ClosingConnection(schema_name)

with get_db("test") as db:
    with db.cursor() as cursor:
        cursor.execute("some sql")
        rv = cursor.fetchall()

Upvotes: 2

Related Questions