org.apache.geode.management.internal.beans.QueryDataFunction.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.geode.management.internal.beans.QueryDataFunction.java

Source

/*
 * 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 org.apache.geode.management.internal.beans;

import org.apache.commons.lang.StringUtils;
import org.apache.geode.SystemFailure;
import org.apache.geode.cache.CacheFactory;
import org.apache.geode.cache.DataPolicy;
import org.apache.geode.cache.Region;
import org.apache.geode.cache.execute.Function;
import org.apache.geode.cache.execute.FunctionAdapter;
import org.apache.geode.cache.execute.FunctionContext;
import org.apache.geode.cache.execute.FunctionException;
import org.apache.geode.cache.execute.FunctionService;
import org.apache.geode.cache.execute.RegionFunctionContext;
import org.apache.geode.cache.execute.ResultCollector;
import org.apache.geode.cache.query.Query;
import org.apache.geode.cache.query.QueryInvalidException;
import org.apache.geode.cache.query.QueryService;
import org.apache.geode.cache.query.SelectResults;
import org.apache.geode.cache.query.internal.CompiledValue;
import org.apache.geode.cache.query.internal.DefaultQuery;
import org.apache.geode.cache.query.internal.QCompiler;
import org.apache.geode.distributed.DistributedMember;
import org.apache.geode.internal.InternalEntity;
import org.apache.geode.internal.cache.BucketRegion;
import org.apache.geode.internal.cache.InternalCache;
import org.apache.geode.internal.cache.LocalDataSet;
import org.apache.geode.internal.cache.PartitionedRegion;
import org.apache.geode.internal.cache.PartitionedRegionHelper;
import org.apache.geode.internal.logging.LogService;
import org.apache.geode.management.DistributedRegionMXBean;
import org.apache.geode.management.ManagementService;
import org.apache.geode.management.internal.ManagementConstants;
import org.apache.geode.management.internal.ManagementStrings;
import org.apache.geode.management.internal.SystemManagementService;
import org.apache.geode.management.internal.cli.commands.DataCommandsUtils;
import org.apache.geode.management.internal.cli.json.GfJsonException;
import org.apache.geode.management.internal.cli.json.GfJsonObject;
import org.apache.geode.management.internal.cli.json.TypedJson;
import org.apache.logging.log4j.Logger;

import java.io.Serializable;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * This function is executed on one or multiple members based on the member input to
 * DistributedSystemMXBean.queryData()
 */
@SuppressWarnings({ "deprecation", "unchecked" })
public class QueryDataFunction extends FunctionAdapter implements InternalEntity {

    private static final long serialVersionUID = 1L;

    private static final Logger logger = LogService.getLogger();

    private static final String MEMBER_KEY = "member";
    private static final String RESULT_KEY = "result";
    private static final String NO_DATA_FOUND = "No Data Found";
    private static final String QUERY_EXEC_SUCCESS = "Query Executed Successfully";
    private static final int DISPLAY_MEMBERWISE = 0;
    private static final int QUERY = 1;
    private static final int REGION = 2;
    private static final int LIMIT = 3;
    private static final int QUERY_RESULTSET_LIMIT = 4;
    private static final int QUERY_COLLECTIONS_DEPTH = 5;
    private static final String SELECT_EXPR = "\\s*SELECT\\s+.+\\s+FROM.+";
    private static final Pattern SELECT_EXPR_PATTERN = Pattern.compile(SELECT_EXPR, Pattern.CASE_INSENSITIVE);
    private static final String SELECT_WITH_LIMIT_EXPR = "\\s*SELECT\\s+.+\\s+FROM(\\s+|(.*\\s+))LIMIT\\s+[0-9]+.*";
    private static final Pattern SELECT_WITH_LIMIT_EXPR_PATTERN = Pattern.compile(SELECT_WITH_LIMIT_EXPR,
            Pattern.CASE_INSENSITIVE);

    @Override
    public boolean hasResult() {
        return true;
    }

