通过 Unified Query 使用 Kyuubi 执行 SparkSQL

package com.example;

import com.portofino.unifiedquery.submission.KyuubiSqlExecutionResultJava;
import com.portofino.unifiedquery.submission.KyuubiSqlOperationRequestException;
import com.portofino.unifiedquery.submission.KyuubiSqlOperationHandleJava;
import com.portofino.unifiedquery.submission.KyuubiSqlSessionClientJava;
import com.portofino.unifiedquery.submission.KyuubiSqlSessionHandleJava;
import java.util.Map;

public final class KyuubiSqlSessionExample {
    private static String requiredEnv(String name) {
        String value = System.getenv(name);
        if (value == null || value.isBlank()) {
            throw new IllegalArgumentException("Set environment variable " + name);
        }
        return value;
    }

    private static KyuubiSqlSessionClientJava newClient() {
        return new KyuubiSqlSessionClientJava(
            requiredEnv("KYUUBI_URI"),
            System.getenv().getOrDefault("KYUUBI_USERNAME", ""),
            System.getenv().getOrDefault("KYUUBI_PASSWORD", ""),
            30);
    }

    private static int waitSeconds() {
        return Integer.parseInt(System.getenv().getOrDefault("UQ_SQL_WAIT_SECONDS", "300"));
    }

    private static void printHandle(
            KyuubiSqlSessionHandleJava session,
            KyuubiSqlOperationHandleJava operation) {
        System.out.printf(
            "PENDING session_id=%s kyuubi_instance=%s operation_id=%s%n",
            session.getSessionId(),
            session.getKyuubiInstance(),
            operation.getOperationId());
    }

    private static void printOutcome(KyuubiSqlExecutionResultJava outcome) {
        if ("COMPLETED".equals(outcome.getKind())) {
            outcome.getColumns().forEach(column ->
                System.out.println(column.getName() + " " + column.getDataType()));
            System.out.println(outcome.getRows());
        } else {
            System.err.println(outcome.getStatus().getDiagnostic());
        }
    }

    private static void submit(KyuubiSqlSessionClientJava client) {
        var session = client.openSession(Map.of(
            "kyuubi.session.engine.type", "SPARK_SQL"));
        boolean keepSessionOpen = false;
        try {
            KyuubiSqlExecutionResultJava outcome;
            try {
                String sql = System.getenv().getOrDefault(
                    "UQ_EXAMPLE_SQL", "SELECT 1 AS value");
                outcome = client.executeSql(session, sql, waitSeconds());
            } catch (KyuubiSqlOperationRequestException error) {
                keepSessionOpen = true;
                var accepted = error.getOperationHandle();
                System.err.printf(
                    "PENDING session_id=%s kyuubi_instance=%s operation_id=%s%n",
                    accepted.session().sessionId(),
                    accepted.session().kyuubiInstance(),
                    accepted.operationId());
                throw error;
            } catch (RuntimeException error) {
                keepSessionOpen = true;
                System.err.printf(
                    "No handle was confirmed for session_id=%s on kyuubi_instance=%s; check Kyuubi before resubmitting.%n",
                    session.getSessionId(), session.getKyuubiInstance());
                throw error;
            }

            if ("PENDING".equals(outcome.getKind())) {
                keepSessionOpen = true;
                printHandle(session, outcome.getHandle());
                return;
            }

            printOutcome(outcome);
            client.closeOperation(outcome.getHandle());
        } finally {
            if (!keepSessionOpen) {
                client.closeSession(session);
            }
        }
    }

    private static void resume(KyuubiSqlSessionClientJava client) {
        var session = new KyuubiSqlSessionHandleJava(
            requiredEnv("UQ_KYUUBI_SESSION_ID"),
            requiredEnv("UQ_KYUUBI_INSTANCE"));
        var handle = new KyuubiSqlOperationHandleJava(
            session.getSessionId(),
            session.getKyuubiInstance(),
            requiredEnv("UQ_KYUUBI_OPERATION_ID"));

        try {
            System.out.println(client.getOperationStatus(handle).getState());
            var outcome = client.getResult(handle, waitSeconds());
            if ("PENDING".equals(outcome.getKind())) {
                printHandle(session, handle);
                return;
            }
            printOutcome(outcome);
            client.closeOperation(handle);
            client.closeSession(session);
        } catch (RuntimeException error) {
            printHandle(session, handle);
            throw error;
        }
    }

