In this blog post, we will discuss how to set session timeouts in Python using Psycopg2. Psycopg2 is a PostgreSQL adapter for Python that allows easy integration between Python and PostgreSQL databases. By setting session timeouts, you can control the maximum duration of a session before it is automatically terminated.

Setting up Psycopg2

Before we start, make sure you have Psycopg2 installed. You can install it using pip:

pip install psycopg2

Setting Session Timeouts

To set session timeouts in Psycopg2, you need to modify the connect method that establishes a connection with the PostgreSQL database. The connect method accepts various parameters, including options that allow you to customize the session settings.

Here’s an example of how to set the session timeout to 10 minutes using the options parameter:

import psycopg2
from psycopg2 import extensions

# Set session timeout in seconds
timeout = 600

# Create the connection object
conn = psycopg2.connect(
    dbname="your_db", 
    user="your_user", 
    password="your_password", 
    host="your_host", 
    port="your_port",
    options=f'-c statement_timeout={timeout * 1000}'
)

# ... Rest of the code ...

In the code above, we use the options parameter to set the statement_timeout parameter, which determines the session timeout in milliseconds. We multiply the timeout value by 1000 to convert it to milliseconds.

Handling Session Timeouts

When a session timeout occurs, an exception psycopg2.errors.AdminShutdown is raised. You can catch and handle this exception to perform any necessary cleanup or take appropriate actions.

try:
    # Perform database operations
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM your_table")
    rows = cursor.fetchall()
    cursor.close()

except psycopg2.errors.AdminShutdown:
    # Handle session timeout
    print("Session timeout occurred")
    # Perform necessary cleanup or actions

finally:
    # Close the connection
    conn.close()

In the code snippet above, we wrap the database operations within a try-except block to catch the psycopg2.errors.AdminShutdown exception. In the except block, you can handle the session timeout by performing necessary cleanup or other actions.

Conclusion

By setting session timeouts in Psycopg2, you can control the duration of a session and ensure that long-running queries or idle connections are automatically terminated. This helps in managing database resources efficiently and improving overall system performance.