    @Override
    public void execute(final FunctionContext context) {
        Object[] functionArgs = (Object[]) context.getArguments();
        boolean showMember = (Boolean) functionArgs[DISPLAY_MEMBERWISE];
        String queryString = (String) functionArgs[QUERY];
        String regionName = (String) functionArgs[REGION];
        int limit = (Integer) functionArgs[LIMIT];

        int queryResultSetLimit = (Integer) functionArgs[QUERY_RESULTSET_LIMIT];

        int queryCollectionsDepth = (Integer) functionArgs[QUERY_COLLECTIONS_DEPTH];

        try {
            context.getResultSender().lastResult(selectWithType(context, queryString, showMember, regionName, limit,
                    queryResultSetLimit, queryCollectionsDepth));
        } catch (Exception e) {
            context.getResultSender().sendException(e);
        }
    }

    @Override
    public String getId() {
        return ManagementConstants.QUERY_DATA_FUNCTION;
    }

    private QueryDataFunctionResult selectWithType(final FunctionContext context, String queryString,
            final boolean showMember, final String regionName, final int limit, final int queryResultSetLimit,
            final int queryCollectionsDepth) throws Exception {
        InternalCache cache = getCache();
        Function localQueryFunc = new LocalQueryFunction("LocalQueryFunction", regionName, showMember)
                .setOptimizeForWrite(true);
        queryString = applyLimitClause(queryString, limit, queryResultSetLimit);

        try {
            TypedJson result = new TypedJson(queryCollectionsDepth);

            Region region = cache.getRegion(regionName);

            if (region == null) {
                throw new Exception(ManagementStrings.QUERY__MSG__REGIONS_NOT_FOUND_ON_MEMBER.toLocalizedString(
                        regionName, cache.getDistributedSystem().getDistributedMember().getId()));
            }

            Object results = null;

            boolean noDataFound = true;

            if (region.getAttributes().getDataPolicy() == DataPolicy.NORMAL) {
                QueryService queryService = cache.getQueryService();

                Query query = queryService.newQuery(queryString);
                results = query.execute();

            } else {
                ResultCollector rcollector;

                PartitionedRegion parRegion = PartitionedRegionHelper.getPartitionedRegion(regionName, cache);
                if (parRegion != null && showMember) {
                    if (parRegion.isDataStore()) {

                        Set<BucketRegion> localPrimaryBucketRegions = parRegion.getDataStore()
                                .getAllLocalPrimaryBucketRegions();
                        Set<Integer> localPrimaryBucketSet = new HashSet<>();
                        for (BucketRegion bRegion : localPrimaryBucketRegions) {
                            localPrimaryBucketSet.add(bRegion.getId());
                        }
                        LocalDataSet lds = new LocalDataSet(parRegion, localPrimaryBucketSet);
                        DefaultQuery query = (DefaultQuery) cache.getQueryService().newQuery(queryString);
                        results = lds.executeQuery(query, null, localPrimaryBucketSet);
                    }
                } else {
                    rcollector = FunctionService.onRegion(cache.getRegion(regionName)).setArguments(queryString)
                            .execute(localQueryFunc);
                    results = rcollector.getResult();
                }
            }

            if (results != null && results instanceof SelectResults) {

                SelectResults selectResults = (SelectResults) results;
                for (Object object : selectResults) {
                    result.add(RESULT_KEY, object);
                    noDataFound = false;
                }
            } else if (results != null && results instanceof ArrayList) {
                ArrayList listResults = (ArrayList) results;
                ArrayList actualResult = (ArrayList) listResults.get(0);
                for (Object object : actualResult) {
                    result.add(RESULT_KEY, object);
                    noDataFound = false;
                }
            }

            if (!noDataFound && showMember) {
                result.add(MEMBER_KEY, cache.getDistributedSystem().getDistributedMember().getId());
            }

            if (noDataFound) {
                return new QueryDataFunctionResult(QUERY_EXEC_SUCCESS,
                        BeanUtilFuncs.compress(new JsonisedErrorMessage(NO_DATA_FOUND).toString()));
            }
            return new QueryDataFunctionResult(QUERY_EXEC_SUCCESS, BeanUtilFuncs.compress(result.toString()));
        } catch (Exception e) {
            logger.warn(e.getMessage(), e);
            throw e;
        }
    }

