如何创建PostgreSQL连接-执行-断开通用函数优化重复代码

Absolutely great catch on that repeated boilerplate code—extracting that into a reusable utility is exactly the right move to clean up your codebase and make it more maintainable. Let's walk through how to build a generic database execution function that handles connection setup, cursor management, cleanup, and error handling all in one place.

The Core Idea: Context Managers + Generic Execution #

Psycopg2 works seamlessly with Python's with statement (context managers), which automatically handles resource cleanup (closing connections/cursors) even if exceptions occur. We can wrap this logic into a single function that accepts your query, parameters, and optional settings, then returns the appropriate result.

Optimized Code Implementation #

First, here's the generic database utility function, followed by your refactored business logic functions:

import psycopg2
import psycopg2.extras
from loguru import logger

DATABASE_CONNECTION = {
    "host": "TEST",
    "database": "TEST",
    "user": "TEST",
    "password": "TEST"
}

def execute_db_query(sql, params=None, cursor_factory=None, autocommit=False):
    """
    Generic function to execute PostgreSQL queries with automatic resource management.
    :param sql: SQL query string to execute
    :param params: Tuple of parameters for parameterized queries (default: empty tuple)
    :param cursor_factory: Optional cursor type (e.g., psycopg2.extras.DictCursor)
    :param autocommit: Enable autocommit for write operations (default: False)
    :return: Query result (rows for SELECT, row count for writes, None on failure)
    """
    params = params or ()
    try:
        # Use context manager for connection: auto-closes on exit
        with psycopg2.connect(**DATABASE_CONNECTION) as conn:
            conn.autocommit = autocommit
            # Use context manager for cursor: auto-closes on exit
            with conn.cursor(cursor_factory=cursor_factory) as cur:
                cur.execute(sql, params)
                
                # Return appropriate result based on query type
                query_type = sql.strip().upper()
                if query_type.startswith("SELECT"):
                    return cur.fetchall()
                elif query_type.startswith(("INSERT", "UPDATE", "DELETE")):
                    return cur.rowcount
                return None
    except (Exception, psycopg2.DatabaseError) as error:
        logger.exception("Database query failed: %s" % error)
        # Context manager handles rollback automatically on exception
        return None

# --- Refactored business logic functions ---
def register_datas(store, data):
    """Register a data to database"""
    sql = "INSERT INTO public.store_items (store, name) VALUES (%s, %s);"
    rowcount = execute_db_query(sql, (store, data["name"]), autocommit=True)
    return bool(rowcount) if rowcount is not None else False

def get_all_keywords(filter_type):
    """Get all keywords by filter type"""
    sql = "SELECT keyword FROM public.keywords WHERE filter_type = %s;"
    results = execute_db_query(sql, (filter_type,), cursor_factory=psycopg2.extras.DictCursor)
    return [row["keyword"] for row in results] if results is not None else []

def check_if_store_exists(store):
    """Check if the store exists in database"""
    sql = "SELECT store FROM public.store_config WHERE store = %s;"
    result = execute_db_query(sql, (store,))
    return bool(result) if result is not None else False

Key Improvements Explained #

• No More Manual Cleanup: The with statements handle closing connections and cursors automatically, even if an error occurs—you can delete all those repetitive close() calls.
• Centralized Error Handling: All exception logging and rollback logic lives in one place, so you don't repeat it in every function.
• Flexibility: The generic function supports different cursor types (like DictCursor) and autocommit for write operations, so it adapts to your specific needs.
• Cleaner Business Logic: Your original functions now focus solely on their core purpose (defining the query and processing results) instead of database plumbing.

If you ever need to add more query types (like fetchone() instead of fetchall()), you can easily extend the generic function with optional parameters to support that.

内容的提问来源于stack exchange，提问作者PythonNewbie