from browseruse_bench.utils import (
    get_db_connection,
    db_cursor,
    init_db_schema,
    create_run_record,
    update_run_status,
    insert_task_results,
    insert_eval_metric,
)

def get_db_connection(
    database_url: Optional[str] = None
) -> psycopg2.extensions.connection

@contextmanager
def db_cursor(
    database_url: Optional[str] = None
) -> Generator[psycopg2.extensions.cursor, None, None]

from browseruse_bench.utils import db_cursor

with db_cursor() as cur:
    cur.execute("SELECT * FROM benchmark_run LIMIT 10")
    rows = cur.fetchall()
# Auto-commits on success, auto-rollbacks on exception

def init_db_schema(database_url: Optional[str] = None) -> None

def create_run_record(
    agent_name: str,
    benchmark_name: str,
    mode: str,
    count: Optional[int],
    task_ids_json: Optional[str],
    started_at: datetime,
    status: str,
    experiment_path: Optional[str],
    database_url: Optional[str] = None
) -> str

def update_run_status(
    run_uuid: str,
    status: str,
    finished_at: datetime,
    error_message: Optional[str] = None,
    database_url: Optional[str] = None
) -> None

def insert_task_results(
    rows: Iterable[Dict[str, Any]],
    database_url: Optional[str] = None
) -> None

def insert_eval_metric(
    row: Dict[str, Any],
    database_url: Optional[str] = None
) -> None