    /**
     * Matches the input query with query with limit pattern. If limit is found in input query this
     * function ignores. Else it will append a default limit .. 1000 If input limit is 0 then also it
     * will append default limit of 1000
     *
     * @param query input query
     * @param limit limit on the result set
     *
     * @return a string having limit clause
     */
    protected static String applyLimitClause(final String query, int limit, final int queryResultSetLimit) {

        Matcher matcher = SELECT_EXPR_PATTERN.matcher(query);

        if (matcher.matches()) {
            Matcher limit_matcher = SELECT_WITH_LIMIT_EXPR_PATTERN.matcher(query);
            boolean queryAlreadyHasLimitClause = limit_matcher.matches();

            if (!queryAlreadyHasLimitClause) {
                if (limit == 0) {
                    limit = queryResultSetLimit;
                }
                String result = query;
                result += " LIMIT " + limit;
                return result;
            }
        }
        return query;
    }

    private static Object callFunction(final Object functionArgs, final Set<DistributedMember> members,
            final boolean zipResult) throws Exception {

        try {
            if (members.size() == 1) {
                DistributedMember member = members.iterator().next();
                ResultCollector collector = FunctionService.onMember(member).setArguments(functionArgs)
                        .execute(ManagementConstants.QUERY_DATA_FUNCTION);
                List list = (List) collector.getResult();
                Object object = null;
                if (list.size() > 0) {
                    object = list.get(0);
                }

                if (object instanceof Throwable) {
                    throw (Throwable) object;
                }

                QueryDataFunctionResult result = (QueryDataFunctionResult) object;
                if (zipResult) { // The result is already compressed
                    return result.compressedBytes;
                } else {
                    Object[] functionArgsList = (Object[]) functionArgs;
                    boolean showMember = (Boolean) functionArgsList[DISPLAY_MEMBERWISE];
                    if (showMember) {// Added to show a single member similar to multiple
                        // member.
                        // Note , if no member is selected this is the code path executed. A
                        // random associated member is chosen.
                        List<String> decompressedList = new ArrayList<>();
                        decompressedList.add(BeanUtilFuncs.decompress(result.compressedBytes));
                        return wrapResult(decompressedList.toString());
                    }
                    return BeanUtilFuncs.decompress(result.compressedBytes);
                }

            } else { // More than 1 Member
                ResultCollector coll = FunctionService.onMembers(members).setArguments(functionArgs)
                        .execute(ManagementConstants.QUERY_DATA_FUNCTION);

                List list = (List) coll.getResult();
                Object object = list.get(0);
                if (object instanceof Throwable) {
                    throw (Throwable) object;
                }

                Iterator<QueryDataFunctionResult> it = list.iterator();
                List<String> decompressedList = new ArrayList<>();

                while (it.hasNext()) {
                    String decompressedStr;
                    decompressedStr = BeanUtilFuncs.decompress(it.next().compressedBytes);
                    decompressedList.add(decompressedStr);
                }

                if (zipResult) {
                    return BeanUtilFuncs.compress(wrapResult(decompressedList.toString()));
                } else {
                    return wrapResult(decompressedList.toString());
                }

            }
        } catch (FunctionException fe) {
            throw new Exception(ManagementStrings.QUERY__MSG__QUERY_EXEC.toLocalizedString(fe.getMessage()));
        } catch (VirtualMachineError e) {
            SystemFailure.initiateFailure(e);
            throw e;
        } catch (Throwable e) {
            SystemFailure.checkFailure();
            throw new Exception(ManagementStrings.QUERY__MSG__QUERY_EXEC.toLocalizedString(e.getMessage()));
        }
    }

    private static String wrapResult(final String str) {
        StringWriter w = new StringWriter();
        synchronized (w.getBuffer()) {
            w.write("{\"result\":");
            w.write(str);
            w.write("}");
            return w.toString();
        }
    }

