All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.milvus.client.MilvusMultiServiceClient Maven / Gradle / Ivy

Go to download

Java SDK for Milvus, a distributed high-performance vector search engine. update grpc to 1.42.1 update protobuf to 3.19.1

There is a newer version: 2.2.2.1
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 io.milvus.client;

import com.google.common.util.concurrent.ListenableFuture;
import io.milvus.connection.ClusterFactory;
import io.milvus.connection.ServerSetting;
import io.milvus.grpc.*;
import io.milvus.param.*;
import io.milvus.param.alias.AlterAliasParam;
import io.milvus.param.alias.CreateAliasParam;
import io.milvus.param.alias.DropAliasParam;
import io.milvus.param.collection.*;
import io.milvus.param.control.*;
import io.milvus.param.credential.*;
import io.milvus.param.dml.*;
import io.milvus.param.index.*;
import io.milvus.param.partition.*;
import lombok.NonNull;
import org.apache.commons.collections4.CollectionUtils;

import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

public class MilvusMultiServiceClient implements MilvusClient {

    private final ClusterFactory clusterFactory;

    /**
     * Sets connect param for multi milvus clusters.
     * @param multiConnectParam multi server connect param
     */
    public MilvusMultiServiceClient(@NonNull MultiConnectParam multiConnectParam) {

        List serverSettings = multiConnectParam.getHosts().stream()
                .map(host -> {

                    MilvusClient milvusClient = buildMilvusClient(host, multiConnectParam);

                    return ServerSetting.newBuilder()
                            .withHost(host)
                            .withMilvusClient(milvusClient).build();

                }).collect(Collectors.toList());

        boolean keepMonitor = serverSettings.size() > 1;

        this.clusterFactory = ClusterFactory.newBuilder()
                .withServerSetting(serverSettings)
                .keepMonitor(keepMonitor)
                .withQueryNodeSingleSearch(multiConnectParam.getQueryNodeSingleSearch())
                .build();
    }

    private MilvusClient buildMilvusClient(ServerAddress host, MultiConnectParam multiConnectParam) {
        long connectTimeoutMsm = multiConnectParam.getConnectTimeoutMs();
        long keepAliveTimeMs = multiConnectParam.getKeepAliveTimeMs();
        long keepAliveTimeoutMs = multiConnectParam.getKeepAliveTimeoutMs();
        boolean keepAliveWithoutCalls = multiConnectParam.isKeepAliveWithoutCalls();
        boolean secure = multiConnectParam.isSecure();
        long idleTimeoutMs = multiConnectParam.getIdleTimeoutMs();

        ConnectParam clusterConnectParam = ConnectParam.newBuilder()
                .withHost(host.getHost())
                .withPort(host.getPort())
                .withConnectTimeout(connectTimeoutMsm, TimeUnit.MILLISECONDS)
                .withKeepAliveTime(keepAliveTimeMs, TimeUnit.MILLISECONDS)
                .withKeepAliveTimeout(keepAliveTimeoutMs, TimeUnit.MILLISECONDS)
                .keepAliveWithoutCalls(keepAliveWithoutCalls)
                .secure(secure)
                .withIdleTimeout(idleTimeoutMs, TimeUnit.MILLISECONDS)
                .withAuthorization(multiConnectParam.getAuthorization())
                .build();
        return new MilvusServiceClient(clusterConnectParam);
    }


    @Override
    public MilvusClient withTimeout(long timeout, TimeUnit timeoutUnit) {
        return clusterFactory.getMaster().getClient().withTimeout(timeout, timeoutUnit);
    }

    @Override
    public void close(long maxWaitSeconds) throws InterruptedException {
        this.clusterFactory.getAvailableServerSettings().parallelStream()
                .forEach(serverSetting -> serverSetting.getClient().close());
        this.clusterFactory.close();
    }

    @Override
    public R hasCollection(HasCollectionParam requestParam) {
        return this.clusterFactory.getMaster().getClient().hasCollection(requestParam);
    }