    public static void main(String[] args) {
        var client = newClient();
        if (args.length == 1 && "--resume".equals(args[0])) {
            resume(client);
        } else if (args.length == 0) {
            submit(client);
        } else {
            throw new IllegalArgumentException(
                "Usage: KyuubiSqlSessionExample [--resume]");
        }
    }
}

import argparse
import os
import sys

from unified_query_api import (
    KyuubiRequestError,
    KyuubiSqlOperationHandle,
    KyuubiSqlOperationRequestError,
    KyuubiSqlSessionClient,
    KyuubiSqlSessionHandle,
)


def required_env(name):
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"Set environment variable {name}")
    return value


def new_client():
    return KyuubiSqlSessionClient(
        required_env("KYUUBI_URI"),
        username=os.environ.get("KYUUBI_USERNAME", ""),
        password=os.environ.get("KYUUBI_PASSWORD", ""),
        request_timeout_seconds=30,
    )


def wait_seconds():
    return int(os.environ.get("UQ_SQL_WAIT_SECONDS", "300"))


def print_handle(handle):
    print(
        "PENDING",
        f"session_id={handle.session.session_id}",
        f"kyuubi_instance={handle.session.kyuubi_instance}",
        f"operation_id={handle.operation_id}",
        flush=True,
    )


def print_outcome(outcome):
    if outcome.kind == "COMPLETED":
        print([(column.name, column.data_type) for column in outcome.columns])
        print(outcome.rows)
    else:
        print(outcome.status.diagnostic or outcome.status.state)


def submit(client):
    session = client.open_session({"kyuubi.session.engine.type": "SPARK_SQL"})
    keep_session_open = False
    try:
        try:
            outcome = client.execute_sql(
                session,
                os.environ.get("UQ_EXAMPLE_SQL", "SELECT 1 AS value"),
                sync_wait_timeout_seconds=wait_seconds(),
            )
        except KyuubiSqlOperationRequestError as error:
            keep_session_open = True
            print_handle(error.handle)
            raise
        except KyuubiRequestError:
            keep_session_open = True
            print(
                f"No handle was confirmed for session_id={session.session_id} and kyuubi_instance={session.kyuubi_instance}; check Kyuubi before resubmitting.",
                flush=True,
            )
            raise

        if outcome.kind == "PENDING":
            keep_session_open = True
            print_handle(outcome.handle)
            return

        print_outcome(outcome)
        client.close_operation(outcome.handle)
    finally:
        if not keep_session_open:
            client.close_session(session)


def resume(client):
    session = KyuubiSqlSessionHandle(
        session_id=required_env("UQ_KYUUBI_SESSION_ID"),
        kyuubi_instance=required_env("UQ_KYUUBI_INSTANCE"),
    )
    handle = KyuubiSqlOperationHandle(
        session=session,
        operation_id=required_env("UQ_KYUUBI_OPERATION_ID"),
    )

    try:
        print(client.get_operation_status(handle).state)
        outcome = client.get_result(handle, sync_wait_timeout_seconds=wait_seconds())
    except KyuubiRequestError:
        print_handle(handle)
        raise

    if outcome.kind == "PENDING":
        print_handle(handle)
        return

    print_outcome(outcome)
    client.close_operation(handle)
    client.close_session(session)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()

    client = new_client()
    if args.resume:
        resume(client)
    else:
        submit(client)


if __name__ == "__main__":
    main()

package com.example;

import com.portofino.unifiedquery.submission.KyuubiSqlExecutionResultJava;
import com.portofino.unifiedquery.submission.KyuubiSqlOperationRequestException;
import com.portofino.unifiedquery.submission.KyuubiSqlOperationHandleJava;
import com.portofino.unifiedquery.submission.KyuubiSqlSessionClientJava;
import com.portofino.unifiedquery.submission.KyuubiSqlSessionHandleJava;
import java.util.Map;
import java.util.UUID;

public final class KyuubiSequentialSqlExample {
    private static boolean retainSessionForRecovery;

