com.microsoft.azure.documentdb.internal.directconnectivity.ServerStoreModel Maven / Gradle / Ivy

Go to download
package com.microsoft.azure.documentdb.internal.directconnectivity;

import java.util.concurrent.ExecutorService;

import com.microsoft.azure.documentdb.RetryOptions;
import com.microsoft.azure.documentdb.internal.RefreshAuthTokenDelegate;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.microsoft.azure.documentdb.ConsistencyLevel;
import com.microsoft.azure.documentdb.DocumentClientException;
import com.microsoft.azure.documentdb.internal.AuthorizationTokenProvider;
import com.microsoft.azure.documentdb.internal.DatabaseAccountConfigurationProvider;
import com.microsoft.azure.documentdb.internal.DocumentServiceRequest;
import com.microsoft.azure.documentdb.internal.DocumentServiceResponse;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.RetryRequestDelegate;
import com.microsoft.azure.documentdb.internal.RetryUtility;
import com.microsoft.azure.documentdb.internal.SessionContainer;
import com.microsoft.azure.documentdb.internal.SessionTokenHelper;
import com.microsoft.azure.documentdb.internal.StoreModel;

/**
 * Used internally to provides functionality to communicate and process response in direct connectivity mode
 * in the Azure Cosmos DB database service.
 */
public class ServerStoreModel implements StoreModel {
    private final SessionContainer sessionContainer;
    private final ReplicatedResourceClient replicatedResourceClient;
    private Integer defaultReplicaIndex;
    private boolean forceAddressRefresh;
    private RefreshAuthTokenDelegate refreshAuthTokenDelegate;
    private RetryOptions retryOptions;

    public ServerStoreModel(
            TransportClient transportClient,
            GlobalAddressResolver globalAddressResolver,
            SessionContainer sessionContainer,
            DatabaseAccountConfigurationProvider databaseAccountConfigurationProvider,
            AuthorizationTokenProvider authorizationTokenProvider,
            ExecutorService executorService,
            RefreshAuthTokenDelegate refreshAuthTokenDelegate,
            RetryOptions retryOptions,
            boolean useMultipleWriteLocations) {

        this.sessionContainer = sessionContainer;
        this.replicatedResourceClient = new ReplicatedResourceClient(
                globalAddressResolver,
                sessionContainer,
                transportClient,
                databaseAccountConfigurationProvider,
                authorizationTokenProvider,
                executorService,
                useMultipleWriteLocations);
        this.refreshAuthTokenDelegate = refreshAuthTokenDelegate;
        this.retryOptions = retryOptions;
    }

    public Integer getDefaultReplicaIndex() {
        return defaultReplicaIndex;
    }

    public void setDefaultReplicaIndex(int defaultReplicaIndex) {
        this.defaultReplicaIndex = defaultReplicaIndex;
    }

    public String getLastReadAddress() {
        return replicatedResourceClient.getLastReadAddress();
    }

    public void setLastReadAddress(String lastReadAddress) {
        replicatedResourceClient.setLastReadAddress(lastReadAddress);
    }

    public boolean isForceAddressRefresh() {
        return forceAddressRefresh;
    }

    public void setForceAddressRefresh(boolean forceAddressRefresh) {
        this.forceAddressRefresh = forceAddressRefresh;
    }

    @Override
    public DocumentServiceResponse processMessage(final DocumentServiceRequest request) throws DocumentClientException {
        if (this.defaultReplicaIndex != null) {
            request.setDefaultReplicaIndex(this.defaultReplicaIndex);
        }

        final String consistency = request.getHeaders().get(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL);
        if (StringUtils.isNotEmpty(consistency)) {
            try {
                request.setOriginalRequestConsistencyLevel(ConsistencyLevel.valueOf(consistency));
            } catch (IllegalArgumentException e) {
                throw new DocumentClientException(HttpConstants.StatusCodes.BADREQUEST, "InvalidHeaderValue ConsistencyLevel " + consistency);
            }
        }

        if (request.isReadingFromMaster()) {
            request.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.Strong.name());
        }

        RetryRequestDelegate processDelegate = new RetryRequestDelegate() {

            @Override
            public DocumentServiceResponse apply(DocumentServiceRequest requestInner) throws DocumentClientException {
                StoreResponse storeResponse;
                try {
                    storeResponse = replicatedResourceClient.invoke(requestInner);
                    if (request.getClientSideRequestStatistics() != null) {
                        storeResponse.setClientSideRequestStatistics(request.getClientSideRequestStatistics());
                    }
                } catch (DocumentClientException dce) {
                    if (request.getClientSideRequestStatistics() != null) {
                        dce.setClientSideRequestStatistics(request.getClientSideRequestStatistics());
                    }
                    SessionTokenHelper.updateSession(ServerStoreModel.this.sessionContainer,
                            request, dce);
                    throw dce;
                }

                return new DocumentServiceResponse(storeResponse);
            }
        };

        return RetryUtility.executeStoreClientRequest(processDelegate, this.retryOptions, request, this.refreshAuthTokenDelegate);
    }
}