com.hazelcast.jet.sql.impl.expression.json.JsonQueryFunction Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hazelcast-sql Show documentation
SQL Parser and Optimizer
There is a newer version: 5.5.0
/*
 * Copyright 2024 Hazelcast Inc.
 *
 * Licensed under the Hazelcast Community License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://hazelcast.com/hazelcast-community-license
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.hazelcast.jet.sql.impl.expression.json;

import com.hazelcast.core.HazelcastJsonValue;
import com.hazelcast.jet.sql.impl.JetSqlSerializerHook;
import com.hazelcast.logging.ILogger;
import com.hazelcast.logging.Logger;
import com.hazelcast.nio.ObjectDataInput;
import com.hazelcast.nio.ObjectDataOutput;
import com.hazelcast.sql.impl.QueryException;
import com.hazelcast.sql.impl.expression.ConcurrentInitialSetCache;
import com.hazelcast.sql.impl.expression.ConstantExpression;
import com.hazelcast.sql.impl.expression.Expression;
import com.hazelcast.sql.impl.expression.ExpressionEvalContext;
import com.hazelcast.sql.impl.expression.VariExpression;
import com.hazelcast.sql.impl.row.Row;
import com.hazelcast.sql.impl.type.QueryDataType;
import com.hazelcast.shaded.org.apache.calcite.sql.SqlJsonQueryEmptyOrErrorBehavior;
import com.hazelcast.shaded.org.apache.calcite.sql.SqlJsonQueryWrapperBehavior;
import com.hazelcast.shaded.org.jsfr.json.path.JsonPath;

import java.io.IOException;
import java.io.ObjectInputStream;
import java.util.Arrays;
import java.util.Collection;
import java.util.Objects;
import java.util.function.Function;

import static com.hazelcast.jet.sql.impl.expression.json.JsonPathUtil.serialize;
import static com.hazelcast.jet.sql.impl.expression.json.JsonPathUtil.wrapToArray;

@SuppressWarnings("checkstyle:MagicNumber")
public class JsonQueryFunction extends VariExpression {
    private static final ILogger LOGGER = Logger.getLogger(JsonQueryFunction.class);
    private static final Function COMPILE_FUNCTION = JsonPathUtil::compile;

    private transient ConcurrentInitialSetCache pathCache;
    private JsonPath constantPathCache;

    private SqlJsonQueryWrapperBehavior wrapperBehavior;
    private SqlJsonQueryEmptyOrErrorBehavior onEmpty;
    private SqlJsonQueryEmptyOrErrorBehavior onError;

    public JsonQueryFunction() { }

    private JsonQueryFunction(
            Expression[] operands,
            SqlJsonQueryWrapperBehavior wrapperBehavior,
            SqlJsonQueryEmptyOrErrorBehavior onEmpty,
            SqlJsonQueryEmptyOrErrorBehavior onError
    ) {
        super(operands);
        this.wrapperBehavior = wrapperBehavior;
        this.onEmpty = onEmpty;
        this.onError = onError;
        prepareCache();
    }

    public static JsonQueryFunction create(
            Expression json,
            Expression path,
            SqlJsonQueryWrapperBehavior wrapperBehavior,
            SqlJsonQueryEmptyOrErrorBehavior onEmpty,
            SqlJsonQueryEmptyOrErrorBehavior onError
    ) {
        final Expression[] operands = new Expression[]{json, path};

        return new JsonQueryFunction(operands, wrapperBehavior, onEmpty, onError);
    }

    @Override
    public int getClassId() {
        return JetSqlSerializerHook.JSON_QUERY;
    }

    @Override
    public HazelcastJsonValue eval(final Row row, final ExpressionEvalContext context) {
        // first evaluate the required parameter
        final String path = (String) operands[1].eval(row, context);
        validatePath(path);

        final Object operand0 = operands[0].eval(row, context);
        String json = operand0 instanceof HazelcastJsonValue
                ? operand0.toString()
                : (String) operand0;
        if (json == null) {
            json = "";
        }

        final JsonPath jsonPath = constantPathCache != null ? constantPathCache :
                pathCache.computeIfAbsent(path, COMPILE_FUNCTION);

        return wrap(execute(json, jsonPath));
    }

    private void prepareCache() {
        if (this.operands[1] instanceof ConstantExpression) {
            String path = (String) this.operands[1].eval(null, null);
            validatePath(path);
            this.constantPathCache = JsonPathUtil.compile(path);
        } else {
            this.pathCache = JsonPathUtil.makePathCache();
        }
    }

    private void validatePath(String path) {
        if (path == null) {
            throw QueryException.error("SQL/JSON path expression cannot be null");
        }
    }

    private String onErrorResponse(final Exception exception) {
        switch (onError) {
            case ERROR:
                // We deliberately don't use the cause here. The reason is that exceptions from ANTLR are not always
                // serializable, they can contain references to parser context and other objects, which are not.
                // That's why we also log the exception here.
                LOGGER.fine("JSON_QUERY failed", exception);
                throw QueryException.error("JSON_QUERY failed: " + exception);
            case EMPTY_ARRAY:
                return "[]";
            case EMPTY_OBJECT:
                return "{}";
            default:
            case NULL:
                return null;
        }
    }

    private String onEmptyResponse() {
        switch (onEmpty) {
            case ERROR:
                throw QueryException.error("JSON_QUERY evaluated to no value");
            case EMPTY_ARRAY:
                return "[]";
            case EMPTY_OBJECT:
                return "{}";
            case NULL:
            default:
                return null;
        }
    }

    private HazelcastJsonValue wrap(String json) {
        if (json == null) {
            return null;
        }
        return new HazelcastJsonValue(json);
    }

    private String execute(final String json, final JsonPath path) {
        Collection