    private static String requiredEnv(String name) {
        String value = System.getenv(name);
        if (value == null || value.isBlank()) {
            throw new IllegalArgumentException("Set environment variable " + name);
        }
        return value;
    }

    private static KyuubiSqlSessionClientJava newClient() {
        return new KyuubiSqlSessionClientJava(
            requiredEnv("KYUUBI_URI"),
            System.getenv().getOrDefault("KYUUBI_USERNAME", ""),
            System.getenv().getOrDefault("KYUUBI_PASSWORD", ""),
            30);
    }

    private static void printHandle(KyuubiSqlOperationHandleJava handle) {
        var session = handle.getSession();
        System.err.printf(
            "PENDING session_id=%s kyuubi_instance=%s operation_id=%s%n",
            session.getSessionId(), session.getKyuubiInstance(), handle.getOperationId());
    }

    private static KyuubiSqlExecutionResultJava executeToTerminal(
            KyuubiSqlSessionClientJava client,
            KyuubiSqlSessionHandleJava session,
            String sql) {
        KyuubiSqlExecutionResultJava outcome;
        try {
            outcome = client.executeSql(session, sql, 300);
        } catch (KyuubiSqlOperationRequestException error) {
            retainSessionForRecovery = true;
            var accepted = error.getOperationHandle();
            System.err.printf(
                "PENDING session_id=%s kyuubi_instance=%s operation_id=%s%n",
                accepted.session().sessionId(),
                accepted.session().kyuubiInstance(),
                accepted.operationId());
            throw error;
        } catch (RuntimeException error) {
            retainSessionForRecovery = true;
            System.err.printf(
                "No handle was confirmed for session_id=%s on kyuubi_instance=%s; check Kyuubi before resubmitting.%n",
                session.getSessionId(), session.getKyuubiInstance());
            throw error;
        }

        while ("PENDING".equals(outcome.getKind())) {
            var handle = outcome.getHandle();
            try {
                System.out.println(client.getOperationStatus(handle).getState());
                outcome = client.getResult(handle, 300);
            } catch (RuntimeException error) {
                retainSessionForRecovery = true;
                printHandle(handle);
                throw error;
            }
        }

        if ("FAILED".equals(outcome.getKind())) {
            client.closeOperation(outcome.getHandle());
            return outcome;
        }

        try {
            client.closeOperation(outcome.getHandle());
        } catch (RuntimeException error) {
            retainSessionForRecovery = true;
            printHandle(outcome.getHandle());
            throw error;
        }
        return outcome;
    }

    public static void main(String[] args) {
        var client = newClient();
        var session = client.openSession(Map.of(
            "kyuubi.session.engine.type", "SPARK_SQL"));
        String table = "uq_example_" + UUID.randomUUID().toString().replace("-", "");
        boolean tableMayExist = false;

        try {
            var outcome = executeToTerminal(
                client, session, "DROP TABLE IF EXISTS " + table);
            if ("FAILED".equals(outcome.getKind())) {
                System.err.println(outcome.getStatus().getDiagnostic());
                return;
            }

            tableMayExist = true;
            outcome = executeToTerminal(client, session,
                "CREATE TABLE " + table + " (id INT, label STRING) USING PARQUET");
            if ("FAILED".equals(outcome.getKind())) {
                System.err.println(outcome.getStatus().getDiagnostic());
                return;
            }

            outcome = executeToTerminal(client, session,
                "INSERT INTO " + table + " VALUES (1, 'first'), (2, 'second')");
            if ("FAILED".equals(outcome.getKind())) {
                System.err.println(outcome.getStatus().getDiagnostic());
                return;
            }

            outcome = executeToTerminal(client, session,
                "SELECT id, label FROM " + table + " ORDER BY id");
            if ("FAILED".equals(outcome.getKind())) {
                System.err.println(outcome.getStatus().getDiagnostic());
                return;
            }
            System.out.println(outcome.getRows());

            outcome = executeToTerminal(client, session, "DROP TABLE IF EXISTS " + table);
            if ("FAILED".equals(outcome.getKind())) {
                System.err.println(outcome.getStatus().getDiagnostic());
                return;
            }
            tableMayExist = false;
        } finally {
            if (!retainSessionForRecovery) {
                if (tableMayExist) {
                    var cleanup = executeToTerminal(
                        client, session, "DROP TABLE IF EXISTS " + table);
                    if ("FAILED".equals(cleanup.getKind())) {
                        System.err.println(cleanup.getStatus().getDiagnostic());
                    }
                }
                if (!retainSessionForRecovery) {
                    client.closeSession(session);
                }
            }
        }
    }
}

