"""
Configuration management for pycopg.
Loads database credentials from environment variables or .env file.
"""
import os
import re
from dataclasses import dataclass, field
from os import PathLike
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
from pycopg.exceptions import ConfigurationError
__all__ = ["Config"]
# python-dotenv is optional
try:
from dotenv import load_dotenv
HAS_DOTENV = True
except (
ImportError
): # pragma: no cover — python-dotenv ImportError only fires when package absent; --all-extras always installs it
HAS_DOTENV = False
def load_dotenv( # type: ignore[misc] # fallback shape need not mirror python-dotenv's full stub (stream/verbose/interpolate/encoding) — genuinely unstubbed optional-dependency edge (D-01)
dotenv_path: str | PathLike[str] | None = None,
*,
override: bool = False,
**kwargs: Any,
) -> bool:
"""No-op when python-dotenv is not installed."""
return False
# Safe GUC option key: must be a valid identifier-like name (WR-04)
_SAFE_OPTION_KEY_RE = re.compile(r"^[a-z_][a-z0-9_.]*$", re.IGNORECASE)
def _validate_libpq_option(key: str, value: str) -> None:
"""Validate a libpq GUC option key/value before interpolating into the options string.
Rejects keys that do not match the GUC identifier pattern and values that
contain characters that are special in the libpq options string (space,
single-quote, backslash). These characters could inject additional -c
flags or break the options string parsed by libpq (WR-04).
Parameters
----------
key : str
GUC parameter name (e.g. 'search_path').
value : str
GUC parameter value.
Raises
------
ConfigurationError
If the key or value contains unsafe characters.
"""
if not _SAFE_OPTION_KEY_RE.match(key):
raise ConfigurationError(
f"Unsafe option key {key!r}: must match pattern [a-z_][a-z0-9_.]* "
"(e.g. 'search_path', 'application_name')"
)
if " " in value or "'" in value or "\\" in value:
raise ConfigurationError(
f"Unsafe option value for {key!r}: value must not contain spaces, "
"single-quotes, or backslashes (libpq options injection)"
)
[docs]
@dataclass
class Config:
"""Database connection configuration.
Can be created from:
- Individual parameters
- DATABASE_URL environment variable
- .env file
"""
host: str = "localhost"
port: int = 5432
database: str = "postgres"
user: str = "postgres"
password: str = ""
sslmode: str | None = None
options: dict[str, str] = field(default_factory=dict)
statement_timeout: int | None = None
default_batch_size: int = 1000
[docs]
@classmethod
def from_url(cls, url: str) -> "Config":
"""Create config from a database URL.
Parameters
----------
url : str
PostgreSQL connection URL. Formats supported:
- postgresql://user:pass@host:port/dbname
- postgresql+asyncpg://user:pass@host:port/dbname
- postgres://user:pass@host:port/dbname
Returns
-------
Config
Config instance.
"""
# Normalize URL scheme
url = url.replace("postgresql+asyncpg://", "postgresql://")
url = url.replace("postgres://", "postgresql://")
parsed = urlparse(url)
# Extract query params as options
options = {}
if parsed.query:
for param in parsed.query.split("&"):
if "=" in param:
key, value = param.split("=", 1)
options[key] = value
# Extract statement_timeout from query params
statement_timeout_str = options.pop("statement_timeout", None)
statement_timeout = (
int(statement_timeout_str) if statement_timeout_str else None
)
return cls(
host=parsed.hostname or "localhost",
port=parsed.port or 5432,
database=parsed.path.lstrip("/") if parsed.path else "postgres",
user=parsed.username or "postgres",
password=parsed.password or "",
sslmode=options.pop("sslmode", None),
options=options,
statement_timeout=statement_timeout,
)
[docs]
@classmethod
def from_env(
cls,
dotenv_path: str | Path | None = None,
*,
load_dotenv_file: bool = True,
) -> "Config":
"""Create config from environment variables.
Looks for DATABASE_URL first, then individual variables:
- DATABASE_URL: Full connection URL
- DB_HOST, PGHOST: Database host
- DB_PORT, PGPORT: Database port
- DB_NAME, PGDATABASE: Database name
- DB_USER, PGUSER: Database user
- DB_PASSWORD, PGPASSWORD: Database password
Parameters
----------
dotenv_path : str or Path, optional
Path to .env file. If None, searches current directory and parents.
load_dotenv_file : bool, optional
Whether to load .env file. Set to False to only use existing
environment variables, by default True.
Returns
-------
Config
Config instance.
"""
if load_dotenv_file:
if dotenv_path:
# When explicit path is given, override existing env vars
load_dotenv(dotenv_path, override=True)
else:
load_dotenv()
# Try DATABASE_URL first
database_url = os.getenv("DATABASE_URL")
if database_url:
return cls.from_url(database_url)
# Fall back to individual env vars
return cls(
host=os.getenv("DB_HOST") or os.getenv("PGHOST") or "localhost",
port=int(os.getenv("DB_PORT") or os.getenv("PGPORT") or "5432"),
database=os.getenv("DB_NAME") or os.getenv("PGDATABASE") or "postgres",
user=os.getenv("DB_USER") or os.getenv("PGUSER") or "postgres",
password=os.getenv("DB_PASSWORD") or os.getenv("PGPASSWORD") or "",
sslmode=os.getenv("DB_SSLMODE") or os.getenv("PGSSLMODE"),
)
@property
def dsn(self) -> str:
"""Generate psycopg-compatible DSN string.
Returns
-------
str
Connection string for psycopg.
"""
parts = [
f"host={self.host}",
f"port={self.port}",
f"dbname={self.database}",
f"user={self.user}",
]
if self.password:
parts.append(f"password={self.password}")
if self.sslmode:
parts.append(f"sslmode={self.sslmode}")
# Add options string if statement_timeout or options are set
options_parts = []
if self.statement_timeout is not None:
options_parts.append(f"-c statement_timeout={self.statement_timeout}")
for key, value in self.options.items():
_validate_libpq_option(key, value)
options_parts.append(f"-c {key}={value}")
if options_parts:
parts.append(f"options={' '.join(options_parts)}")
return " ".join(parts)
@property
def url(self) -> str:
"""Generate SQLAlchemy-compatible URL using psycopg v3.
Returns
-------
str
PostgreSQL URL for SQLAlchemy with psycopg driver.
"""
auth = f"{self.user}:{self.password}" if self.password else self.user
# Use postgresql+psycopg for psycopg v3 (not psycopg2)
base = f"postgresql+psycopg://{auth}@{self.host}:{self.port}/{self.database}"
if self.sslmode:
base += f"?sslmode={self.sslmode}"
return base
@property
def async_url(self) -> str:
"""Generate SQLAlchemy-compatible async URL using psycopg v3.
Mirrors :attr:`url` exactly but emits the async psycopg driver scheme
so SQLAlchemy builds an async engine. ``url`` itself is unchanged and
still uses the sync driver.
Returns
-------
str
PostgreSQL URL for SQLAlchemy with the async psycopg driver.
"""
auth = f"{self.user}:{self.password}" if self.password else self.user
# Use postgresql+psycopg_async for the async psycopg v3 driver
base = (
f"postgresql+psycopg_async://{auth}@{self.host}:{self.port}/{self.database}"
)
if self.sslmode:
base += f"?sslmode={self.sslmode}"
return base
[docs]
def connect_params(self) -> dict[str, Any]:
"""Get connection parameters as dict for psycopg.connect().
Returns
-------
dict
Dict with host, port, dbname, user, password.
"""
params: dict[str, Any] = {
"host": self.host,
"port": self.port,
"dbname": self.database,
"user": self.user,
}
if self.password:
params["password"] = self.password
if self.sslmode:
params["sslmode"] = self.sslmode
# Add PostgreSQL options string if needed
options_parts = []
if self.statement_timeout is not None:
options_parts.append(f"-c statement_timeout={self.statement_timeout}")
for key, value in self.options.items():
_validate_libpq_option(key, value)
options_parts.append(f"-c {key}={value}")
if options_parts:
params["options"] = " ".join(options_parts)
return params
[docs]
def with_database(self, database: str) -> "Config":
"""Create a new config pointing to a different database.
Parameters
----------
database : str
Target database name.
Returns
-------
Config
New Config instance with updated database.
"""
return Config(
host=self.host,
port=self.port,
database=database,
user=self.user,
password=self.password,
sslmode=self.sslmode,
options=self.options.copy(),
statement_timeout=self.statement_timeout,
default_batch_size=self.default_batch_size,
)
def __repr__(self) -> str:
"""Return string representation of the Config instance."""
return f"Config(host={self.host!r}, port={self.port}, database={self.database!r}, user={self.user!r})"