auditlog/src/main/scala/com/gerritforge/analytics/auditlog/spark/dataframe/ops/DataFrameOps.scala

// Copyright (C) 2018 GerritForge Ltd
//
// 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.gerritforge.analytics.auditlog.spark.dataframe.ops

import com.gerritforge.analytics.auditlog.broadcast.{AdditionalUsersInfo, GerritProjects, GerritUserIdentifiers}
import com.gerritforge.analytics.auditlog.spark.sql.udf.SparkExtractors.{extractCommandArgumentsUDF, extractCommandUDF, extractSubCommandUDF}
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.{udf, _}
import org.apache.spark.sql.{Column, DataFrame}

import scala.util.Try


object DataFrameOps {

  implicit class PimpedAuditLogDataFrame(dataFrame: DataFrame) {

    private def hasColumn(path: String) = Try(dataFrame(path)).isSuccess

    private def ifExistThenGetOrNull(column: String, targetColumn: => Column) = if (hasColumn(column)) targetColumn else lit(null)

    def hydrateWithUserIdentifierColumn(userIdentifierCol: String, gerritAccounts: GerritUserIdentifiers): DataFrame = {
      def extractIdentifier: UserDefinedFunction = udf((who: Int) => gerritAccounts.getIdentifier(who))

      dataFrame.withColumn(userIdentifierCol, ifExistThenGetOrNull("who", extractIdentifier(col("who"))))

    }

    def withTimeBucketColumn(timeBucketCol: String, timeAggregation: String): DataFrame = {
      dataFrame
        .withColumn(timeBucketCol, date_trunc(format=timeAggregation, from_unixtime(col("time_at_start").divide(1000))))
    }

    def withCommandColumns(commandCol: String, commandArgsCol: String): DataFrame = {
      dataFrame
        .withColumn(commandCol,
          extractCommandUDF(
            col("what"),
            col("access_path"),
            ifExistThenGetOrNull("http_method", col("http_method"))))
        .withColumn(commandArgsCol, extractCommandArgumentsUDF(col("what"), col("access_path")))
    }

    def withSubCommandColumns(subCommandCol: String): DataFrame = {
      dataFrame.withColumn(subCommandCol,
        extractSubCommandUDF(
          col("what"),
          col("access_path")
        )
      )
    }

    def withProjectColumn(projectCol: String, gerritProjects: GerritProjects): DataFrame = {
      def extractProjectUDF: UserDefinedFunction = udf((what: String, accessPath: String) => gerritProjects.extractProject(what, accessPath))

      dataFrame
        .withColumn(projectCol, extractProjectUDF(col("what"), col("access_path")))
    }

    def withUserTypeColumn(commandCol: String, additionalUsersInfo: AdditionalUsersInfo): DataFrame = {
      def extractUserType: UserDefinedFunction = udf((who: Int) => additionalUsersInfo.getUserType(who))

      dataFrame.withColumn(commandCol, ifExistThenGetOrNull("who", extractUserType(col("who"))))
    }

    def aggregateNumEventsColumn(numEventsCol: String, cols: List[String]): DataFrame = {
      dataFrame.groupBy(cols.map(c => col(c)): _*)
        .agg(count("*")
        .alias(numEventsCol))
    }
  }
}