smoke.py

# A simple smoke test for the DB driver.

import argparse
import concurrent.futures
import functools
import logging
import sys

import pandas
from rich.console import Console
from rich.table import Table

from wherobots.db import connect, connect_direct, errors, ProgressInfo
from wherobots.db.connection import Connection
from wherobots.db.constants import DEFAULT_ENDPOINT, DEFAULT_SESSION_TYPE
from wherobots.db.region import Region
from wherobots.db.runtime import Runtime
from wherobots.db.session_type import SessionType
from wherobots.db.models import Store, StoreResult


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--api-key-file", help="File containing the API key")
    parser.add_argument("--token-file", help="File containing the token")
    parser.add_argument("--region", help="Region to connect to (ie. aws-us-west-2)")
    parser.add_argument("--runtime", help="Runtime type (ie. tiny)")
    parser.add_argument("--version", help="Runtime version (ie. latest)")
    parser.add_argument("--force-new", action="store_true")
    parser.add_argument("--store", action="store_true")
    parser.add_argument(
        "--session-type",
        help="Type of session to create",
        default=DEFAULT_SESSION_TYPE,
        choices=[st.value for st in SessionType],
    )
    parser.add_argument(
        "--debug",
        help="Enable debug logging",
        action="store_const",
        const=logging.DEBUG,
        default=logging.INFO,
    )
    parser.add_argument(
        "--api-endpoint",
        help="Wherobots API endpoint to request a SQL session from",
        default=DEFAULT_ENDPOINT,
    )
    parser.add_argument("--ws-url", help="Direct URL to connect to")
    parser.add_argument(
        "--shutdown-after-inactive-seconds",
        help="Request a specific SQL Session expiration timeout (in seconds)",
    )
    parser.add_argument(
        "--wide", help="Enable wide output", action="store_const", const=80, default=30
    )
    parser.add_argument(
        "--progress",
        help="Enable execution progress reporting",
        action="store_true",
    )
    parser.add_argument("sql", nargs="+", help="SQL query to execute")
    args = parser.parse_args()

    logging.basicConfig(
        stream=sys.stdout,
        level=args.debug,
        format="%(asctime)s.%(msecs)03d %(levelname)s %(name)20s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    logging.getLogger("websockets.protocol").setLevel(args.debug)

    api_key = None
    token = None
    headers = None
    store = None

    if args.api_key_file:
        with open(args.api_key_file) as f:
            api_key = f.read().strip()
        headers = {"X-API-Key": api_key}

    if args.token_file:
        with open(args.token_file) as f:
            token = f.read().strip()
        headers = {"Authorization": f"Bearer {token}"}

    if args.store:
        store = Store.for_download()
        logging.info("Will requests for results to be stored in cloud storage.")

    if args.ws_url:
        conn_func = functools.partial(connect_direct, uri=args.ws_url, headers=headers)
    else:
        conn_func = functools.partial(
            connect,
            host=args.api_endpoint,
            token=token,
            api_key=api_key,
            shutdown_after_inactive_seconds=args.shutdown_after_inactive_seconds,
            wait_timeout=900,
            runtime=Runtime(args.runtime) if args.runtime else Runtime.MICRO,
            region=Region(args.region) if args.region else Region.AWS_US_WEST_2,
            version=args.version,
            force_new=args.force_new,
            session_type=SessionType(args.session_type),
        )

    def render_df(df: pandas.DataFrame) -> Table:
        table = Table(show_header=True)
        table.add_column("#")
        for column in df.columns:
            table.add_column(column, max_width=args.wide, no_wrap=True)
        for row in df.itertuples(name=None):
            r = [str(x) for x in row]
            table.add_row(*r)
        return table

    def render_stored(sr: StoreResult) -> Table:
        table = Table(show_header=True)
        table.add_column("URI")
        table.add_column("Size", justify="right")
        table.add_row(sr.result_uri, str(sr.size))
        return table

    def render(results: pandas.DataFrame | StoreResult) -> None:
        if isinstance(results, StoreResult):
            Console().print(render_stored(results))
        else:
            Console().print(render_df(results))

    def execute(conn: Connection, sql: str) -> pandas.DataFrame | StoreResult:
        with conn.cursor() as cursor:
            cursor.execute(sql, store=store)
            if args.store:
                return cursor.get_store_result()
            else:
                return cursor.fetchall()

    try:
        with conn_func() as conn:
            if args.progress:
                console = Console(stderr=True)

                def _on_progress(info: ProgressInfo) -> None:
                    pct = (
                        f"{info.tasks_completed / info.tasks_total * 100:.0f}%"
                        if info.tasks_total
                        else "?"
                    )
                    console.print(
                        f"  [dim]\\[progress][/dim] "
                        f"[bold]{pct}[/bold]  "
                        f"{info.tasks_completed}/{info.tasks_total} tasks "
                        f"[dim]({info.tasks_active} active)[/dim]  "
                        f"[dim]{info.execution_id[:8]}[/dim]",
                        highlight=False,
                    )

                conn.set_progress_handler(_on_progress)

            with concurrent.futures.ThreadPoolExecutor() as pool:
                futures = [pool.submit(execute, conn, s) for s in args.sql]
                for future in concurrent.futures.as_completed(futures):
                    render(future.result())
    except errors.Error as e:
        sys.stderr.write(f"\n{e}\n")
        sys.exit(1)