com.couchbase.client.kotlin.query.internal.QueryExecutor.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of kotlin-client Show documentation
The official Couchbase Kotlin SDK
There is a newer version: 1.4.7
/*
 * Copyright 2021 Couchbase, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.couchbase.client.kotlin.query.internal

import com.couchbase.client.core.api.query.CoreQueryContext
import com.couchbase.client.core.api.query.CoreQueryOps
import com.couchbase.client.core.api.query.CoreQueryOptions
import com.couchbase.client.core.api.query.CoreQueryProfile
import com.couchbase.client.core.api.query.CoreQueryScanConsistency
import com.couchbase.client.core.api.shared.CoreMutationState
import com.couchbase.client.core.deps.com.fasterxml.jackson.databind.JsonNode
import com.couchbase.client.core.deps.com.fasterxml.jackson.databind.node.ArrayNode
import com.couchbase.client.core.deps.com.fasterxml.jackson.databind.node.ObjectNode
import com.couchbase.client.core.endpoint.http.CoreCommonOptions
import com.couchbase.client.core.json.Mapper
import com.couchbase.client.core.transaction.config.CoreSingleQueryTransactionOptions
import com.couchbase.client.kotlin.CommonOptions
import com.couchbase.client.kotlin.codec.JsonSerializer
import com.couchbase.client.kotlin.codec.typeRef
import com.couchbase.client.kotlin.query.QueryFlowItem
import com.couchbase.client.kotlin.query.QueryMetadata
import com.couchbase.client.kotlin.query.QueryParameters
import com.couchbase.client.kotlin.query.QueryProfile
import com.couchbase.client.kotlin.query.QueryRow
import com.couchbase.client.kotlin.query.QueryScanConsistency
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.reactive.asFlow
import kotlinx.coroutines.reactive.awaitSingle
import kotlin.time.toJavaDuration

internal class QueryExecutor(
    private val queryOps: CoreQueryOps,
    private val queryContext: CoreQueryContext?,
    private val defaultSerializer: JsonSerializer,
) {
    fun query(
        statement: String,
        common: CommonOptions,
        parameters: QueryParameters,
        preserveExpiry: Boolean,

        serializer: JsonSerializer?,

        consistency: QueryScanConsistency,
        readonly: Boolean,
        adhoc: Boolean,
        flexIndex: Boolean,

        metrics: Boolean,
        profile: QueryProfile,

        maxParallelism: Int?,
        scanCap: Int?,
        pipelineBatch: Int?,
        pipelineCap: Int?,

        clientContextId: String?,
        raw: Map,
        useReplica: Boolean?,
    ): Flow {
        val actualSerializer = serializer ?: defaultSerializer

        val coreQueryOpts = CoreQueryOptions(
            common = common,
            parameters = parameters,
            preserveExpiry = preserveExpiry,
            actualSerializer = actualSerializer,
            consistency = consistency,
            readonly = readonly,
            adhoc = adhoc,
            flexIndex = flexIndex,
            metrics = metrics,
            profile = profile,
            maxParallelism = maxParallelism,
            scanCap = scanCap,
            pipelineBatch = pipelineBatch,
            pipelineCap = pipelineCap,
            clientContextId = clientContextId,
            raw = raw,
            useReplica = useReplica,
        )

        return flow {
            val response = queryOps.queryReactive(
                statement,
                coreQueryOpts,
                queryContext,
                null,
                null,
            ).awaitSingle()

            emitAll(
                response.rows().asFlow().map {
                    QueryRow(it.data(), actualSerializer)
                }
            )

            emit(QueryMetadata(response.metaData().awaitSingle()))
        }
    }
}

internal fun CoreQueryOptions(
    common: CommonOptions,
    parameters: QueryParameters,
    preserveExpiry: Boolean,

    actualSerializer: JsonSerializer,

    consistency: QueryScanConsistency,
    readonly: Boolean,
    adhoc: Boolean,
    flexIndex: Boolean,

    metrics: Boolean,
    profile: QueryProfile,

    maxParallelism: Int?,
    scanCap: Int?,
    pipelineBatch: Int?,
    pipelineCap: Int?,

    clientContextId: String?,
    raw: Map,
    useReplica: Boolean?,
): CoreQueryOptions {
    return object : CoreQueryOptions {
        override fun adhoc(): Boolean = adhoc
        override fun clientContextId(): String? = clientContextId
        override fun consistentWith(): CoreMutationState? = (consistency as? QueryScanConsistency.ConsistentWith)?.toCore()
        override fun maxParallelism(): Int? = maxParallelism
        override fun metrics(): Boolean = metrics

        override fun namedParameters(): ObjectNode? = (parameters as? QueryParameters.Named)?.let { params ->
            // Let the user's serializer serialize arguments
            val map = mutableMapOf()
            params.inject(map)
            val queryBytes = actualSerializer.serialize(map, typeRef())
            Mapper.decodeIntoTree(queryBytes) as ObjectNode
        }

        override fun pipelineBatch(): Int? = pipelineBatch
        override fun pipelineCap(): Int? = pipelineCap

        override fun positionalParameters(): ArrayNode? = (parameters as? QueryParameters.Positional)?.let { params ->
            // Let the user's serializer serialize arguments
            val map = mutableMapOf()
            params.inject(map)
            val queryBytes = actualSerializer.serialize(map, typeRef())
            Mapper.decodeIntoTree(queryBytes).get("args") as? ArrayNode
        }

        override fun profile(): CoreQueryProfile = profile.core

        override fun raw(): JsonNode? {
            if (raw.isEmpty()) return null
            val rawBytes = actualSerializer.serialize(raw, typeRef())
            return Mapper.decodeIntoTree(rawBytes)
        }

        override fun readonly(): Boolean = readonly
        override fun scanWait(): java.time.Duration? = consistency.scanWait?.toJavaDuration()
        override fun scanCap(): Int? = scanCap

        override fun scanConsistency(): CoreQueryScanConsistency? = when (consistency) {
            is QueryScanConsistency.NotBounded -> CoreQueryScanConsistency.NOT_BOUNDED
            is QueryScanConsistency.RequestPlus -> CoreQueryScanConsistency.REQUEST_PLUS
            else -> null // ConsistentWith is handled by separate consistentWith() accessor
        }

        override fun flexIndex(): Boolean = flexIndex
        override fun preserveExpiry(): Boolean? = if (preserveExpiry) true else null
        override fun asTransactionOptions(): CoreSingleQueryTransactionOptions? = null
        override fun commonOptions(): CoreCommonOptions = common.toCore()
        override fun useReplica(): Boolean? = useReplica
    }
}