aa_psycopg 0.1.8

Internal Partenamut library for PostgreSQL access using Psycopg 3

aa-psycopg

Internal Partenamut library for PostgreSQL access using Psycopg 3.
Provides a simple and efficient interface for interacting with PostgreSQL using:

• A direct connection (PostgreSQLConnection)
• A connection pool (PostgreSQLPool)

Supports query execution (SELECT, INSERT, UPDATE, DELETE), transactions, schema caching, and query runtime statistics.

---

Installation

If you want to install only psycopg functionality (see: Usage of Psycopg for Direct Connection or Pooling)

pip install aa-psycopg

If you also want to install features for validating query configuration with pydantic (see Usage of QueryConfig for Query Management)

pip install aa-psycopg[pydantic]

Requirements

• Python 3.11+
• psycopg 3
• psycopg_pool

---

Features

• Easy-to-use wrapper for psycopg3.
• Connection pooling via psycopg_pool.
• Method for chunked/streamed query execution (read_in_chunks) to avoid loading large datasets into memory.
• Automatic query runtime tracking.
• Safe connection string handling (masks passwords in logs).
• Schema caching for executed queries.

---

API overview

• ping(retries=0, timeout=60, query_name="ping") -> bool
  Test database connectivity.
• read(query, params=None, query_name=None) -> list[dict]
  Execute a SELECT query. Set query_name to track query runtime statistics.
• read_in_chunks(query, params=None, chunk_size=500000, query_name=None) -> Generator[list[dict]]
  Execute a SELECT query and yield results in batches (streaming mode).
• write(query, params=None, returning=False, query_name=None) -> list[dict] | None
  Execute INSERT/UPDATE/DELETE (optionally returning results). Set query_name to track query runtime statistics.
• execute_transaction(queries_params, query_name=None)
  Run multiple queries inside a transaction. Set query_name to track query runtime statistics.
• get_stats() -> dict
  Retrieve runtime stats (execution time & call count per query_name).
• get_schema(query) -> list[psycopg.Column]
  Get schema for a previously executed query.

Usage of Psycopg for Direct Connection or Pooling

Direct Connection

import logging
import json

from aa_psycopg.connection import PostgreSQLConnection

# Configure logging (can adjust level or handlers as needed)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

with PostgreSQLConnection(
    user="myuser",
    password="mypassword",
    host="localhost",
    port=5432,
    db="mydatabase"
) as client:
    # Check if database is alive
    client.ping()

    # Fetch results
    results = client.read("SELECT * FROM my_table", query_name="fetch_all")
    logger.info("Fetched results: %s", results[:5])

    # Stream rows in chunks (efficient for very large result sets)
    for i, chunk in client.read_in_chunks(
        "SELECT * FROM big_table",
        chunk_size=5000,  # how many rows to fetch per chunk
        query_name="stream_big_table",
    ):
        logger.info("Fetched results of chunk %s: %s", i + 1, chunk[:5])

    # Write data
    client.write(
        "INSERT INTO my_table (name) VALUES (%(name)s)",
        params={"name": "example"},
        query_name="insert_row"
    )
    logger.info("Inserted row into my_table")

    # Transaction
    client.execute_transaction([
        ("INSERT INTO my_table (name) VALUES (%(name)s)", {"name": "row1"}),
        ("INSERT INTO my_table (name) VALUES (%(name)s)", {"name": "row2"}),
    ], query_name="bulk_insert")
    logger.info("Executed bulk insert transaction")

    # Query stats
    stats = client.get_stats()
    logger.info("Query stats: %s", json.dumps(stats, indent=4))

---

Connection Pool

import logging
import json
from aa_psycopg.pool import PostgreSQLPool

# Configure logging (can be redirected to CloudWatch or file)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

with PostgreSQLPool(
    user="myuser",
    password="mypassword",
    host="localhost",
    port=5432,
    db="mydatabase",
    min_size=0,  # pool starts with 0 connections, only creating connection with first query
    max_size=1  # pool will not exceed 1 connection
) as client:
    # Check connectivity
    client.ping()

    # Fetch results
    results = client.read("SELECT * FROM users WHERE active = true", query_name="active_users")
    logger.info("Fetched results: %s", results[:5])

    # Stream results in chunks (efficient for very large result sets)
    for i, chunk in client.read_in_chunks(
        "SELECT * FROM big_table",
        chunk_size=5000,  # how many rows to fetch per chunk
        query_name="stream_big_table",
    ):
        logger.info("Fetched results of chunk %s: %s", i + 1, chunk[:5])

    # Insert with RETURNING
    new_ids = client.write(
        "INSERT INTO users (name) VALUES (%(name)s) RETURNING id",
        params={"name": "Alice"},
        returning=True,
        query_name="insert_user"
    )
    logger.info("Inserted new user IDs: %s", new_ids)

    # Bulk transaction
    client.execute_transaction([
        ("UPDATE users SET active=false WHERE id=%(id)s", {"id": 1}),
        ("DELETE FROM users WHERE active=false", None),
    ], query_name="cleanup")
    logger.info("Executed cleanup transaction")

    # Stats (dump as JSON)
    stats = client.get_stats()
    logger.info("Query stats: %s", json.dumps(stats, indent=4))

---

Usage of QueryConfig for Query Management

QueryConfig is a Pydantic-based configuration class that allows you to manage SQL queries and their parameters from either a .sql file or a raw query string. It can also load configurations from YAML files for easier organization.

---

Installation of optional dependency

QueryConfig is optional. To enable it, install the extra dependencies:

uv pip install aa-psycopg[pydantic]

This will pull in pydantic.

---

Basic Usage

Load from a .sql file

from pathlib import Path
from aa_psycopg.query_config import QueryConfig
cfg = QueryConfig(path_query=Path("queries/get_users.sql"), params={"user_id": 123})
print(cfg.query)   # Contents of the SQL file
print(cfg.params)  # {"user_id": 123}

Load directly from a raw query string

from aa_psycopg.query_config import QueryConfig
cfg = QueryConfig(query="SELECT * FROM users WHERE id = %(user_id)s", params={"user_id": 123})
print(cfg.query)  # The query string

---

Loading from YAML

You can define a query configuration in a .yaml file:

# query_config.yaml
path_query: queries/get_users.sql
params:
  user_id: 123

OR

# query_config.yaml
query: |
  SELECT * FROM users 
  WHERE id = %(user_id)s
params:
  user_id: 123

Load it using from_yaml:

from aa_psycopg.query_config import QueryConfig
cfg = QueryConfig.from_yaml("query_config.yaml")
print(cfg.query)   # SQL contents from the file
print(cfg.params)  # {"user_id": 123}

---

Validation

• You must provide exactly one of query or path_query.
• If a .sql file is given, it must exist and have a .sql extension.
• All parameters defined in the SQL (using %(param)s style) must match the keys in params. If they don’t, a ValueError is raised with details. --> """

# Example query with a parameter mismatch (SQL expects "id", params gives "user_id")
from aa_psycopg.query_config import QueryConfig
try:
    cfg = QueryConfig(
        query="SELECT * FROM users WHERE id = %(id)s",
        params={"user_id": 123},  # This will cause a validation error
    )
except ValueError as e:
    print("Validation failed, mismatch detected!")
    print(e)

---

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

Please make sure to update tests as appropriate.

---