com.facebook.presto.session.SessionMatchSpec Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of presto-session-property-managers Show documentation
Presto - Session Property Managers
There is a newer version: 0.289
Show newest version
/*
 * 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.facebook.presto.session;

import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
import com.facebook.presto.spi.session.SessionConfigurationContext;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.regex.Pattern;

import static java.util.Objects.requireNonNull;

public class SessionMatchSpec
{
    private final Optional userRegex;
    private final Optional sourceRegex;
    private final Set clientTags;
    private final Optional queryType;
    private final Optional clientInfoRegex;
    private final Optional resourceGroupRegex;
    private final Optional overrideSessionProperties;
    private final Map sessionProperties;

    @JsonCreator
    public SessionMatchSpec(
            @JsonProperty("user") Optional userRegex,
            @JsonProperty("source") Optional sourceRegex,
            @JsonProperty("clientTags") Optional> clientTags,
            @JsonProperty("queryType") Optional queryType,
            @JsonProperty("group") Optional resourceGroupRegex,
            @JsonProperty("clientInfo") Optional clientInfoRegex,
            @JsonProperty("overrideSessionProperties") Optional overrideSessionProperties,
            @JsonProperty("sessionProperties") Map sessionProperties)
    {
        this.userRegex = requireNonNull(userRegex, "userRegex is null");
        this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
        requireNonNull(clientTags, "clientTags is null");
        this.clientTags = ImmutableSet.copyOf(clientTags.orElse(ImmutableList.of()));
        this.queryType = requireNonNull(queryType, "queryType is null");
        this.resourceGroupRegex = requireNonNull(resourceGroupRegex, "resourceGroupRegex is null");
        this.clientInfoRegex = requireNonNull(clientInfoRegex, "clientInfoRegex is null");
        this.overrideSessionProperties = requireNonNull(overrideSessionProperties, "overrideSessionProperties is null");
        requireNonNull(sessionProperties, "sessionProperties is null");
        this.sessionProperties = ImmutableMap.copyOf(sessionProperties);
    }

    public Map match(SessionConfigurationContext context)
    {
        if (userRegex.isPresent() && !userRegex.get().matcher(context.getUser()).matches()) {
            return ImmutableMap.of();
        }
        if (sourceRegex.isPresent()) {
            String source = context.getSource().orElse("");
            if (!sourceRegex.get().matcher(source).matches()) {
                return ImmutableMap.of();
            }
        }
        if (!clientTags.isEmpty() && !context.getClientTags().containsAll(clientTags)) {
            return ImmutableMap.of();
        }

        if (queryType.isPresent()) {
            String contextQueryType = context.getQueryType().orElse("");
            if (!queryType.get().equalsIgnoreCase(contextQueryType)) {
                return ImmutableMap.of();
            }
        }

        if (clientInfoRegex.isPresent()) {
            String clientInfo = context.getClientInfo().orElse("");
            if (!clientInfoRegex.get().matcher(clientInfo).matches()) {
                return ImmutableMap.of();
            }
        }

        if (resourceGroupRegex.isPresent()) {
            String resourceGroupId = context.getResourceGroupId().map(ResourceGroupId::toString).orElse("");
            if (!resourceGroupRegex.get().matcher(resourceGroupId).matches()) {
                return ImmutableMap.of();
            }
        }

        return sessionProperties;
    }

    @JsonProperty("user")
    public Optional getUserRegex()
    {
        return userRegex;
    }

    @JsonProperty("source")
    public Optional getSourceRegex()
    {
        return sourceRegex;
    }

    @JsonProperty
    public Set getClientTags()
    {
        return clientTags;
    }

    @JsonProperty
    public Optional getQueryType()
    {
        return queryType;
    }

    @JsonProperty("group")
    public Optional getResourceGroupRegex()
    {
        return resourceGroupRegex;
    }

    @JsonProperty
    public Optional getClientInfoRegex()
    {
        return clientInfoRegex;
    }

    @JsonProperty
    public Optional getOverrideSessionProperties()
    {
        return overrideSessionProperties;
    }

    @JsonProperty
    public Map getSessionProperties()
    {
        return sessionProperties;
    }
}