Compare the MySQL and MariaDB ADBC drivers against a MariaDB server

compare_adbc_mysql_mariadb.py

#!/usr/bin/env python3
"""Compare the MySQL and MariaDB ADBC drivers against a MariaDB server.

The demo only creates a TEMPORARY table, once per connection.  It prints the
Arrow type, field metadata, and value (or error) returned by each driver.

See the blog post: https://lefred.be/content/when-the-sea-lion-teaches-the-duck-how-to-bark/

This script requires adbc-driver-manager and pyarrow

"""

from __future__ import annotations

import argparse
import os
from dataclasses import dataclass, field
from typing import Any

from adbc_driver_manager import dbapi


CREATE_TABLE = """
CREATE TEMPORARY TABLE adbc_driver_comparison (
    id UUID NOT NULL,
    embedding VECTOR(3),
    ipv4 INET4,
    location POINT,
    precise_value DECIMAL(65,30),
    elapsed TIME(6),
    event_year YEAR,
    document JSON
)
"""

INSERT_ROW = """
INSERT INTO adbc_driver_comparison VALUES (
    UUID_v7(),
    VEC_FromText('[1.25,-2.5,3.75]'),
    '192.168.1.10',
    ST_PointFromText('POINT(4.35 50.85)'),
    12345678901234567890123456789012345.123456789012345678901234567890,
    '-838:59:59.123456',
    2026,
    '{"name":"MariaDB","valid":true}'
)
"""

CASES = {
    "UUID (protocol reports CHAR)": "id",
    "VECTOR (protocol reports VARBINARY)": "embedding",
    "INET4 (protocol reports CHAR)": "ipv4",
    "MariaDB geometry": "location",
    "65-digit DECIMAL": "precise_value",
    "negative TIME duration": "elapsed",
    "YEAR": "event_year",
    "JSON query result (protocol reports LONGTEXT)": "document",
}

METADATA_TABLE = f"adbc_metadata_comparison_{os.getpid()}"

CREATE_METADATA_TABLE = """
CREATE TABLE {table} (
    base INT PRIMARY KEY,
    hidden_note VARCHAR(20) INVISIBLE COMMENT 'secret',
    doubled INT AS (base * 2) VIRTUAL,
    document JSON,
    native_uuid UUID,
    embedding VECTOR(3),
    address INET4,
    status ENUM('new', 'active', 'disabled'),
    flags SET('read', 'write', 'admin')
)
"""


@dataclass
class Result:
    arrow_type: str = ""
    metadata: dict[str, str] = field(default_factory=dict)
    value: str = ""
    error: str = ""


def uri(driver: str, args: argparse.Namespace) -> str:
    # Each driver deliberately accepts its own URI scheme.
    return (
        f"{driver}://{args.user}:{args.password}@"
        f"{args.host}:{args.port}/{args.database}"
    )


def display_value(value: Any) -> str:
    if isinstance(value, bytes):
        return "0x" + value.hex()
    return repr(value)


def run_driver(driver: str, args: argparse.Namespace) -> dict[str, Result]:
    results: dict[str, Result] = {}
    try:
        connection = dbapi.connect(
            driver=driver, uri=uri(driver, args), autocommit=True
        )
    except Exception as exc:  # Keep the other driver useful if one is absent.
        return {name: Result(error=f"connection failed: {exc}") for name in CASES}

    with connection, connection.cursor() as cursor:
        cursor.execute(CREATE_TABLE)
        cursor.execute(INSERT_ROW)
        for name, column in CASES.items():
            try:
                cursor.execute(f"SELECT `{column}` FROM adbc_driver_comparison")
                table = cursor.fetch_arrow_table()
                field = table.schema.field(0)
                metadata = {
                    key.decode(): value.decode()
                    for key, value in (field.metadata or {}).items()
                }
                results[name] = Result(
                    arrow_type=str(field.type),
                    metadata=metadata,
                    value=display_value(table.column(0)[0].as_py()),
                )
            except Exception as exc:
                results[name] = Result(error=str(exc).replace("\n", " "))
    return results


def print_results(all_results: dict[str, dict[str, Result]]) -> None:
    for case in CASES:
        print(f"\n=== {case} ===")
        for driver in ("mysql", "mariadb"):
            result = all_results[driver][case]
            print(f"{driver:7}:")
            if result.error:
                print(f"  ERROR: {result.error}")
            else:
                print(f"  Arrow type: {result.arrow_type}")
                print(f"  value:      {result.value}")
                print("  metadata:")
                print_metadata(result.metadata, indent="    ")


def print_metadata(metadata: dict[str, Any], indent: str) -> None:
    if not metadata:
        print(f"{indent}(none)")
        return
    for key, value in metadata.items():
        print(f"{indent}{key}")
        print(f"{indent}  = {value}")


def find_columns(objects: Any) -> dict[str, dict[str, Any]]:
    """Flatten the standard nested ADBC GetObjects response."""
    columns: dict[str, dict[str, Any]] = {}
    for catalog in objects:
        for schema in catalog.get("catalog_db_schemas") or []:
            for table in schema.get("db_schema_tables") or []:
                for column in table.get("table_columns") or []:
                    columns[column["column_name"]] = column
    return columns


