io.trino.plugin.session.SessionMatchSpec Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trino-session-property-managers Show documentation
Trino - Session property managers
There is a newer version: 465
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 io.trino.plugin.session;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import io.trino.spi.session.SessionConfigurationContext;
import org.jdbi.v3.core.mapper.RowMapper;
import org.jdbi.v3.core.statement.StatementContext;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.regex.Pattern;

import static com.google.common.base.Preconditions.checkArgument;
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 resourceGroupRegex;
    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("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");
        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 (resourceGroupRegex.isPresent() && !resourceGroupRegex.get().matcher(context.getResourceGroupId().toString()).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 Map getSessionProperties()
    {
        return sessionProperties;
    }

    public static class Mapper
            implements RowMapper
    {
        @Override
        public SessionMatchSpec map(ResultSet resultSet, StatementContext context)
                throws SQLException
        {
            Map sessionProperties = getProperties(
                    Optional.ofNullable(resultSet.getString("session_property_names")),
                    Optional.ofNullable(resultSet.getString("session_property_values")));

            return new SessionMatchSpec(
                    Optional.ofNullable(resultSet.getString("user_regex")).map(Pattern::compile),
                    Optional.ofNullable(resultSet.getString("source_regex")).map(Pattern::compile),
                    Optional.ofNullable(resultSet.getString("client_tags")).map(tag -> Splitter.on(",").splitToList(tag)),
                    Optional.ofNullable(resultSet.getString("query_type")),
                    Optional.ofNullable(resultSet.getString("group_regex")).map(Pattern::compile),
                    sessionProperties);
        }

        private static Map getProperties(Optional names, Optional values)
        {
            if (names.isEmpty()) {
                return ImmutableMap.of();
            }

            checkArgument(values.isPresent(), "names are present, but values are not");
            List sessionPropertyNames = Splitter.on(",").splitToList(names.get());
            List sessionPropertyValues = Splitter.on(",").splitToList(values.get());
            checkArgument(sessionPropertyNames.size() == sessionPropertyValues.size(),
                    "The number of property names and values should be the same");

            Map sessionProperties = new HashMap<>();
            for (int i = 0; i < sessionPropertyNames.size(); i++) {
                sessionProperties.put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
            }

            return sessionProperties;
        }
    }
}