Learn to connect a Python application to a PostgreSQL database using psycopg2, covering connection setup, SQL execution, error handling, and retry mechanisms.
In this lesson, you’ll learn how to connect a Python application to a PostgreSQL database using the psycopg2 library. We will cover setting up a connection, executing SQL commands, and handling connection errors gracefully. Additionally, you’ll discover how to implement a retry mechanism if the initial connection fails.The following Python code demonstrates how to import the library, establish a connection, and perform basic SQL operations:
Copy
Ask AI
import psycopg2# Establish a connection to the PostgreSQL databaseconn = psycopg2.connect("dbname=test user=postgres")cur = conn.cursor()# Create a table, insert data, and fetch itcur.execute("CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);")cur.execute("INSERT INTO test (num, data) VALUES (%s, %s)", (100, "abc'def"))cur.execute("SELECT * FROM test;")print(cur.fetchall())# Commit the transaction and close connectionsconn.commit()cur.close()conn.close()
The workflow illustrated above includes:
Importing the psycopg2 library.
Establishing a connection using specific connection parameters.
Creating a cursor for executing SQL commands.
Executing SQL commands such as creating a table, inserting data, and selecting data.
Committing the transaction and closing the connection.
When developing an API with a framework like FastAPI, managing database connections effectively becomes essential. The snippet below illustrates how to integrate psycopg2 with FastAPI while leveraging Pydantic models for data validation:
Copy
Ask AI
from typing import Optionalfrom fastapi import FastAPI, Response, status, HTTPException, Bodyfrom pydantic import BaseModelfrom random import randrangeimport psycopg2app = FastAPI()class Post(BaseModel): title: str content: str published: bool = True# Dummy data for postsmy_posts = [ {"title": "title of post 1", "content": "content of post 1", "id": 1}, {"title": "favorite foods", "content": "I like pizza", "id": 2}]
Proper error handling is critical when connecting to the database. In this section, we configure the connection to return query results as dictionaries using the RealDictCursor. This makes it easier to work with query results. Additionally, the code demonstrates the use of a try/except block to capture and handle connection errors:
Copy
Ask AI
import psycopg2from psycopg2.extras import RealDictCursorfrom fastapi import FastAPIfrom pydantic import BaseModelapp = FastAPI()class Post(BaseModel): title: str content: str published: bool = Truetry: # Connect to the 'fastapi' database on localhost with given user credentials conn = psycopg2.connect( host='localhost', database='fastapi', user='postgres', password='password123', cursor_factory=RealDictCursor ) cursor = conn.cursor() print("Database connection was successful!")except Exception as error: print("Connecting to database failed") print("Error:", error)my_posts = [ {"title": "title of post 1", "content": "content of post 1", "id": 1}, {"title": "favorite foods", "content": "I like pizza", "id": 2}]
For enhanced readability and maintainability, consider using environment variables to store sensitive database credentials instead of hard-coding them.
Handling Connection Failures with a Retry Mechanism
Database connection attempts may fail temporarily—for example, if the database service has not fully started. In these cases, implementing a retry mechanism can be especially useful. The code below demonstrates a looping construct that continuously attempts to connect until successful, with a 2-second delay between each attempt.
Copy
Ask AI
import timeimport psycopg2from psycopg2.extras import RealDictCursorfrom fastapi import FastAPIfrom pydantic import BaseModelapp = FastAPI()class Post(BaseModel): title: str content: str published: bool = True# Continuously attempt to connect to the database until successfulwhile True: try: conn = psycopg2.connect( host='localhost', database='fastapi', user='postgres', password='password123', cursor_factory=RealDictCursor ) cursor = conn.cursor() print("Database connection was successful!") break # Exit loop on successful connection except Exception as error: print("Connecting to database failed") print("Error:", error) time.sleep(2) # Wait 2 seconds before retryingmy_posts = [ {"title": "title of post 1", "content": "content of post 1", "id": 1}, {"title": "favorite foods", "content": "I like pizza", "id": 2}]
Implementing a retry mechanism not only ensures your application waits for a stable connection but also provides resilience against transient network or service issues.
Hard-coding database credentials (such as host, database name, user, and password) is considered a security risk. This approach:
Exposes sensitive information if committed to version control.
Makes it challenging to switch between development and production environments where credentials differ.
For production-level applications, use environment variables or a configuration manager to handle sensitive information securely. This practice enhances both security and flexibility.With a robust database connection and proper error handling in place, you can now extend your application by executing more SQL commands and building API endpoints that interact seamlessly with your PostgreSQL database.