File size: 4,342 Bytes
9abace2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """
PostgreSQL Operations
Part of SOVEREIGN PYTHON LLM ENGINE
PostgreSQL database operations using asyncpg.
"""
from typing import Any
from dataclasses import dataclass
from ...core.evidence import WORMLedger
@dataclass
class PostgresQueryResult:
"""Postgres query result"""
rows: list[dict[str, Any]]
row_count: int
columns: list[str]
class PostgresOperations:
"""
PostgreSQL database operations.
Uses asyncpg for async native Postgres protocol.
"""
def __init__(
self,
connection_string: str,
worm_ledger: WORMLedger | None = None
):
"""
Initialize Postgres operations.
Args:
connection_string: Postgres connection string
worm_ledger: Optional WORM ledger
"""
self.connection_string = connection_string
self.worm_ledger = worm_ledger
async def query(self, sql: str, *params) -> PostgresQueryResult:
"""
Execute SELECT query.
Args:
sql: SQL query
params: Query parameters
Returns:
PostgresQueryResult with rows
"""
try:
import asyncpg
except ImportError:
raise ImportError("asyncpg is required. Install: pip install asyncpg")
conn = await asyncpg.connect(self.connection_string)
try:
rows = await conn.fetch(sql, *params)
# Convert to dicts
columns = list(rows[0].keys()) if rows else []
row_dicts = [dict(row) for row in rows]
result = PostgresQueryResult(
rows=row_dicts,
row_count=len(row_dicts),
columns=columns
)
# Log to WORM
if self.worm_ledger:
await self.worm_ledger.append({
"event": "postgres_query",
"sql": sql,
"row_count": result.row_count
})
return result
finally:
await conn.close()
async def execute(self, sql: str, *params) -> str:
"""
Execute INSERT/UPDATE/DELETE.
Args:
sql: SQL statement
params: Parameters
Returns:
Status string (e.g., "INSERT 0 1")
"""
try:
import asyncpg
except ImportError:
raise ImportError("asyncpg is required")
conn = await asyncpg.connect(self.connection_string)
try:
status = await conn.execute(sql, *params)
# Log to WORM
if self.worm_ledger:
await self.worm_ledger.append({
"event": "postgres_execute",
"sql": sql,
"status": status
})
return status
finally:
await conn.close()
async def transaction(self, statements: list[tuple[str, tuple]]) -> None:
"""
Execute multiple statements in transaction.
Args:
statements: List of (sql, params) tuples
"""
try:
import asyncpg
except ImportError:
raise ImportError("asyncpg is required")
conn = await asyncpg.connect(self.connection_string)
try:
async with conn.transaction():
for sql, params in statements:
await conn.execute(sql, *params)
finally:
await conn.close()
# Tool registration helpers
async def postgres_query_tool(connection_string: str, sql: str, params: list | None = None) -> dict:
"""Execute Postgres query"""
ops = PostgresOperations(connection_string)
result = await ops.query(sql, *(params or []))
return {
"rows": result.rows,
"row_count": result.row_count,
"columns": result.columns
}
async def postgres_execute_tool(connection_string: str, sql: str, params: list | None = None) -> dict:
"""Execute Postgres statement"""
ops = PostgresOperations(connection_string)
status = await ops.execute(sql, *(params or []))
return {
"status": status
}
|