    public static Object queryData(final String query, final String members, final int limit,
            final boolean zipResult, final int queryResultSetLimit, final int queryCollectionsDepth)
            throws Exception {

        if (query == null || query.isEmpty()) {
            return new JsonisedErrorMessage(ManagementStrings.QUERY__MSG__QUERY_EMPTY.toLocalizedString())
                    .toString();
        }

        Set<DistributedMember> inputMembers = null;
        if (StringUtils.isNotBlank(members)) {
            inputMembers = new HashSet<>();
            StringTokenizer st = new StringTokenizer(members, ",");
            while (st.hasMoreTokens()) {
                String member = st.nextToken();
                DistributedMember distributedMember = BeanUtilFuncs.getDistributedMemberByNameOrId(member);
                inputMembers.add(distributedMember);
                if (distributedMember == null) {
                    return new JsonisedErrorMessage(
                            ManagementStrings.QUERY__MSG__INVALID_MEMBER.toLocalizedString(member)).toString();
                }
            }
        }

        InternalCache cache = (InternalCache) CacheFactory.getAnyInstance();
        try {

            SystemManagementService service = (SystemManagementService) ManagementService
                    .getExistingManagementService(cache);
            Set<String> regionsInQuery = compileQuery(cache, query);

            // Validate region existence
            if (regionsInQuery.size() > 0) {
                for (String regionPath : regionsInQuery) {
                    DistributedRegionMXBean regionMBean = service.getDistributedRegionMXBean(regionPath);
                    if (regionMBean == null) {
                        return new JsonisedErrorMessage(
                                ManagementStrings.QUERY__MSG__REGIONS_NOT_FOUND.toLocalizedString(regionPath))
                                        .toString();
                    } else {
                        Set<DistributedMember> associatedMembers = DataCommandsUtils
                                .getRegionAssociatedMembers(regionPath, cache, true);

                        if (inputMembers != null && inputMembers.size() > 0) {
                            if (!associatedMembers.containsAll(inputMembers)) {
                                return new JsonisedErrorMessage(
                                        ManagementStrings.QUERY__MSG__REGIONS_NOT_FOUND_ON_MEMBERS
                                                .toLocalizedString(regionPath)).toString();
                            }
                        }
                    }
                }
            } else {
                return new JsonisedErrorMessage(ManagementStrings.QUERY__MSG__INVALID_QUERY
                        .toLocalizedString("Region mentioned in query probably missing /")).toString();
            }

            // Validate
            if (regionsInQuery.size() > 1 && inputMembers == null) {
                for (String regionPath : regionsInQuery) {
                    DistributedRegionMXBean regionMBean = service.getDistributedRegionMXBean(regionPath);

                    if (regionMBean.getRegionType().equals(DataPolicy.PARTITION.toString())
                            || regionMBean.getRegionType().equals(DataPolicy.PERSISTENT_PARTITION.toString())) {
                        return new JsonisedErrorMessage(
                                ManagementStrings.QUERY__MSG__JOIN_OP_EX.toLocalizedString()).toString();
                    }
                }
            }

            String randomRegion = regionsInQuery.iterator().next();

            Set<DistributedMember> associatedMembers = DataCommandsUtils
                    .getQueryRegionsAssociatedMembers(regionsInQuery, cache, false);// First
            // available
            // member

            if (associatedMembers != null && associatedMembers.size() > 0) {
                Object[] functionArgs = new Object[6];
                if (inputMembers != null && inputMembers.size() > 0) {// on input
                    // members

                    functionArgs[DISPLAY_MEMBERWISE] = true;
                    functionArgs[QUERY] = query;
                    functionArgs[REGION] = randomRegion;
                    functionArgs[LIMIT] = limit;
                    functionArgs[QUERY_RESULTSET_LIMIT] = queryResultSetLimit;
                    functionArgs[QUERY_COLLECTIONS_DEPTH] = queryCollectionsDepth;
                    return callFunction(functionArgs, inputMembers, zipResult);
                } else { // Query on any random member
                    functionArgs[DISPLAY_MEMBERWISE] = false;
                    functionArgs[QUERY] = query;
                    functionArgs[REGION] = randomRegion;
                    functionArgs[LIMIT] = limit;
                    functionArgs[QUERY_RESULTSET_LIMIT] = queryResultSetLimit;
                    functionArgs[QUERY_COLLECTIONS_DEPTH] = queryCollectionsDepth;
                    return callFunction(functionArgs, associatedMembers, zipResult);
                }

            } else {
                return new JsonisedErrorMessage(ManagementStrings.QUERY__MSG__REGIONS_NOT_FOUND
                        .toLocalizedString(regionsInQuery.toString())).toString();
            }

        } catch (QueryInvalidException qe) {
            return new JsonisedErrorMessage(
                    ManagementStrings.QUERY__MSG__INVALID_QUERY.toLocalizedString(qe.getMessage())).toString();
        }
    }

