Examples with Portal io.crate.protocols.postgres.Portal used on opensource projects

Example 1 with Portal

use of io.crate.protocols.postgres.Portal in project crate by crate.

the class Session method singleExec.

@VisibleForTesting
CompletableFuture<?> singleExec(Portal portal, ResultReceiver<?> resultReceiver, int maxRows) {
    var activeConsumer = portal.activeConsumer();
    if (activeConsumer != null && activeConsumer.suspended()) {
        activeConsumer.replaceResultReceiver(resultReceiver, maxRows);
        activeConsumer.resume();
        return resultReceiver.completionFuture();
    }
    var jobId = UUIDs.dirtyUUID();
    var routingProvider = new RoutingProvider(Randomness.get().nextInt(), planner.getAwarenessAttributes());
    var clusterState = executor.clusterService().state();
    var txnCtx = new CoordinatorTxnCtx(sessionContext);
    var nodeCtx = executor.nodeContext();
    var params = new RowN(portal.params().toArray());
    var plannerContext = new PlannerContext(clusterState, routingProvider, jobId, txnCtx, nodeCtx, maxRows, params);
    var analyzedStmt = portal.analyzedStatement();
    String rawStatement = portal.preparedStmt().rawStatement();
    if (analyzedStmt == null) {
        String errorMsg = "Statement must have been analyzed: " + rawStatement;
        jobsLogs.logPreExecutionFailure(jobId, rawStatement, errorMsg, sessionContext.sessionUser());
        throw new IllegalStateException(errorMsg);
    }
    Plan plan;
    try {
        plan = planner.plan(analyzedStmt, plannerContext);
    } catch (Throwable t) {
        jobsLogs.logPreExecutionFailure(jobId, rawStatement, SQLExceptions.messageOf(t), sessionContext.sessionUser());
        throw t;
    }
    if (!analyzedStmt.isWriteOperation()) {
        resultReceiver = new RetryOnFailureResultReceiver(executor.clusterService(), clusterState, indexName -> executor.clusterService().state().metadata().hasIndex(indexName), resultReceiver, jobId, (newJobId, resultRec) -> retryQuery(newJobId, analyzedStmt, routingProvider, new RowConsumerToResultReceiver(resultRec, maxRows, new JobsLogsUpdateListener(newJobId, jobsLogs)), params, txnCtx, nodeCtx));
    }
    jobsLogs.logExecutionStart(jobId, rawStatement, sessionContext.sessionUser(), StatementClassifier.classify(plan));
    RowConsumerToResultReceiver consumer = new RowConsumerToResultReceiver(resultReceiver, maxRows, new JobsLogsUpdateListener(jobId, jobsLogs));
    portal.setActiveConsumer(consumer);
    plan.execute(executor, plannerContext, consumer, params, SubQueryResults.EMPTY);
    return resultReceiver.completionFuture();
}

Example 2 with Portal

use of io.crate.protocols.postgres.Portal in project crate by crate.

the class Session method execute.

@Nullable
public CompletableFuture<?> execute(String portalName, int maxRows, ResultReceiver<?> resultReceiver) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("method=execute portalName={} maxRows={}", portalName, maxRows);
    }
    Portal portal = getSafePortal(portalName);
    var analyzedStmt = portal.analyzedStatement();
    if (isReadOnly && analyzedStmt.isWriteOperation()) {
        throw new ReadOnlyException(portal.preparedStmt().rawStatement());
    }
    if (analyzedStmt instanceof AnalyzedBegin) {
        currentTransactionState = TransactionState.IN_TRANSACTION;
        resultReceiver.allFinished(false);
    } else if (analyzedStmt instanceof AnalyzedCommit) {
        currentTransactionState = TransactionState.IDLE;
        resultReceiver.allFinished(false);
        return resultReceiver.completionFuture();
    } else if (analyzedStmt instanceof AnalyzedDeallocate) {
        String stmtToDeallocate = ((AnalyzedDeallocate) analyzedStmt).preparedStmtName();
        if (stmtToDeallocate != null) {
            close((byte) 'S', stmtToDeallocate);
        } else {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("deallocating all prepared statements");
            }
            preparedStatements.clear();
        }
        resultReceiver.allFinished(false);
    } else if (analyzedStmt instanceof AnalyzedDiscard) {
        AnalyzedDiscard discard = (AnalyzedDiscard) analyzedStmt;
        // See https://www.postgresql.org/docs/current/sql-discard.html
        if (discard.target() == Target.ALL) {
            close();
        }
        resultReceiver.allFinished(false);
    } else if (analyzedStmt.isWriteOperation()) {
        /* We defer the execution for any other statements to `sync` messages so that we can efficiently process
             * bulk operations. E.g. If we receive `INSERT INTO (x) VALUES (?)` bindings/execute multiple times
             * We want to create bulk requests internally:                                                          /
             * - To reduce network overhead
             * - To have 1 disk flush per shard instead of 1 disk flush per item
             *
             * Many clients support this by doing something like this:
             *
             *      var preparedStatement = conn.prepareStatement("...")
             *      for (var args in manyArgs):
             *          preparedStatement.execute(args)
             *      conn.commit()
             */
        deferredExecutionsByStmt.compute(portal.preparedStmt().parsedStatement(), (key, oldValue) -> {
            DeferredExecution deferredExecution = new DeferredExecution(portal, maxRows, resultReceiver);
            if (oldValue == null) {
                ArrayList<DeferredExecution> deferredExecutions = new ArrayList<>();
                deferredExecutions.add(deferredExecution);
                return deferredExecutions;
            } else {
                oldValue.add(deferredExecution);
                return oldValue;
            }
        });
        return resultReceiver.completionFuture();
    } else {
        if (!deferredExecutionsByStmt.isEmpty()) {
            throw new UnsupportedOperationException("Only write operations are allowed in Batch statements");
        }
        if (activeExecution == null) {
            activeExecution = singleExec(portal, resultReceiver, maxRows);
        } else {
            activeExecution = activeExecution.thenCompose(ignored -> singleExec(portal, resultReceiver, maxRows));
        }
        return activeExecution;
    }
    return null;
}

