ladybug_adbc_client.py

"""Columnar.tech-style clients for the LadybugDB Flight server.

Two flavours, same zero-JSON Arrow result:

1. **ADBC (preferred, exactly the recipe you posted)** -- ``adbc_driver_flightsql``
   DBAPI. ``cursor.execute()`` *typically* takes SQL, but the Flight SQL wire
   field is just a string, so we pass **Cypher** straight through::

       import adbc_driver_flightsql.dbapi as flight_sql
       with flight_sql.connect("grpc://localhost:50051") as conn:
           with conn.cursor() as cur:
               cur.execute("MATCH (u:User) RETURN u.name, u.age ORDER BY u.id")
               table = cur.fetch_arrow_table()   # <- columnar.tech method

2. **Plain Flight** -- ``pyarrow.flight.FlightClient`` for environments where
   the ADBC driver is unavailable. Same server, same tables.

Run ``python ladybug_flight_server.py`` first, then::

    uv run python ladybug_adbc_client.py
    uv run python ladybug_adbc_client.py --query "MATCH (u:User) RETURN u.name"
    uv run python ladybug_adbc_client.py --demo --no-adbc   # plain Flight only
"""

from __future__ import annotations

import argparse

import pyarrow as pa
import pyarrow.flight as flight

from ladybug_flight_server import DEMO_QUERIES


def query_cypher_adbc(uri: str, cypher: str) -> pa.Table:
    """Run Cypher over ADBC Flight SQL; return native Arrow (columnar.tech)."""
    import adbc_driver_flightsql.dbapi as flight_sql  # lazy: plain-Flight use skips it

    with flight_sql.connect(uri) as conn:
        with conn.cursor() as cur:
            cur.execute(cypher)  # SQL-typed param, Cypher content -- works, see module doc
            return cur.fetch_arrow_table()


def query_cypher_flight(uri: str, cypher: str) -> pa.Table:
    """Run Cypher over plain Arrow Flight; return native Arrow."""
    client = flight.FlightClient(uri)
    descriptor = flight.FlightDescriptor.for_command(cypher.encode("utf-8"))
    info = client.get_flight_info(descriptor)
    reader = client.do_get(info.endpoints[0].ticket)
    return reader.read_all()


def main() -> None:
    ap = argparse.ArgumentParser(description="LadybugDB ADBC/Flight client demo")
    ap.add_argument("--uri", default="grpc://localhost:50051")
    ap.add_argument("--query", default="MATCH (u:User) RETURN u.name, u.age ORDER BY u.id")
    ap.add_argument("--demo", action="store_true", help="run all DEMO_QUERIES")
    ap.add_argument("--no-adbc", action="store_true", help="use plain Flight, not ADBC")
    args = ap.parse_args()

    run = query_cypher_flight if args.no_adbc else query_cypher_adbc
    mode = "plain Flight" if args.no_adbc else "ADBC FlightSQL"
    queries = DEMO_QUERIES if args.demo else {"query": args.query}
    for label, cypher in queries.items():
        print(f"\n🔌 [{mode}] {label}: {cypher[:100]}")
        table = run(args.uri, cypher)
        print(f"📊 {table.num_rows} rows x {len(table.schema)} cols, "
              f"{table.get_total_buffer_size()} bytes")
        try:
            print(table.to_pandas().to_string(index=False))
        except ImportError:
            for row in table.to_pylist():
                print(row)


if __name__ == "__main__":
    main()