def find_constraints(objects: Any) -> list[dict[str, Any]]:
    for catalog in objects:
        for schema in catalog.get("catalog_db_schemas") or []:
            for table in schema.get("db_schema_tables") or []:
                if table.get("table_name") == METADATA_TABLE:
                    return table.get("table_constraints") or []
    return []


def decoded_metadata(field: Any) -> dict[str, str]:
    return {
        key.decode(): value.decode()
        for key, value in (field.metadata or {}).items()
    }


def compact_type(data_type: Any) -> str:
    value = str(data_type)
    if value.startswith("extension<arrow.opaque["):
        marker = "type_name="
        start = value.find(marker) + len(marker)
        end = value.find(",", start)
        return f"opaque({value[start:end]})"
    if value == "fixed_size_list<item: float>[3]":
        return "fixed_list<float32>[3]"
    return value


def metadata_for_driver(driver: str, args: argparse.Namespace) -> dict[str, Any]:
    with dbapi.connect(driver=driver, uri=uri(driver, args), autocommit=True) as connection:
        objects = connection.adbc_get_objects(
            depth="columns",
            catalog_filter=args.database,
            table_name_filter=METADATA_TABLE,
        ).read_all().to_pylist()
        columns = find_columns(objects)
        constraints = find_constraints(objects)
        schema = connection.adbc_get_table_schema(
            METADATA_TABLE, catalog_filter=args.database
        )
        json_field = schema.field("document")
        return {
            "invisible remarks": columns["hidden_note"].get("remarks"),
            "generated flag": columns["doubled"].get("xdbc_is_generatedcolumn"),
            "JSON Arrow type": str(json_field.type),
            "JSON metadata": decoded_metadata(json_field),
            "native schema types": {
                name: compact_type(schema.field(name).type)
                for name in ("native_uuid", "embedding", "address")
            },
            "ENUM metadata": decoded_metadata(schema.field("status")),
            "SET metadata": decoded_metadata(schema.field("flags")),
            "constraint types": [item["constraint_type"] for item in constraints],
        }


def run_metadata_comparison(args: argparse.Namespace) -> None:
    """Compare discovery APIs using a short-lived regular table.

    INFORMATION_SCHEMA does not list temporary tables, so this part must use a
    regular table. Its process-ID-suffixed name is removed in a finally block.
    """
    setup_uri = uri("mariadb", args)
    try:
        with dbapi.connect(driver="mariadb", uri=setup_uri, autocommit=True) as connection:
            with connection.cursor() as cursor:
                cursor.execute(f"DROP TABLE IF EXISTS `{METADATA_TABLE}`")
                cursor.execute(CREATE_METADATA_TABLE.format(table=f"`{METADATA_TABLE}`"))

        print("\n=== ADBC metadata discovery ===")
        for driver in ("mysql", "mariadb"):
            print(f"{driver:7}:")
            try:
                metadata = metadata_for_driver(driver, args)
                print(f"  invisible remarks: {metadata['invisible remarks']!r}")
                print(f"  generated flag:    {metadata['generated flag']!r}")
                print(f"  JSON Arrow type:   {metadata['JSON Arrow type']}")
                print("  JSON metadata:")
                print_metadata(metadata["JSON metadata"], indent="    ")
                print("  Native schema types:")
                print_metadata(metadata["native schema types"], indent="    ")
                print("  ENUM/SET allowed values:")
                enum_values = metadata["ENUM metadata"].get("mariadb.enum_values")
                set_values = metadata["SET metadata"].get("mariadb.set_values")
                print(f"    ENUM = {enum_values or '(not reported)'}")
                print(f"    SET  = {set_values or '(not reported)'}")
                constraints = metadata["constraint types"]
                print(f"  Constraints: {', '.join(constraints) if constraints else '(not reported)'}")
            except Exception as exc:
                print(f"  ERROR: {str(exc).replace(chr(10), ' ')}")
    finally:
        with dbapi.connect(driver="mariadb", uri=setup_uri, autocommit=True) as connection:
            with connection.cursor() as cursor:
                cursor.execute(f"DROP TABLE IF EXISTS `{METADATA_TABLE}`")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default=os.getenv("MARIADB_HOST", "127.0.0.1"))
    parser.add_argument("--port", type=int, default=int(os.getenv("MARIADB_PORT", "13100")))
    parser.add_argument("--user", default=os.getenv("MARIADB_USER", "msandbox"))
    parser.add_argument("--password", default=os.getenv("MARIADB_PASSWORD", "msandbox"))
    parser.add_argument("--database", default=os.getenv("MARIADB_DATABASE", "test"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    print(f"MariaDB server: {args.host}:{args.port}/{args.database}")
    print("The password is intentionally not printed.")
    results = {driver: run_driver(driver, args) for driver in ("mysql", "mariadb")}
    print_results(results)
    run_metadata_comparison(args)


if __name__ == "__main__":
    main()