    private InternalCache getCache() {
        return (InternalCache) CacheFactory.getAnyInstance();
    }

    private static class JsonisedErrorMessage {

        private static String message = "message";

        private GfJsonObject gFJsonObject = new GfJsonObject();

        public JsonisedErrorMessage(final String errorMessage) throws Exception {
            try {
                gFJsonObject.put(message, errorMessage);
            } catch (GfJsonException e) {
                throw new Exception(e);
            }
        }

        @Override
        public String toString() {
            return gFJsonObject.toString();
        }
    }

    /**
     * Compile the query and return a set of regions involved in the query It throws an
     * QueryInvalidException if the query is not proper
     *
     * @param cache current cache
     * @param query input query
     *
     * @return a set of regions involved in the query
     */
    private static Set<String> compileQuery(final InternalCache cache, final String query)
            throws QueryInvalidException {
        QCompiler compiler = new QCompiler();
        Set<String> regionsInQuery;
        try {
            CompiledValue compiledQuery = compiler.compileQuery(query);
            Set<String> regions = new HashSet<>();
            compiledQuery.getRegionsInQuery(regions, null);
            regionsInQuery = Collections.unmodifiableSet(regions);
            return regionsInQuery;
        } catch (QueryInvalidException qe) {
            logger.error("{} Failed, Error {}", query, qe.getMessage(), qe);
            throw qe;
        }
    }

    /**
     * Function to gather data locally. This function is required to execute query with region context
     */
    private class LocalQueryFunction extends FunctionAdapter {

        private static final long serialVersionUID = 1L;

        private final String id;

        private boolean optimizeForWrite = false;
        private boolean showMembers = false;
        private String regionName;

        public LocalQueryFunction(final String id, final String regionName, final boolean showMembers) {
            super();
            this.id = id;
            this.regionName = regionName;
            this.showMembers = showMembers;
        }

        @Override
        public boolean hasResult() {
            return true;
        }

        @Override
        public boolean isHA() {
            return false;
        }

        @Override
        public boolean optimizeForWrite() {
            return optimizeForWrite;
        }

        public LocalQueryFunction setOptimizeForWrite(final boolean optimizeForWrite) {
            this.optimizeForWrite = optimizeForWrite;
            return this;
        }

        @Override
        public void execute(final FunctionContext context) {
            InternalCache cache = getCache();
            QueryService queryService = cache.getQueryService();
            String qstr = (String) context.getArguments();
            Region r = cache.getRegion(regionName);
            try {
                Query query = queryService.newQuery(qstr);
                SelectResults sr;
                if (r.getAttributes().getPartitionAttributes() != null && showMembers) {
                    sr = (SelectResults) query.execute((RegionFunctionContext) context);
                    context.getResultSender().lastResult(sr.asList());
                } else {
                    sr = (SelectResults) query.execute();
                    context.getResultSender().lastResult(sr.asList());
                }

            } catch (Exception e) {
                throw new FunctionException(e);
            }
        }

        private InternalCache getCache() {
            return (InternalCache) CacheFactory.getAnyInstance();
        }

        @Override
        public String getId() {
            return this.id;
        }
    }

    private static class QueryDataFunctionResult implements Serializable {

        private static final long serialVersionUID = 1L;

        private final String message;
        private final byte[] compressedBytes;

        public QueryDataFunctionResult(final String message, final byte[] compressedBytes) {
            this.message = message;
            this.compressedBytes = compressedBytes;
        }

        public String getMessage() {
            return message;
        }

        public byte[] getCompressedBytes() {
            return compressedBytes;
        }
    }
}