Example 3 with Portal

use of io.crate.protocols.postgres.Portal in project crate by crate.

the class Session method getOutputTypes.

@Nullable
public List<? extends DataType> getOutputTypes(String portalName) {
    Portal portal = getSafePortal(portalName);
    var analyzedStatement = portal.analyzedStatement();
    List<Symbol> fields = analyzedStatement.outputs();
    if (fields != null) {
        return Symbols.typeView(fields);
    }
    return null;
}

Example 4 with Portal

use of io.crate.protocols.postgres.Portal in project crate by crate.

the class Session method bind.

public void bind(String portalName, String statementName, List<Object> params, @Nullable FormatCodes.FormatCode[] resultFormatCodes) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("method=bind portalName={} statementName={} params={}", portalName, statementName, params);
    }
    PreparedStmt preparedStmt;
    try {
        preparedStmt = getSafeStmt(statementName);
    } catch (Throwable t) {
        jobsLogs.logPreExecutionFailure(UUIDs.dirtyUUID(), null, SQLExceptions.messageOf(t), sessionContext.sessionUser());
        throw t;
    }
    Portal portal = new Portal(portalName, preparedStmt, params, preparedStmt.analyzedStatement(), resultFormatCodes);
    Portal oldPortal = portals.put(portalName, portal);
    if (oldPortal != null) {
        // According to the wire protocol spec named portals should be removed explicitly and only
        // unnamed portals are implicitly closed/overridden.
        // We don't comply with the spec because we allow batching of statements, see #execute
        oldPortal.closeActiveConsumer();
    }
}

Example 5 with Portal

use of io.crate.protocols.postgres.Portal in project crate by crate.

the class Session method describe.

public DescribeResult describe(char type, String portalOrStatement) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("method=describe type={} portalOrStatement={}", type, portalOrStatement);
    }
    switch(type) {
        case 'P':
            Portal portal = getSafePortal(portalOrStatement);
            var analyzedStmt = portal.analyzedStatement();
            return new DescribeResult(portal.preparedStmt().parameterTypes(), analyzedStmt.outputs(), resolveTableFromSelect(analyzedStmt));
        case 'S':
            /*
                 * describe might be called without prior bind call.
                 *
                 * If the client uses server-side prepared statements this is usually the case.
                 *
                 * E.g. the statement is first prepared:
                 *
                 *      parse stmtName=S_1 query=insert into t (x) values ($1) paramTypes=[integer]
                 *      describe type=S portalOrStatement=S_1
                 *      sync
                 *
                 * and then used with different bind calls:
                 *
                 *      bind portalName= statementName=S_1 params=[0]
                 *      describe type=P portalOrStatement=
                 *      execute
                 *
                 *      bind portalName= statementName=S_1 params=[1]
                 *      describe type=P portalOrStatement=
                 *      execute
                 */
            PreparedStmt preparedStmt = preparedStatements.get(portalOrStatement);
            AnalyzedStatement analyzedStatement = preparedStmt.analyzedStatement();
            return new DescribeResult(preparedStmt.parameterTypes(), analyzedStatement.outputs(), resolveTableFromSelect(analyzedStatement));
        default:
            throw new AssertionError("Unsupported type: " + type);
    }
}