#445

Committed 24 Jul 2025 02:27PM UTC coverage: 69.257% (-7.1%) from 76.395%

Build # #445

Build Type

push

Committed by

Commit Message

Merge branch 'hive-3.x' of github.com:ExpediaGroup/waggle-dance into hive-3.x

Run Details

2620 of 3783 relevant lines covered (69.26%)

0.69 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.73

/waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/mapping/model/ASTQueryMapping.java

/**
 * Copyright (C) 2016-2025 Expedia, Inc.
 *
 * 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.hotels.bdp.waggledance.mapping.model;

import static org.apache.hadoop.hive.ql.parse.BaseSemanticAnalyzer.unescapeIdentifier;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.SortedSet;
import java.util.Stack;
import java.util.TreeSet;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.antlr.runtime.CommonToken;
import org.antlr.runtime.tree.Tree;
import org.apache.hadoop.hive.ql.lib.Node;
import org.apache.hadoop.hive.ql.parse.ASTNode;
import org.apache.hadoop.hive.ql.parse.HiveParser;
import org.apache.hadoop.hive.ql.parse.ParseException;
import org.apache.hadoop.hive.ql.parse.ParseUtils;

import com.hotels.bdp.waggledance.api.WaggleDanceException;

public enum ASTQueryMapping implements QueryMapping {

  INSTANCE;

  private static final String PRESTO_VIEW_MARKER = "/* Presto View";
  private final static String RE_WORD_BOUNDARY = "\\b";
  private final static Comparator<CommonToken> ON_START_INDEX = Comparator.comparingInt(CommonToken::getStartIndex);

  @Override
  public String transformOutboundDatabaseName(MetaStoreMapping metaStoreMapping, String query) {
    if (hasNonHiveViewMarker(query)) {
      // skipping queries that are not "Hive" view queries. We can't parse those.
      return query;
    }
    ASTNode root;
    try {
      root = ParseUtils.parse(query);
    } catch (ParseException e) {
      throw new WaggleDanceException("Can't parse query: '" + query + "'", e);
    }

    StringBuilder result = transformDatabaseTableTokens(metaStoreMapping, root, query);
    transformFunctionTokens(metaStoreMapping, root, result);
    return result.toString();
  }

  private boolean hasNonHiveViewMarker(String query) {
    if (query != null && query.trim().startsWith(PRESTO_VIEW_MARKER)) {
      return true;
    }
    return false;
  }

  private StringBuilder transformDatabaseTableTokens(MetaStoreMapping metaStoreMapping, ASTNode root, String query) {
    StringBuilder result = new StringBuilder();
    SortedSet<CommonToken> dbNameTokens = extractDbNameTokens(root);
    int startIndex = 0;
    for (CommonToken dbNameNode : dbNameTokens) {
      final String dbName = dbNameNode.getText();
      final boolean escaped = dbName.startsWith("`") && dbName.endsWith("`");
      String transformedDbName = metaStoreMapping.transformOutboundDatabaseName(unescapeIdentifier(dbName));
      if (escaped) {
        transformedDbName = "`" + transformedDbName + "`";
      }
      result.append(query, startIndex, dbNameNode.getStartIndex());
      result.append(transformedDbName);
      startIndex = dbNameNode.getStopIndex() + 1;
    }
    result.append(query.substring(startIndex));
    return result;
  }

  private void transformFunctionTokens(MetaStoreMapping metaStoreMapping, ASTNode root, StringBuilder result) {
    // Done differently from the extractDbNameTokens as the Function tokens do not contain a correct start index. We'll
    // have to fall back to search and replace.
    List<CommonToken> functionTokens = extractFunctionTokens(root);
    for (CommonToken functionNode : functionTokens) {
      String functionName = functionNode.getText();
      // No dbname, no need to replace.
      if (!functionName.contains(".")) {
        continue;
      }
      if (functionName.startsWith("`") && functionName.endsWith("`")) {
        functionName = functionName.replaceAll("`", "");
      }
      Pattern pattern = Pattern.compile(RE_WORD_BOUNDARY + functionName + RE_WORD_BOUNDARY);
      Matcher matcher = pattern.matcher(result);
      if (matcher.find()) {
        int index = matcher.start();
        String prefix = metaStoreMapping.getDatabasePrefix();
        result.replace(index, index, prefix);
      }
    }
  }

  private SortedSet<CommonToken> extractDbNameTokens(ASTNode root) {
    SortedSet<CommonToken> dbNameTokens = new TreeSet<>(ON_START_INDEX);
    Stack<ASTNode> stack = new Stack<>();
    stack.push(root);
    while (!stack.isEmpty()) {
      ASTNode current = stack.pop();
      for (ASTNode child : getChildren(current)) {
        stack.push(child);
      }
      if (current.getType() == HiveParser.TOK_TABNAME) {
        if (current.getChildCount() == 2 && childrenAreIdentifiers(current)) {
          // First child of TOK_TABNAME node is the database name node
          CommonToken dbNameNode = (CommonToken) ((ASTNode) current.getChild(0)).getToken();
          dbNameTokens.add(dbNameNode);
        }
        // Otherwise TOK_TABNAME node only has one child which contains just the table name
      }
    }
    return dbNameTokens;
  }

  private List<CommonToken> extractFunctionTokens(ASTNode root) {
    List<CommonToken> tokens = new ArrayList<>();
    Stack<ASTNode> stack = new Stack<>();
    stack.push(root);

    while (!stack.isEmpty()) {
      ASTNode current = stack.pop();
      for (ASTNode child : getChildren(current)) {
        stack.push(child);
      }
      if (current.getType() == HiveParser.TOK_FUNCTION) {
        if (current.getChildCount() < 1) {
          continue;
        }
        Tree child = current.getChild(0);
        if (child.getType() == HiveParser.Identifier) {
          CommonToken dbNameDotFunctionNameNode = (CommonToken) ((ASTNode)child).getToken();
          tokens.add(dbNameDotFunctionNameNode);
        }
      }
    }
    return tokens;
  }

  private boolean childrenAreIdentifiers(ASTNode current) {
    for (ASTNode child : getChildren(current)) {
      if (child.getType() != HiveParser.Identifier) {
        return false;
      }
    }
    return true;
  }

  private static List<ASTNode> getChildren(ASTNode pt) {
    List<ASTNode> rt = new ArrayList<>();
    List<Node> children = pt.getChildren();
    if (children != null) {
      for (Node child : pt.getChildren()) {
        rt.add((ASTNode) child);
      }
    }
    return rt;
  }

}

1	/**
2	* Copyright (C) 2016-2025 Expedia, Inc.
3	*
4	* Licensed under the Apache License, Version 2.0 (the "License");
5	* you may not use this file except in compliance with the License.
6	* You may obtain a copy of the License at
7	*
8	* http://www.apache.org/licenses/LICENSE-2.0
9	*
10	* Unless required by applicable law or agreed to in writing, software
11	* distributed under the License is distributed on an "AS IS" BASIS,
12	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	* See the License for the specific language governing permissions and
14	* limitations under the License.
15	*/
16	package com.hotels.bdp.waggledance.mapping.model;
17
18	import static org.apache.hadoop.hive.ql.parse.BaseSemanticAnalyzer.unescapeIdentifier;
19
20	import java.util.ArrayList;
21	import java.util.Comparator;
22	import java.util.List;
23	import java.util.SortedSet;
24	import java.util.Stack;
25	import java.util.TreeSet;
26	import java.util.regex.Matcher;
27	import java.util.regex.Pattern;
28
29	import org.antlr.runtime.CommonToken;
30	import org.antlr.runtime.tree.Tree;
31	import org.apache.hadoop.hive.ql.lib.Node;
32	import org.apache.hadoop.hive.ql.parse.ASTNode;
33	import org.apache.hadoop.hive.ql.parse.HiveParser;
34	import org.apache.hadoop.hive.ql.parse.ParseException;
35	import org.apache.hadoop.hive.ql.parse.ParseUtils;
36
37	import com.hotels.bdp.waggledance.api.WaggleDanceException;
38
39	public enum ASTQueryMapping implements QueryMapping {	1✔
40
41	INSTANCE;	1✔
42
43	private static final String PRESTO_VIEW_MARKER = "/* Presto View";
44	private final static String RE_WORD_BOUNDARY = "\\b";
45	private final static Comparator<CommonToken> ON_START_INDEX = Comparator.comparingInt(CommonToken::getStartIndex);	1✔
46
47	@Override
48	public String transformOutboundDatabaseName(MetaStoreMapping metaStoreMapping, String query) {
49	if (hasNonHiveViewMarker(query)) {	1✔
50	// skipping queries that are not "Hive" view queries. We can't parse those.
51	return query;	1✔
52	}
53	ASTNode root;
54	try {
55	root = ParseUtils.parse(query);	1✔
56	} catch (ParseException e) {	1✔
57	throw new WaggleDanceException("Can't parse query: '" + query + "'", e);	1✔
58	}	1✔
59
60	StringBuilder result = transformDatabaseTableTokens(metaStoreMapping, root, query);	1✔
61	transformFunctionTokens(metaStoreMapping, root, result);	1✔
62	return result.toString();	1✔
63	}
64
65	private boolean hasNonHiveViewMarker(String query) {
66	if (query != null && query.trim().startsWith(PRESTO_VIEW_MARKER)) {	1✔
67	return true;	1✔
68	}
69	return false;	1✔
70	}
71
72	private StringBuilder transformDatabaseTableTokens(MetaStoreMapping metaStoreMapping, ASTNode root, String query) {
73	StringBuilder result = new StringBuilder();	1✔
74	SortedSet<CommonToken> dbNameTokens = extractDbNameTokens(root);	1✔
75	int startIndex = 0;	1✔
76	for (CommonToken dbNameNode : dbNameTokens) {	1✔
77	final String dbName = dbNameNode.getText();	1✔
78	final boolean escaped = dbName.startsWith("`") && dbName.endsWith("`");	1✔
79	String transformedDbName = metaStoreMapping.transformOutboundDatabaseName(unescapeIdentifier(dbName));	1✔
80	if (escaped) {	1✔
81	transformedDbName = "`" + transformedDbName + "`";	1✔
82	}
83	result.append(query, startIndex, dbNameNode.getStartIndex());	1✔
84	result.append(transformedDbName);	1✔
85	startIndex = dbNameNode.getStopIndex() + 1;	1✔
86	}	1✔
87	result.append(query.substring(startIndex));	1✔
88	return result;	1✔
89	}
90
91	private void transformFunctionTokens(MetaStoreMapping metaStoreMapping, ASTNode root, StringBuilder result) {
92	// Done differently from the extractDbNameTokens as the Function tokens do not contain a correct start index. We'll
93	// have to fall back to search and replace.
94	List<CommonToken> functionTokens = extractFunctionTokens(root);	1✔
95	for (CommonToken functionNode : functionTokens) {	1✔
96	String functionName = functionNode.getText();	1✔
97	// No dbname, no need to replace.
98	if (!functionName.contains(".")) {	1✔
99	continue;	1✔
100	}
101	if (functionName.startsWith("`") && functionName.endsWith("`")) {	1✔
102	functionName = functionName.replaceAll("`", "");	1✔
103	}
104	Pattern pattern = Pattern.compile(RE_WORD_BOUNDARY + functionName + RE_WORD_BOUNDARY);	1✔
105	Matcher matcher = pattern.matcher(result);	1✔
106	if (matcher.find()) {	1✔
107	int index = matcher.start();	1✔
108	String prefix = metaStoreMapping.getDatabasePrefix();	1✔
109	result.replace(index, index, prefix);	1✔
110	}
111	}	1✔
112	}	1✔
113
114	private SortedSet<CommonToken> extractDbNameTokens(ASTNode root) {
115	SortedSet<CommonToken> dbNameTokens = new TreeSet<>(ON_START_INDEX);	1✔
116	Stack<ASTNode> stack = new Stack<>();	1✔
117	stack.push(root);	1✔
118	while (!stack.isEmpty()) {	1✔
119	ASTNode current = stack.pop();	1✔
120	for (ASTNode child : getChildren(current)) {	1✔
121	stack.push(child);	1✔
122	}	1✔
123	if (current.getType() == HiveParser.TOK_TABNAME) {	1✔
124	if (current.getChildCount() == 2 && childrenAreIdentifiers(current)) {	1✔
125	// First child of TOK_TABNAME node is the database name node
126	CommonToken dbNameNode = (CommonToken) ((ASTNode) current.getChild(0)).getToken();	1✔
127	dbNameTokens.add(dbNameNode);	1✔
128	}
129	// Otherwise TOK_TABNAME node only has one child which contains just the table name
130	}
131	}	1✔
132	return dbNameTokens;	1✔
133	}
134
135	private List<CommonToken> extractFunctionTokens(ASTNode root) {
136	List<CommonToken> tokens = new ArrayList<>();	1✔
137	Stack<ASTNode> stack = new Stack<>();	1✔
138	stack.push(root);	1✔
139
140	while (!stack.isEmpty()) {	1✔
141	ASTNode current = stack.pop();	1✔
142	for (ASTNode child : getChildren(current)) {	1✔
143	stack.push(child);	1✔
144	}	1✔
145	if (current.getType() == HiveParser.TOK_FUNCTION) {	1✔
146	if (current.getChildCount() < 1) {	1✔
147	continue;	×
148	}
149	Tree child = current.getChild(0);	1✔
150	if (child.getType() == HiveParser.Identifier) {	1✔
151	CommonToken dbNameDotFunctionNameNode = (CommonToken) ((ASTNode)child).getToken();	1✔
152	tokens.add(dbNameDotFunctionNameNode);	1✔
153	}
154	}
155	}	1✔
156	return tokens;	1✔
157	}
158
159	private boolean childrenAreIdentifiers(ASTNode current) {
160	for (ASTNode child : getChildren(current)) {	1✔
161	if (child.getType() != HiveParser.Identifier) {	1✔
162	return false;	×
163	}
164	}	1✔
165	return true;	1✔
166	}
167
168	private static List<ASTNode> getChildren(ASTNode pt) {
169	List<ASTNode> rt = new ArrayList<>();	1✔
170	List<Node> children = pt.getChildren();	1✔
171	if (children != null) {	1✔
172	for (Node child : pt.getChildren()) {	1✔
173	rt.add((ASTNode) child);	1✔
174	}	1✔
175	}
176	return rt;	1✔
177	}
178
179	}

HotelsDotCom / waggle-dance / #445

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous