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()

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()

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

Answer 3

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

@contextmanager
def _establish_connection():
    db_connection = psycopg2.connect(...)

    try:
        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 ...

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.