Function query

misc/python/materialize/cli/mz_workload_capture.py:46–78

(
    conn: psycopg.Connection,
    sql: LiteralString | bytes | SQL | Composed,
    timeout: int | None = None,
)

def query(
    conn: psycopg.Connection,
    sql: LiteralString | bytes | SQL | Composed,
    timeout: int | None = None,
) -> list[Any]:
    if VERBOSE:
        print(f"\n> {to_sql_string(sql, conn)}")

    cancel_timer: threading.Timer | None = None

    try:
        if timeout is not None:
            cancel_timer = threading.Timer(timeout, conn.cancel)
            cancel_timer.start()

        with conn.cursor() as cur:
            cur.execute(sql)
            return cur.fetchall()
    except psycopg.errors.QueryCanceled:
        if VERBOSE:
            print("Too slow")
        return []
    except KeyboardInterrupt:
        conn.cancel()
        if VERBOSE:
            raise
        print(f"\n> {to_sql_string(sql, conn)}")
        with conn.cursor() as cur:
            cur.execute(sql)
            return cur.fetchall()
    finally:
        if cancel_timer is not None:
            cancel_timer.cancel()


def attach_source_statistics(