    @Override
    public R createCollection(CreateCollectionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().createCollection(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R dropCollection(DropCollectionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().dropCollection(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R loadCollection(LoadCollectionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().loadCollection(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R releaseCollection(ReleaseCollectionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().releaseCollection(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R describeCollection(DescribeCollectionParam requestParam) {
        return this.clusterFactory.getMaster().getClient().describeCollection(requestParam);
    }

    @Override
    public R getCollectionStatistics(GetCollectionStatisticsParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getCollectionStatistics(requestParam);
    }

    @Override
    public R showCollections(ShowCollectionsParam requestParam) {
        return this.clusterFactory.getMaster().getClient().showCollections(requestParam);
    }

    @Override
    public R flush(FlushParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().flush(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R createPartition(CreatePartitionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().createPartition(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R dropPartition(DropPartitionParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().dropPartition(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R hasPartition(HasPartitionParam requestParam) {
        return this.clusterFactory.getMaster().getClient().hasPartition(requestParam);
    }

    @Override
    public R loadPartitions(LoadPartitionsParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().loadPartitions(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R releasePartitions(ReleasePartitionsParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().releasePartitions(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R getPartitionStatistics(GetPartitionStatisticsParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getPartitionStatistics(requestParam);
    }

    @Override
    public R showPartitions(ShowPartitionsParam requestParam) {
        return this.clusterFactory.getMaster().getClient().showPartitions(requestParam);
    }

    @Override
    public R createAlias(CreateAliasParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().createAlias(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R dropAlias(DropAliasParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().dropAlias(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R alterAlias(AlterAliasParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().alterAlias(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R createIndex(CreateIndexParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().createIndex(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R dropIndex(DropIndexParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().dropIndex(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R describeIndex(DescribeIndexParam requestParam) {
        return this.clusterFactory.getMaster().getClient().describeIndex(requestParam);
    }

    @Override
    public R getIndexState(GetIndexStateParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getIndexState(requestParam);
    }

    @Override
    public R getIndexBuildProgress(GetIndexBuildProgressParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getIndexBuildProgress(requestParam);
    }

    @Override
    public R insert(InsertParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().insert(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public ListenableFuture> insertAsync(InsertParam requestParam) {
        List>> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().insertAsync(requestParam))
                .collect(Collectors.toList());
        return response.get(0);
    }

    @Override
    public R delete(DeleteParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().stream()
                .map(serverSetting -> serverSetting.getClient().delete(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

//    @Override
//    public R bulkload(@NonNull BulkloadParam requestParam) {
//        List> response = this.clusterFactory.getAvailableServerSettings().stream()
//                .map(serverSetting -> serverSetting.getClient().bulkload(requestParam))
//                .collect(Collectors.toList());
//        return handleResponse(response);
//    }
//
//    @Override
//    public R getBulkloadState(GetBulkloadStateParam requestParam) {
//        return this.clusterFactory.getMaster().getClient().getBulkloadState(requestParam);
//    }
//
//    @Override
//    public R listBulkloadTasks(@NonNull ListBulkloadTasksParam requestParam) {
//        return this.clusterFactory.getMaster().getClient().listBulkloadTasks(requestParam);
//    }

    @Override
    public R search(SearchParam requestParam) {
        return this.clusterFactory.getMaster().getClient().search(requestParam);
    }

    @Override
    public ListenableFuture> searchAsync(SearchParam requestParam) {
        return this.clusterFactory.getMaster().getClient().searchAsync(requestParam);
    }

    @Override
    public R query(QueryParam requestParam) {
        return this.clusterFactory.getMaster().getClient().query(requestParam);
    }

    @Override
    public ListenableFuture> queryAsync(QueryParam requestParam) {
        return this.clusterFactory.getMaster().getClient().queryAsync(requestParam);
    }

//    @Override
//    public R calcDistance(CalcDistanceParam requestParam) {
//        return this.clusterFactory.getMaster().getClient().calcDistance(requestParam);
//    }

    @Override
    public R getMetrics(GetMetricsParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getMetrics(requestParam);
    }

    @Override
    public R getFlushState(GetFlushStateParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getFlushState(requestParam);
    }

    @Override
    public R getPersistentSegmentInfo(GetPersistentSegmentInfoParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getPersistentSegmentInfo(requestParam);
    }

    @Override
    public R getQuerySegmentInfo(GetQuerySegmentInfoParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getQuerySegmentInfo(requestParam);
    }

    @Override
    public R getReplicas(GetReplicasParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getReplicas(requestParam);
    }

    @Override
    public R loadBalance(LoadBalanceParam requestParam) {
        List> response = this.clusterFactory.getAvailableServerSettings().parallelStream()
                .map(serverSetting -> serverSetting.getClient().loadBalance(requestParam))
                .collect(Collectors.toList());
        return handleResponse(response);
    }

    @Override
    public R getCompactionState(GetCompactionStateParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getCompactionState(requestParam);
    }

    @Override
    public R manualCompact(ManualCompactParam requestParam) {
        return null;
    }

    @Override
    public R getCompactionStateWithPlans(GetCompactionPlansParam requestParam) {
        return this.clusterFactory.getMaster().getClient().getCompactionStateWithPlans(requestParam);
    }

    @Override
    public R createCredential(CreateCredentialParam requestParam) {
        return this.clusterFactory.getMaster().getClient().createCredential(requestParam);
    }

    @Override
    public R updateCredential(UpdateCredentialParam requestParam) {
        return this.clusterFactory.getMaster().getClient().updateCredential(requestParam);
    }

    @Override
    public R deleteCredential(DeleteCredentialParam requestParam) {
        return this.clusterFactory.getMaster().getClient().deleteCredential(requestParam);
    }

    @Override
    public R listCredUsers(ListCredUsersParam requestParam) {
        return this.clusterFactory.getMaster().getClient().listCredUsers(requestParam);
    }

    private  R handleResponse(List> response) {
        if (CollectionUtils.isNotEmpty(response)) {
            R rSuccess = null;
            for (R singleRes : response) {
                if (R.Status.Success.getCode() == singleRes.getStatus()) {
                    rSuccess = singleRes;
                } else {
                    return singleRes;
                }
            }
            if (null != rSuccess) {
                return rSuccess;
            }
        }
        return R.failed(R.Status.Unknown, "Response is empty.");
    }
}





© 2015 - 2024 Weber Informatics LLC | Privacy Policy