What is proper way to close connection in psyopcg2 with "with statement"?

Question

I want to konw, what is a proper way to closing connection with Postgres database using with statement and psyopcg2.

import pandas as pd
import psycopg2
def create_df_from_postgres(params: dict,
                                   columns: str,
                                   tablename: str,
                                   ) -> pd.DataFrame:

    with psycopg2.connect(**params) as conn:
        data_sql = pd.read_sql_query(
          "SELECT " + columns + ", SUM(total)"
          " AS total FROM " + str(tablename),
          con=conn
          )
    # i need to close conection here:
        # conn.close()

    # or here:
    conn.close()
    return data_sql

Is this a better way to handle connection ?

def get_ci_method_and_date(params: dict,
                           columns: str,
                           tablename: str,
                           ) -> pd.DataFrame:

    try:
        connection = psycopg2.connect(**params)
        data_sql = pd.read_sql_query('SELECT ' + columns +
                                     ' FROM ' + str(tablename),
                                     con=connection
                                     )
    finally:
        if(connection):
            connection.close()
    return data_sql

From official psycopg docs

Warning Unlike file objects or other resources, exiting the connection’s with block doesn’t close the connection, but only the transaction associated to it. If you want to make sure the connection is closed after a certain point, you should still use a try-catch block:

conn = psycopg2.connect(DSN)
try:
    # connection usage
finally:
    conn.close()

Antisepsis

Answer 1

Proper way to close a connection:

From official psycopg docs:

Warning Unlike file objects or other resources, exiting the connection’s with block doesn’t close the connection, but only the transaction associated to it. If you want to make sure the connection is closed after a certain point, you should still use a try-catch block:

conn = psycopg2.connect(DSN)
try:
    # connection usage
finally:
    conn.close()

Antisepsis

is it not advisable to put the "conn = psycopg2.connect(DSN)" inside the try statement incase there is an except with trying to connect? ie. wrong password or something. – Mercantilism

@Zaffer, I don't think it is necessary because the connection would not have been made in the first place. You could use another try/except block to catch exceptions related to the connection though – Othilie

Answer 2

I thought the Connection ContextManager closes the connection, but according to the docs, it does not:

Connections can be used as context managers. Note that a context wraps a transaction: if the context exits with success the transaction is committed, if it exits with an exception the transaction is rolled back. Note that the connection is not closed by the context and it can be used for several contexts.

Proposed usage is:

conn = psycopg2.connect(DSN)

with conn:
    with conn.cursor() as curs:
        curs.execute(SQL1)

with conn:
    with conn.cursor() as curs:
        curs.execute(SQL2)

# leaving contexts doesn't close the connection
conn.close()

source: https://www.psycopg.org/docs/connection.html

Lubra

I think that full answer for my question should be: The is no proper way to close the connection with "with statements" because the connection is not closed by the context. with conn: – Vange

You're right. I mean, you could build an own ContextManager to do this, but I think the main idea is that you don't close the connection straight away on purpose. Just wanted to document what I found and what the docs say because I came here through google and this was the missing info for me, so I wanted to document it. – Lubra

Answer 3

Depends on your code structure and logic, but you can also use:

@contextmanager
def _establish_connection():
    try:
        db_connection = psycopg2.connect(...)
        yield db_connection
    finally:
        # Extra safety check if the transaction was not rolled back by some reason
        if db_connection.status == psycopg2.extensions.STATUS_IN_TRANSACTION:
            db_connection.rollback()

        db_connection.close()

# After use your function like that

with _establish_connection() as conn:
    # Do your logic here
    return ...

Gastropod

Answer 4

The whole point of a with statement is that the resources are cleaned up automatically when it exits. So there is no need to call conn.close() explicitly at all.

Ox

this is the answer, if the conn supports an __exit__ call to invoke close() then it will take care of it for you. #1984825 – Vyse

So assuming that conn doesn't supoort __exit__ , my with statement is out of sense? – Vange

I don't understand your question. It does support it. – Ox

I try to use with statements as much as I can. But I do have some unexpected EOF on client connection with an open transaction warning from PSQL with this way of doing. Is it possible that this happens, or should I go see elsewhere because it's impossible? – Extravascular

@DanielRoseman, Psycopg2 Docs say that "the connection is not closed by the context" Docs - Connection Class – Othilie

I thought this answer is right, but according to the docs, it's not. I've created an answer to lay out why. – Lubra