tech.ydb.jdbc.context.BaseYdbExecutor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ydb-jdbc-driver Show documentation
JDBC Driver over YDB Java SDK
There is a newer version: 2.3.5
package tech.ydb.jdbc.context;

import java.sql.SQLException;
import java.time.Duration;
import java.util.Collection;
import java.util.concurrent.LinkedBlockingQueue;

import tech.ydb.core.Result;
import tech.ydb.core.UnexpectedResultException;
import tech.ydb.jdbc.exception.ExceptionFactory;
import tech.ydb.jdbc.query.QueryType;
import tech.ydb.jdbc.query.YdbQuery;
import tech.ydb.table.Session;
import tech.ydb.table.TableClient;
import tech.ydb.table.query.ExplainDataQueryResult;
import tech.ydb.table.query.Params;
import tech.ydb.table.result.ResultSetReader;
import tech.ydb.table.result.impl.ProtoValueReaders;
import tech.ydb.table.settings.ExecuteScanQuerySettings;
import tech.ydb.table.settings.ExecuteSchemeQuerySettings;
import tech.ydb.table.settings.ExplainDataQuerySettings;

/**
 *
 * @author Aleksandr Gorshenin
 */
public abstract class BaseYdbExecutor implements YdbExecutor {
    private final Duration sessionTimeout;
    private final TableClient tableClient;

    public BaseYdbExecutor(YdbContext ctx) {
        this.sessionTimeout = ctx.getOperationProperties().getSessionTimeout();
        this.tableClient = ctx.getTableClient();
    }

    protected Session createNewTableSession(YdbValidator validator) throws SQLException {
        try {
            Result session = tableClient.createSession(sessionTimeout).join();
            validator.addStatusIssues(session.getStatus());
            return session.getValue();
        } catch (UnexpectedResultException ex) {
            throw ExceptionFactory.createException("Cannot create session with " + ex.getStatus(), ex);
        }
    }

    @Override
    public void executeSchemeQuery(YdbContext ctx, YdbValidator validator, YdbQuery query) throws SQLException {
        // Scheme query does not affect transactions or result sets
        ExecuteSchemeQuerySettings settings = ctx.withDefaultTimeout(new ExecuteSchemeQuerySettings());
        final String yql = query.getYqlQuery(null);

        try (Session session = createNewTableSession(validator)) {
            validator.execute(QueryType.SCHEME_QUERY + " >>\n" + yql, () -> session.executeSchemeQuery(yql, settings));
        }
    }

    @Override
    public ExplainDataQueryResult executeExplainQuery(YdbContext ctx, YdbValidator validator, YdbQuery query)
            throws SQLException {
        ensureOpened();

        String yql = query.getYqlQuery(null);
        ExplainDataQuerySettings settings = ctx.withDefaultTimeout(new ExplainDataQuerySettings());
        try (Session session = createNewTableSession(validator)) {
            String msg = QueryType.EXPLAIN_QUERY + " >>\n" + yql;
            return validator.call(msg, () -> session.explainDataQuery(yql, settings));
        }
    }

    @Override
    public ResultSetReader executeScanQuery(YdbContext ctx, YdbValidator validator, YdbQuery query, Params params)
            throws SQLException {
        ensureOpened();

        String yql = query.getYqlQuery(params);
        Collection resultSets = new LinkedBlockingQueue<>();
        Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
        ExecuteScanQuerySettings settings = ExecuteScanQuerySettings.newBuilder()
                .withRequestTimeout(scanQueryTimeout)
                .build();
        try (Session session = createNewTableSession(validator)) {
            validator.execute(QueryType.SCAN_QUERY + " >>\n" + yql,
                    () -> session.executeScanQuery(yql, params, settings).start(resultSets::add));
        }

        return ProtoValueReaders.forResultSets(resultSets);
    }


}