import os
import sys
from uuid import uuid4

from unified_query_api import (
    KyuubiRequestError,
    KyuubiSqlOperationRequestError,
    KyuubiSqlSessionClient,
)


def required_env(name):
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"Set environment variable {name}")
    return value


def wait_seconds():
    return int(os.environ.get("UQ_SQL_WAIT_SECONDS", "300"))


def print_handle(handle):
    print(
        "PENDING",
        f"session_id={handle.session.session_id}",
        f"kyuubi_instance={handle.session.kyuubi_instance}",
        f"operation_id={handle.operation_id}",
        file=sys.stderr,
        flush=True,
    )


def main():
    client = KyuubiSqlSessionClient(
        required_env("KYUUBI_URI"),
        username=os.environ.get("KYUUBI_USERNAME", ""),
        password=os.environ.get("KYUUBI_PASSWORD", ""),
        request_timeout_seconds=30,
    )
    session = client.open_session({"kyuubi.session.engine.type": "SPARK_SQL"})
    table = "uq_example_" + uuid4().hex
    retain_session_for_recovery = False
    table_may_exist = False

    def execute_to_terminal(sql):
        nonlocal retain_session_for_recovery
        try:
            outcome = client.execute_sql(
                session, sql, sync_wait_timeout_seconds=wait_seconds()
            )
        except KyuubiSqlOperationRequestError as error:
            retain_session_for_recovery = True
            print_handle(error.handle)
            raise
        except KyuubiRequestError:
            retain_session_for_recovery = True
            print(
                f"No handle was confirmed for session_id={session.session_id} and kyuubi_instance={session.kyuubi_instance}; check Kyuubi before resubmitting.",
                file=sys.stderr,
            )
            raise

        while outcome.kind == "PENDING":
            handle = outcome.handle
            try:
                print(client.get_operation_status(handle).state)
                outcome = client.get_result(
                    handle, sync_wait_timeout_seconds=wait_seconds()
                )
            except KyuubiRequestError:
                retain_session_for_recovery = True
                print_handle(handle)
                raise

        if outcome.kind == "FAILED":
            client.close_operation(outcome.handle)
            return outcome

        try:
            client.close_operation(outcome.handle)
        except KyuubiRequestError:
            retain_session_for_recovery = True
            print_handle(outcome.handle)
            raise
        return outcome

    try:
        outcome = execute_to_terminal(f"DROP TABLE IF EXISTS {table}")
        if outcome.kind == "FAILED":
            print(outcome.status.diagnostic or outcome.status.state)
            return 1

        table_may_exist = True
        outcome = execute_to_terminal(
            f"CREATE TABLE {table} (id INT, label STRING) USING PARQUET"
        )
        if outcome.kind == "FAILED":
            print(outcome.status.diagnostic or outcome.status.state)
            return 1

        outcome = execute_to_terminal(
            f"INSERT INTO {table} VALUES (1, 'first'), (2, 'second')"
        )
        if outcome.kind == "FAILED":
            print(outcome.status.diagnostic or outcome.status.state)
            return 1

        outcome = execute_to_terminal(
            f"SELECT id, label FROM {table} ORDER BY id"
        )
        if outcome.kind == "FAILED":
            print(outcome.status.diagnostic or outcome.status.state)
            return 1
        print(outcome.rows)

        outcome = execute_to_terminal(f"DROP TABLE IF EXISTS {table}")
        if outcome.kind == "FAILED":
            print(outcome.status.diagnostic or outcome.status.state)
            return 1
        table_may_exist = False
        return 0
    finally:
        if not retain_session_for_recovery:
            if table_may_exist:
                cleanup = execute_to_terminal(f"DROP TABLE IF EXISTS {table}")
                if cleanup.kind == "FAILED":
                    print(cleanup.status.diagnostic or cleanup.status.state)
            if not retain_session_for_recovery:
                client.close_session(session)


if __name__ == "__main__":
    raise SystemExit(main())