/**
* 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.hadoop.hive.ql.hooks;
import com.google.common.collect.Lists;
import com.google.common.hash.Hasher;
import com.google.common.hash.Hashing;
import com.google.gson.stream.JsonWriter;
import org.apache.commons.collections.SetUtils;
import org.apache.commons.io.output.StringBuilderWriter;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hive.common.ObjectPair;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.metastore.api.FieldSchema;
import org.apache.hadoop.hive.metastore.api.Table;
import org.apache.hadoop.hive.ql.QueryPlan;
import org.apache.hadoop.hive.ql.exec.ColumnInfo;
import org.apache.hadoop.hive.ql.exec.SelectOperator;
import org.apache.hadoop.hive.ql.exec.TaskRunner;
import org.apache.hadoop.hive.ql.exec.Utilities;
import org.apache.hadoop.hive.ql.hooks.HookContext.HookType;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.BaseColumnInfo;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.Dependency;
import org.apache.hadoop.hive.ql.hooks.LineageInfo.Predicate;
import org.apache.hadoop.hive.ql.optimizer.lineage.LineageCtx.Index;
import org.apache.hadoop.hive.ql.plan.HiveOperation;
import org.apache.hadoop.hive.ql.session.SessionState;
import org.apache.hadoop.hive.ql.session.SessionState.LogHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
/**
* Implementation of a post execute hook that logs lineage info to a log file.
*/
public class LineageLogger implements ExecuteWithHookContext {
private static final Logger LOG = LoggerFactory.getLogger(LineageLogger.class);
private static final HashSet<String> OPERATION_NAMES = new HashSet<String>();
static {
OPERATION_NAMES.add(HiveOperation.QUERY.getOperationName());
OPERATION_NAMES.add(HiveOperation.CREATETABLE_AS_SELECT.getOperationName());
OPERATION_NAMES.add(HiveOperation.ALTERVIEW_AS.getOperationName());
OPERATION_NAMES.add(HiveOperation.CREATEVIEW.getOperationName());
}
private static final String FORMAT_VERSION = "1.0";
final static class Edge {
public static enum Type {
PROJECTION, PREDICATE
}
private Set<Vertex> sources;
private Set<Vertex> targets;
private String expr;
private Type type;
Edge(Set<Vertex> sources, Set<Vertex> targets, String expr, Type type) {
this.sources = sources;
this.targets = targets;
this.expr = expr;
this.type = type;
}
}
final static class Vertex {
public static enum Type {
COLUMN, TABLE
}
private Type type;
private String label;
private int id;
Vertex(String label) {
this(label, Type.COLUMN);
}
Vertex(String label, Type type) {
this.label = label;
this.type = type;
}
@Override
public int hashCode() {
return label.hashCode() + type.hashCode() * 3;
}
@Override
public boolean equals(Object obj) {
if (this == obj) {
return true;
}
if (!(obj instanceof Vertex)) {
return false;
}
Vertex vertex = (Vertex) obj;
return label.equals(vertex.label) && type == vertex.type;
}
}
@Override
public void run(HookContext hookContext) {
assert(hookContext.getHookType() == HookType.POST_EXEC_HOOK);
QueryPlan plan = hookContext.getQueryPlan();
Index index = hookContext.getIndex();
SessionState ss = SessionState.get();
if (ss != null && index != null
&& OPERATION_NAMES.contains(plan.getOperationName())
&& !plan.isExplain()) {
try {
StringBuilderWriter out = new StringBuilderWriter(1024);
JsonWriter writer = new JsonWriter(out);
String queryStr = plan.getQueryStr().trim();
writer.beginObject();
writer.name("version").value(FORMAT_VERSION);
HiveConf conf = ss.getConf();
boolean testMode = conf.getBoolVar(HiveConf.ConfVars.HIVE_IN_TEST);
if (!testMode) {
// Don't emit user/timestamp info in test mode,
// so that the test golden output file is fixed.
long queryTime = plan.getQueryStartTime().longValue();
if (queryTime == 0) queryTime = System.currentTimeMillis();
long duration = System.currentTimeMillis() - queryTime;
writer.name("user").value(hookContext.getUgi().getUserName());
writer.name("timestamp").value(queryTime/1000);
writer.name("duration").value(duration);
writer.name("jobIds");
writer.beginArray();
List<TaskRunner> tasks = hookContext.getCompleteTaskList();
if (tasks != null && !tasks.isEmpty()) {
for (TaskRunner task: tasks) {
String jobId = task.getTask().getJobID();
if (jobId != null) {
writer.value(jobId);
}
}
}
writer.endArray();
}
writer.name("engine").value(
HiveConf.getVar(conf, HiveConf.ConfVars.HIVE_EXECUTION_ENGINE));
writer.name("database").value(ss.getCurrentDatabase());
writer.name("hash").value(getQueryHash(queryStr));
writer.name("queryText").value(queryStr);
List<Edge> edges = getEdges(plan, index);
Set<Vertex> vertices = getVertices(edges);
writeEdges(writer, edges);
writeVertices(writer, vertices);
writer.endObject();
writer.close();
// Logger the lineage info
String lineage = out.toString();
if (testMode) {
// Logger to console
log(lineage);
} else {
// In non-test mode, emit to a log file,
// which can be different from the normal hive.log.
// For example, using NoDeleteRollingFileAppender to
// log to some file with different rolling policy.
LOG.info(lineage);
}
} catch (Throwable t) {
// Don't fail the query just because of any lineage issue.
log("Failed to log lineage graph, query is not affected\n"
+ org.apache.hadoop.util.StringUtils.stringifyException(t));
}
}
}
/**
* Logger an error to console if available.
*/
private void log(String error) {
LogHelper console = SessionState.getConsole();
if (console != null) {
console.printError(error);
}
}
/**
* Based on the final select operator, find out all the target columns.
* For each target column, find out its sources based on the dependency index.
*/
private List<Edge> getEdges(QueryPlan plan, Index index) {
LinkedHashMap<String, ObjectPair<SelectOperator,
org.apache.hadoop.hive.ql.metadata.Table>> finalSelOps = index.getFinalSelectOps();
Map<String, Vertex> vertexCache = new LinkedHashMap<String, Vertex>();
List<Edge> edges = new ArrayList<Edge>();
for (ObjectPair<SelectOperator,
org.apache.hadoop.hive.ql.metadata.Table> pair: finalSelOps.values()) {
List<FieldSchema> fieldSchemas = plan.getResultSchema().getFieldSchemas();
SelectOperator finalSelOp = pair.getFirst();
org.apache.hadoop.hive.ql.metadata.Table t = pair.getSecond();
String destTableName = null;
List<String> colNames = null;
if (t != null) {
destTableName = t.getDbName() + "." + t.getTableName();
fieldSchemas = t.getCols();
} else {
// Based on the plan outputs, find out the target table name and column names.
for (WriteEntity output : plan.getOutputs()) {
Entity.Type entityType = output.getType();
if (entityType == Entity.Type.TABLE
|| entityType == Entity.Type.PARTITION) {
t = output.getTable();
destTableName = t.getDbName() + "." + t.getTableName();
List<FieldSchema> cols = t.getCols();
if (cols != null && !cols.isEmpty()) {
colNames = Utilities.getColumnNamesFromFieldSchema(cols);
}
break;
}
}
}
Map<ColumnInfo, Dependency> colMap = index.getDependencies(finalSelOp);
List<Dependency> dependencies = colMap != null ? Lists.newArrayList(colMap.values()) : null;
int fields = fieldSchemas.size();
if (t != null && colMap != null && fields < colMap.size()) {
// Dynamic partition keys should be added to field schemas.
List<FieldSchema> partitionKeys = t.getPartitionKeys();
int dynamicKeyCount = colMap.size() - fields;
int keyOffset = partitionKeys.size() - dynamicKeyCount;
if (keyOffset >= 0) {
fields += dynamicKeyCount;
for (int i = 0; i < dynamicKeyCount; i++) {
FieldSchema field = partitionKeys.get(keyOffset + i);
fieldSchemas.add(field);
if (colNames != null) {
colNames.add(field.getName());
}
}
}
}
if (dependencies == null || dependencies.size() != fields) {
log("Result schema has " + fields
+ " fields, but we don't get as many dependencies");
} else {
// Go through each target column, generate the lineage edges.
Set<Vertex> targets = new LinkedHashSet<Vertex>();
for (int i = 0; i < fields; i++) {
Vertex target = getOrCreateVertex(vertexCache,
getTargetFieldName(i, destTableName, colNames, fieldSchemas),
Vertex.Type.COLUMN);
targets.add(target);
Dependency dep = dependencies.get(i);
addEdge(vertexCache, edges, dep.getBaseCols(), target,
dep.getExpr(), Edge.Type.PROJECTION);
}
Set<Predicate> conds = index.getPredicates(finalSelOp);
if (conds != null && !conds.isEmpty()) {
for (Predicate cond: conds) {
addEdge(vertexCache, edges, cond.getBaseCols(),
new LinkedHashSet<Vertex>(targets), cond.getExpr(),
Edge.Type.PREDICATE);
}
}
}
}
return edges;
}
private void addEdge(Map<String, Vertex> vertexCache, List<Edge> edges,
Set<BaseColumnInfo> srcCols, Vertex target, String expr, Edge.Type type) {
Set<Vertex> targets = new LinkedHashSet<Vertex>();
targets.add(target);
addEdge(vertexCache, edges, srcCols, targets, expr, type);
}
/**
* Find an edge from all edges that has the same source vertices.
* If found, add the more targets to this edge's target vertex list.
* Otherwise, create a new edge and add to edge list.
*/
private void addEdge(Map<String, Vertex> vertexCache, List<Edge> edges,
Set<BaseColumnInfo> srcCols, Set<Vertex> targets, String expr, Edge.Type type) {
Set<Vertex> sources = createSourceVertices(vertexCache, srcCols);
Edge edge = findSimilarEdgeBySources(edges, sources, expr, type);
if (edge == null) {
edges.add(new Edge(sources, targets, expr, type));
} else {
edge.targets.addAll(targets);
}
}
/**
* Convert a list of columns to a set of vertices.
* Use cached vertices if possible.
*/
private Set<Vertex> createSourceVertices(
Map<String, Vertex> vertexCache, Collection<BaseColumnInfo> baseCols) {
Set<Vertex> sources = new LinkedHashSet<Vertex>();
if (baseCols != null && !baseCols.isEmpty()) {
for(BaseColumnInfo col: baseCols) {
Table table = col.getTabAlias().getTable();
if (table.isTemporary()) {
// Ignore temporary tables
continue;
}
Vertex.Type type = Vertex.Type.TABLE;
String tableName = table.getDbName() + "." + table.getTableName();
FieldSchema fieldSchema = col.getColumn();
String label = tableName;
if (fieldSchema != null) {
type = Vertex.Type.COLUMN;
label = tableName + "." + fieldSchema.getName();
}
sources.add(getOrCreateVertex(vertexCache, label, type));
}
}
return sources;
}
/**
* Find a vertex from a cache, or create one if not.
*/
private Vertex getOrCreateVertex(
Map<String, Vertex> vertices, String label, Vertex.Type type) {
Vertex vertex = vertices.get(label);
if (vertex == null) {
vertex = new Vertex(label, type);
vertices.put(label, vertex);
}
return vertex;
}
/**
* Find an edge that has the same type, expression, and sources.
*/
private Edge findSimilarEdgeBySources(
List<Edge> edges, Set<Vertex> sources, String expr, Edge.Type type) {
for (Edge edge: edges) {
if (edge.type == type && StringUtils.equals(edge.expr, expr)
&& SetUtils.isEqualSet(edge.sources, sources)) {
return edge;
}
}
return null;
}
/**
* Generate normalized name for a given target column.
*/
private String getTargetFieldName(int fieldIndex,
String destTableName, List<String> colNames, List<FieldSchema> fieldSchemas) {
String fieldName = fieldSchemas.get(fieldIndex).getName();
String[] parts = fieldName.split("\\.");
if (destTableName != null) {
String colName = parts[parts.length - 1];
if (colNames != null && !colNames.contains(colName)) {
colName = colNames.get(fieldIndex);
}
return destTableName + "." + colName;
}
if (parts.length == 2 && parts[0].startsWith("_u")) {
return parts[1];
}
return fieldName;
}
/**
* Get all the vertices of all edges. Targets at first,
* then sources. Assign id to each vertex.
*/
private Set<Vertex> getVertices(List<Edge> edges) {
Set<Vertex> vertices = new LinkedHashSet<Vertex>();
for (Edge edge: edges) {
vertices.addAll(edge.targets);
}
for (Edge edge: edges) {
vertices.addAll(edge.sources);
}
// Assign ids to all vertices,
// targets at first, then sources.
int id = 0;
for (Vertex vertex: vertices) {
vertex.id = id++;
}
return vertices;
}
/**
* Write out an JSON array of edges.
*/
private void writeEdges(JsonWriter writer, List<Edge> edges) throws IOException {
writer.name("edges");
writer.beginArray();
for (Edge edge: edges) {
writer.beginObject();
writer.name("sources");
writer.beginArray();
for (Vertex vertex: edge.sources) {
writer.value(vertex.id);
}
writer.endArray();
writer.name("targets");
writer.beginArray();
for (Vertex vertex: edge.targets) {
writer.value(vertex.id);
}
writer.endArray();
if (edge.expr != null) {
writer.name("expression").value(edge.expr);
}
writer.name("edgeType").value(edge.type.name());
writer.endObject();
}
writer.endArray();
}
/**
* Write out an JSON array of vertices.
*/
private void writeVertices(JsonWriter writer, Set<Vertex> vertices) throws IOException {
writer.name("vertices");
writer.beginArray();
for (Vertex vertex: vertices) {
writer.beginObject();
writer.name("id").value(vertex.id);
writer.name("vertexType").value(vertex.type.name());
writer.name("vertexId").value(vertex.label);
writer.endObject();
}
writer.endArray();
}
/**
* Generate query string md5 hash.
*/
private String getQueryHash(String queryStr) {
Hasher hasher = Hashing.md5().newHasher();
hasher.putString(queryStr);
return hasher.hash().toString();
}
}