Examples with TransactionSelector com.google.spanner.v1.TransactionSelector used on opensource projects

Example 1 with TransactionSelector

use of com.google.spanner.v1.TransactionSelector in project java-spanner by googleapis.

the class MockSpannerServiceImpl method partition.

private void partition(String sessionName, TransactionSelector transactionSelector, StreamObserver<PartitionResponse> responseObserver) {
    Session session = sessions.get(sessionName);
    if (session == null) {
        setSessionNotFound(sessionName, responseObserver);
        return;
    }
    sessionLastUsed.put(session.getName(), Instant.now());
    try {
        ByteString transactionId = getTransactionId(session, transactionSelector);
        responseObserver.onNext(PartitionResponse.newBuilder().addPartitions(Partition.newBuilder().setPartitionToken(generatePartitionToken(session.getName(), transactionId)).build()).build());
        responseObserver.onCompleted();
    } catch (StatusRuntimeException e) {
        responseObserver.onError(e);
    } catch (Throwable t) {
        responseObserver.onError(Status.INTERNAL.asRuntimeException());
    }
}

Example 2 with TransactionSelector

use of com.google.spanner.v1.TransactionSelector in project java-spanner by googleapis.

the class MockSpannerServiceImpl method returnResultSet.

private void returnResultSet(ResultSet resultSet, ByteString transactionId, TransactionSelector transactionSelector, StreamObserver<ResultSet> responseObserver) {
    ResultSetMetadata metadata = resultSet.getMetadata();
    if (transactionId != null) {
        metadata = metadata.toBuilder().setTransaction(ignoreNextInlineBeginRequest.getAndSet(false) ? Transaction.getDefaultInstance() : Transaction.newBuilder().setId(transactionId).build()).build();
    } else if (transactionSelector.hasBegin() || transactionSelector.hasSingleUse()) {
        Transaction transaction = getTemporaryTransactionOrNull(transactionSelector);
        metadata = metadata.toBuilder().setTransaction(transaction).build();
    }
    resultSet = resultSet.toBuilder().setMetadata(metadata).build();
    responseObserver.onNext(resultSet);
}

Example 3 with TransactionSelector

use of com.google.spanner.v1.TransactionSelector in project java-spanner by googleapis.

the class AbstractReadContext method readInternalWithOptions.

ResultSet readInternalWithOptions(String table, @Nullable String index, KeySet keys, Iterable<String> columns, final Options readOptions, ByteString partitionToken) {
    beforeReadOrQuery();
    final ReadRequest.Builder builder = ReadRequest.newBuilder().setSession(session.getName()).setTable(checkNotNull(table)).addAllColumns(columns);
    if (readOptions.hasLimit()) {
        builder.setLimit(readOptions.limit());
    }
    keys.appendToProto(builder.getKeySetBuilder());
    if (index != null) {
        builder.setIndex(index);
    }
    if (partitionToken != null) {
        builder.setPartitionToken(partitionToken);
    }
    final int prefetchChunks = readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
    ResumableStreamIterator stream = new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.READ, span) {

        @Override
        CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
            GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
            TransactionSelector selector = null;
            if (resumeToken != null) {
                builder.setResumeToken(resumeToken);
                selector = getTransactionSelector();
            } else if (!builder.hasTransaction()) {
                selector = getTransactionSelector();
            }
            if (selector != null) {
                builder.setTransaction(selector);
            }
            builder.setRequestOptions(buildRequestOptions(readOptions));
            SpannerRpc.StreamingCall call = rpc.read(builder.build(), stream.consumer(), session.getOptions());
            call.request(prefetchChunks);
            stream.setCall(call, /* withBeginTransaction = */
            builder.getTransaction().hasBegin());
            return stream;
        }
    };
    return new GrpcResultSet(stream, this);
}

Example 4 with TransactionSelector

use of com.google.spanner.v1.TransactionSelector in project java-spanner by googleapis.

the class AbstractReadContext method executeQueryInternalWithOptions.

ResultSet executeQueryInternalWithOptions(final Statement statement, final com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode, final Options options, final ByteString partitionToken) {
    beforeReadOrQuery();
    final int prefetchChunks = options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks;
    final ExecuteSqlRequest.Builder request = getExecuteSqlRequestBuilder(statement, queryMode, options, /* withTransactionSelector = */
    false);
    ResumableStreamIterator stream = new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {

        @Override
        CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
            GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
            if (partitionToken != null) {
                request.setPartitionToken(partitionToken);
            }
            TransactionSelector selector = null;
            if (resumeToken != null) {
                request.setResumeToken(resumeToken);
                selector = getTransactionSelector();
            } else if (!request.hasTransaction()) {
                selector = getTransactionSelector();
            }
            if (selector != null) {
                request.setTransaction(selector);
            }
            SpannerRpc.StreamingCall call = rpc.executeQuery(request.build(), stream.consumer(), session.getOptions());
            call.request(prefetchChunks);
            stream.setCall(call, request.getTransaction().hasBegin());
            return stream;
        }
    };
    return new GrpcResultSet(stream, this);
}

Example 5 with TransactionSelector

use of com.google.spanner.v1.TransactionSelector in project java-spanner by googleapis.

the class AbstractReadContext method getExecuteBatchDmlRequestBuilder.

ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(Iterable<Statement> statements, Options options) {
    ExecuteBatchDmlRequest.Builder builder = ExecuteBatchDmlRequest.newBuilder().setSession(session.getName());
    int idx = 0;
    for (Statement stmt : statements) {
        builder.addStatementsBuilder();
        builder.getStatementsBuilder(idx).setSql(stmt.getSql());
        Map<String, Value> stmtParameters = stmt.getParameters();
        if (!stmtParameters.isEmpty()) {
            com.google.protobuf.Struct.Builder paramsBuilder = builder.getStatementsBuilder(idx).getParamsBuilder();
            for (Map.Entry<String, Value> param : stmtParameters.entrySet()) {
                paramsBuilder.putFields(param.getKey(), Value.toProto(param.getValue()));
                if (param.getValue() != null) {
                    builder.getStatementsBuilder(idx).putParamTypes(param.getKey(), param.getValue().getType().toProto());
                }
            }
        }
        idx++;
    }
    TransactionSelector selector = getTransactionSelector();
    if (selector != null) {
        builder.setTransaction(selector);
    }
    builder.setSeqno(getSeqNo());
    builder.setRequestOptions(buildRequestOptions(options));
    return builder;
}