repository/src/main/scala/org/apache/atlas/query/GremlinEvaluator.scala

/*
 * 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.atlas.query


import org.apache.atlas.query.Expressions._
import org.apache.atlas.query.TypeUtils.ResultWithPathStruct
import org.apache.atlas.repository.graphdb.AtlasGraph
import org.apache.atlas.typesystem.json._
import org.apache.atlas.typesystem.types.DataTypes.TypeCategory
import org.apache.atlas.typesystem.types._
import org.json4s._
import org.json4s.native.Serialization._

import scala.language.existentials

case class GremlinQueryResult(query: String,
                              resultDataType: IDataType[_],
                              rows: java.util.List[_]) {
    def toJson = JsonHelper.toJson(this)
}

class GremlinEvaluator(qry: GremlinQuery, persistenceStrategy: GraphPersistenceStrategies, g: AtlasGraph[_,_]) {

   /**
     *
     * @param gResultObj is the object returned from gremlin. This must be a List
     * @param qryResultObj is the object constructed for the output w/o the Path.
     * @return a ResultWithPathStruct
     */
    def addPathStruct(gResultObj: AnyRef, qryResultObj: Any): Any = {
        if (!qry.isPathExpression) {
            qryResultObj
        } else {
            import scala.collection.JavaConversions._
            import scala.collection.JavaConverters._
            
            val iPaths = gResultObj.asInstanceOf[java.util.List[AnyRef]].init

            val oPaths = iPaths.map { value =>
                persistenceStrategy.constructInstance(TypeSystem.getInstance().getIdType.getStructType, value)
            }.toList.asJava
            val sType = qry.expr.dataType.asInstanceOf[StructType]
            val sInstance = sType.createInstance()
            sInstance.set(ResultWithPathStruct.pathAttrName, oPaths)
            sInstance.set(ResultWithPathStruct.resultAttrName, qryResultObj)
            sInstance
        }
    }

    def instanceObject(v: AnyRef): AnyRef = {
        if (qry.isPathExpression) {
            import scala.collection.JavaConversions._
            v.asInstanceOf[java.util.List[AnyRef]].last
        } else {
            v
        }
    }

    def evaluate(): GremlinQueryResult = {
        import scala.collection.JavaConversions._
         val debug:Boolean = false
        val rType = qry.expr.dataType
        val oType = if (qry.isPathExpression) {
            qry.expr.children(0).dataType
        }
        else {
            rType
        }
        val rawRes = g.executeGremlinScript(qry.queryStr, qry.isPathExpression);
        if(debug) {
            println(" rawRes " +rawRes)
        }
        if (!qry.hasSelectList && ! qry.isGroupBy) {
            val rows = rawRes.asInstanceOf[java.util.List[AnyRef]].map { v =>
                val instObj = instanceObject(v)
                val o = persistenceStrategy.constructInstance(oType, instObj)
                addPathStruct(v, o)
            }
            GremlinQueryResult(qry.expr.toString, rType, rows.toList)
        } else {
            val sType = oType.asInstanceOf[StructType]
            val rows = rawRes.asInstanceOf[java.util.List[AnyRef]].map { r =>
                val rV = instanceObject(r)
                val sInstance = sType.createInstance()
                val selObj = SelectExpressionHelper.extractSelectExpression(qry.expr)
                if (selObj.isDefined) {
                    val selExpr = selObj.get.asInstanceOf[Expressions.SelectExpression]
                    selExpr.selectListWithAlias.foreach { aE =>
                        val cName = aE.alias
                        val (src, idx) = qry.resultMaping(cName)
                        val v = getColumnValue(rV, src, idx)
                        //if select clause is selecting the entire object then return only the instance id (guid, version, state and typeName)
                        if (aE.dataType.getTypeCategory == TypeCategory.CLASS) {
                             sInstance.set(cName, persistenceStrategy.constructClassInstanceId(aE.dataType.asInstanceOf[ClassType], v))
                        } else {
                            sInstance.set(cName, persistenceStrategy.constructInstance(aE.dataType, v))
                        }
                    }
                }
                else if(qry.isGroupBy) {
                    //the order in the result will always match the order in the select list
                    val selExpr = qry.expr.asInstanceOf[GroupByExpression].selExpr
                    var idx = 0;
                    val row : java.util.List[Object] = rV.asInstanceOf[java.util.List[Object]]
                    selExpr.selectListWithAlias.foreach { aE =>
                        val cName = aE.alias
                        val cValue = row.get(idx);

                        sInstance.set(cName, persistenceStrategy.constructInstance(aE.dataType, cValue))
                        idx += 1;
                    }
                }
                addPathStruct(r, sInstance)
            }
            GremlinQueryResult(qry.expr.toString, rType, rows.toList)
        }

    }
    
    private def getColumnValue(rowValue: AnyRef, colName: String, idx: Integer) : AnyRef  = {

        var rawColumnValue: AnyRef = null;
        if(rowValue.isInstanceOf[java.util.Map[_,_]]) {
            val columnsMap = rowValue.asInstanceOf[java.util.Map[String,AnyRef]];
            rawColumnValue = columnsMap.get(colName);
        }
        else {
            //when there is only one column, result does not come back as a map
            rawColumnValue = rowValue;
        }

        var value : AnyRef = null;
        if(rawColumnValue.isInstanceOf[java.util.List[_]] && idx >= 0) {
            val arr = rawColumnValue.asInstanceOf[java.util.List[AnyRef]];
            value = arr.get(idx);
        }
        else {
            value = rawColumnValue;
        }

        return value;
    }
}

object JsonHelper {

    class GremlinQueryResultSerializer()
        extends Serializer[GremlinQueryResult] {
        def deserialize(implicit format: Formats) = {
            throw new UnsupportedOperationException("Deserialization of GremlinQueryResult not supported")
        }

        def serialize(implicit f: Formats) = {
            case GremlinQueryResult(query, rT, rows) =>
                JObject(JField("query", JString(query)),
                    JField("dataType", TypesSerialization.toJsonValue(rT)(f)),
                    JField("rows", Extraction.decompose(rows)(f))
                )
        }
    }

    implicit val formats = org.json4s.native.Serialization.formats(NoTypeHints) + new TypedStructSerializer +
        new TypedReferenceableInstanceSerializer + new BigDecimalSerializer + new BigIntegerSerializer +
        new GremlinQueryResultSerializer

    def toJson(r: GremlinQueryResult): String = {
        writePretty(r)
    }
}