149. Context Managers for Database Connections

Here are 10 Python code snippets demonstrating how to use context managers to manage database connections. These examples will help you manage resources (like database connections) more efficiently by ensuring that connections are properly opened and closed, even when exceptions occur.

1. Basic Context Manager for Database Connections

This snippet creates a basic context manager for managing a database connection.

import sqlite3

class DatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_name)
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if self.connection:
            self.connection.close()

# Usage
with DatabaseConnection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM users')
    print(cursor.fetchall())

Explanation:

• __enter__: Establishes the connection.

• __exit__: Ensures that the connection is closed after the block of code is executed.

---

2. Using Context Manager with Transactions

Automatically managing transactions in a database connection.

import sqlite3

class TransactionalDatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_name)
        self.connection.isolation_level = None  # Autocommit mode
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            self.connection.rollback()  # Rollback on exception
        else:
            self.connection.commit()  # Commit if no exception
        self.connection.close()

# Usage
with TransactionalDatabaseConnection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute('INSERT INTO users (name) VALUES (?)', ('Alice',))
    print("Transaction successful")

Explanation:

• The context manager handles committing the transaction unless there is an exception, in which case it rolls back.

---

3. Using a Context Manager with PostgreSQL

Managing a PostgreSQL database connection using the psycopg2 library.

import psycopg2

class PostgreSQLConnection:
    def __init__(self, db_config):
        self.db_config = db_config
        self.connection = None

    def __enter__(self):
        self.connection = psycopg2.connect(**self.db_config)
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if self.connection:
            self.connection.close()

# Usage
db_config = {
    'dbname': 'example_db',
    'user': 'postgres',
    'password': 'password',
    'host': 'localhost'
}

with PostgreSQLConnection(db_config) as conn:
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM users')
    print(cursor.fetchall())

Explanation:

• This example uses psycopg2 for PostgreSQL, and the context manager handles connection opening and closing automatically.

---

4. Context Manager with Error Handling

Automatically handling exceptions when interacting with the database.

import sqlite3

class SafeDatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_name)
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print(f"Error occurred: {exc_value}")
            self.connection.rollback()
        else:
            self.connection.commit()
        self.connection.close()

# Usage
with SafeDatabaseConnection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM non_existent_table')  # This will raise an exception

Explanation:

• This context manager ensures that errors are logged, and the connection is rolled back if there is a problem with the query.

---

5. Context Manager for MongoDB (PyMongo)

Managing a MongoDB connection using a custom context manager.

from pymongo import MongoClient

class MongoDBConnection:
    def __init__(self, uri):
        self.uri = uri
        self.client = None

    def __enter__(self):
        self.client = MongoClient(self.uri)
        return self.client

    def __exit__(self, exc_type, exc_value, traceback):
        if self.client:
            self.client.close()

# Usage
with MongoDBConnection('mongodb://localhost:27017') as client:
    db = client['example_db']
    collection = db['users']
    result = collection.find_one({'name': 'Alice'})
    print(result)

Explanation:

• This example uses PyMongo for MongoDB connections, ensuring that the connection is automatically closed after use.

---

6. Context Manager for SQLite with Custom Queries

Using a context manager to execute queries in a SQLite database.

import sqlite3

class QueryExecutor:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_name)
        return self.connection.cursor()

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print(f"Query failed: {exc_value}")
        self.connection.commit()
        self.connection.close()

# Usage
with QueryExecutor('example.db') as cursor:
    cursor.execute('CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)')
    cursor.execute('INSERT INTO users (name) VALUES (?)', ('John',))

Explanation:

• This context manager allows for executing queries and ensures the connection is properly closed.

---

7. Context Manager for Automatic Database Cleanup

Automatically cleaning up database tables after use.

import sqlite3

class CleanupDatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_name)
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            self.connection.rollback()
        else:
            # Clean up: delete temporary data or reset tables
            self.connection.execute('DELETE FROM temp_table')
            self.connection.commit()
        self.connection.close()

# Usage
with CleanupDatabaseConnection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute('CREATE TABLE IF NOT EXISTS temp_table (id INTEGER PRIMARY KEY)')
    cursor.execute('INSERT INTO temp_table (id) VALUES (1)')

Explanation:

• This context manager cleans up temporary data in the database, ensuring that no leftover data remains.

---

8. Context Manager for Connection Pooling

Using connection pooling in a database connection manager.

import psycopg2
from psycopg2 import pool

class ConnectionPoolManager:
    def __init__(self, db_config):
        self.db_config = db_config
        self.pool = None

    def __enter__(self):
        self.pool = psycopg2.pool.SimpleConnectionPool(1, 10, **self.db_config)
        return self.pool.getconn()

    def __exit__(self, exc_type, exc_value, traceback):
        if self.pool:
            self.pool.putconn(self.pool.getconn())

# Usage
db_config = {
    'dbname': 'example_db',
    'user': 'postgres',
    'password': 'password',
    'host': 'localhost'
}

with ConnectionPoolManager(db_config) as conn:
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM users')
    print(cursor.fetchall())

Explanation:

• This example demonstrates how to use connection pooling to manage connections to the database.

---

9. Context Manager for MongoDB Transactions

Using a context manager to handle MongoDB transactions.

from pymongo import MongoClient
from pymongo.errors import OperationFailure

class MongoDBTransaction:
    def __init__(self, uri):
        self.uri = uri
        self.client = None

    def __enter__(self):
        self.client = MongoClient(self.uri)
        self.db = self.client['example_db']
        self.session = self.client.start_session()
        return self.db

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            self.session.abort_transaction()
        else:
            self.session.commit_transaction()
        self.client.close()

# Usage
with MongoDBTransaction('mongodb://localhost:27017') as db:
    try:
        db.users.insert_one({'name': 'Alice'})
        db.orders.insert_one({'order': '123', 'user': 'Alice'})
    except OperationFailure:
        print("Transaction failed")

Explanation:

• This example uses PyMongo for MongoDB transactions with a context manager.

---

10. Context Manager for Redis Connections

Managing Redis connections with a custom context manager.

import redis

class RedisConnection:
    def __init__(self, host='localhost', port=6379):
        self.host = host
        self.port = port
        self.connection = None

    def __enter__(self):
        self.connection = redis.StrictRedis(host=self.host, port=self.port)
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if self.connection:
            self.connection.close()

# Usage
with RedisConnection() as redis_conn:
    redis_conn.set('key', 'value')
    print(redis_conn.get('key'))