/** * 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.zeppelin.jdbc; import java.io.IOException; import java.sql.SQLException; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import org.apache.commons.lang.StringUtils; import org.apache.zeppelin.completer.CompletionType; import org.apache.zeppelin.interpreter.thrift.InterpreterCompletion; import org.junit.Assert; import org.junit.Before; import org.junit.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.google.common.base.Joiner; import jline.console.completer.ArgumentCompleter; import static com.google.common.collect.Sets.newHashSet; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; /** * SQL completer unit tests */ public class SqlCompleterTest { public class CompleterTester { private SqlCompleter completer; private String buffer; private int fromCursor; private int toCursor; private Set<InterpreterCompletion> expectedCompletions; public CompleterTester(SqlCompleter completer) { this.completer = completer; } public CompleterTester buffer(String buffer) { this.buffer = buffer; return this; } public CompleterTester from(int fromCursor) { this.fromCursor = fromCursor; return this; } public CompleterTester to(int toCursor) { this.toCursor = toCursor; return this; } public CompleterTester expect(Set<InterpreterCompletion> expectedCompletions) { this.expectedCompletions = expectedCompletions; return this; } public void test() { for (int c = fromCursor; c <= toCursor; c++) { expectedCompletions(buffer, c, expectedCompletions); } } private void expectedCompletions(String buffer, int cursor, Set<InterpreterCompletion> expected) { if (StringUtils.isNotEmpty(buffer) && buffer.length() > cursor) { buffer = buffer.substring(0, cursor + 1); } List<InterpreterCompletion> candidates = new ArrayList<>(); completer.complete(buffer, cursor, candidates); String explain = explain(buffer, cursor, candidates); logger.info(explain); Assert.assertEquals("Buffer [" + buffer.replace(" ", ".") + "] and Cursor[" + cursor + "] " + explain, expected, newHashSet(candidates)); } private String explain(String buffer, int cursor, List<InterpreterCompletion> candidates) { List<String> cndidateStrings = new ArrayList<>(); for (InterpreterCompletion candidate : candidates) { cndidateStrings.add(candidate.getValue()); } StringBuffer sb = new StringBuffer(); for (int i = 0; i <= Math.max(cursor, buffer.length()); i++) { if (i == cursor) { sb.append("("); } if (i >= buffer.length()) { sb.append("_"); } else { if (Character.isWhitespace(buffer.charAt(i))) { sb.append("."); } else { sb.append(buffer.charAt(i)); } } if (i == cursor) { sb.append(")"); } } sb.append(" >> [").append(Joiner.on(",").join(cndidateStrings)).append("]"); return sb.toString(); } } private Logger logger = LoggerFactory.getLogger(SqlCompleterTest.class); private final static Set<String> EMPTY = new HashSet<>(); private CompleterTester tester; private ArgumentCompleter.WhitespaceArgumentDelimiter delimiter = new ArgumentCompleter.WhitespaceArgumentDelimiter(); private SqlCompleter sqlCompleter = new SqlCompleter(); @Before public void beforeTest() throws IOException, SQLException { Map<String, Set<String>> tables = new HashMap<>(); Map<String, Set<String>> columns = new HashMap<>(); Set<String> schemas = new HashSet<>(); Set<String> keywords = new HashSet<>(); keywords.add("SUM"); keywords.add("SUBSTRING"); keywords.add("SUBCLASS_ORIGIN"); keywords.add("ORDER"); keywords.add("SELECT"); keywords.add("LIMIT"); keywords.add("FROM"); schemas.add("prod_dds"); schemas.add("prod_emart"); Set<String> prod_dds_tables = new HashSet<>(); prod_dds_tables.add("financial_account"); prod_dds_tables.add("customer"); Set<String> prod_emart_tables = new HashSet<>(); prod_emart_tables.add("financial_account"); tables.put("prod_dds", prod_dds_tables); tables.put("prod_emart", prod_emart_tables); Set<String> prod_dds_financial_account_columns = new HashSet<>(); prod_dds_financial_account_columns.add("account_rk"); prod_dds_financial_account_columns.add("account_id"); Set<String> prod_dds_customer_columns = new HashSet<>(); prod_dds_customer_columns.add("customer_rk"); prod_dds_customer_columns.add("name"); prod_dds_customer_columns.add("birth_dt"); Set<String> prod_emart_financial_account_columns = new HashSet<>(); prod_emart_financial_account_columns.add("account_rk"); prod_emart_financial_account_columns.add("balance_amt"); columns.put("prod_dds.financial_account", prod_dds_financial_account_columns); columns.put("prod_dds.customer", prod_dds_customer_columns); columns.put("prod_emart.financial_account", prod_emart_financial_account_columns); sqlCompleter.init(schemas, tables, columns, keywords); tester = new CompleterTester(sqlCompleter); } @Test public void testFindAliasesInSQL_Simple() { String sql = "select * from prod_emart.financial_account a"; Map<String, String> res = sqlCompleter.findAliasesInSQL(delimiter.delimit(sql, 0).getArguments()); assertEquals(1, res.size()); assertTrue(res.get("a").equals("prod_emart.financial_account")); } @Test public void testFindAliasesInSQL_Two() { String sql = "select * from prod_dds.financial_account a, prod_dds.customer b"; Map<String, String> res = sqlCompleter.findAliasesInSQL(sqlCompleter.getSqlDelimiter().delimit(sql, 0).getArguments()); assertEquals(2, res.size()); assertTrue(res.get("a").equals("prod_dds.financial_account")); assertTrue(res.get("b").equals("prod_dds.customer")); } @Test public void testFindAliasesInSQL_WrongTables() { String sql = "select * from prod_ddsxx.financial_account a, prod_dds.customerxx b"; Map<String, String> res = sqlCompleter.findAliasesInSQL(sqlCompleter.getSqlDelimiter().delimit(sql, 0).getArguments()); assertEquals(0, res.size()); } @Test public void testCompleteName_Empty() { String buffer = ""; int cursor = 0; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); sqlCompleter.completeName(buffer, cursor, candidates, aliases, true); assertEquals(17, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("prod_dds", "prod_dds", CompletionType.schema.name()))); assertTrue(candidates.contains(new InterpreterCompletion("prod_emart", "prod_emart", CompletionType.schema.name()))); assertTrue(candidates.contains(new InterpreterCompletion("SUM", "SUM", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("SUBSTRING", "SUBSTRING", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("SUBCLASS_ORIGIN", "SUBCLASS_ORIGIN", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("SELECT", "SELECT", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("ORDER", "ORDER", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("LIMIT", "LIMIT", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("FROM", "FROM", CompletionType.keyword.name()))); assertTrue(candidates.contains(new InterpreterCompletion("financial_account", "financial_account", CompletionType.table.name()))); assertTrue(candidates.contains(new InterpreterCompletion("customer", "customer", CompletionType.table.name()))); assertTrue(candidates.contains(new InterpreterCompletion("account_id", "account_id", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("customer_rk", "customer_rk", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("account_rk", "account_rk", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("name", "name", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("birth_dt", "birth_dt", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("balance_amt", "balance_amt", CompletionType.column.name()))); } @Test public void testCompleteName_SimpleSchema() { String buffer = "prod_"; int cursor = 3; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); sqlCompleter.completeName(buffer, cursor, candidates, aliases, false); assertEquals(2, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("prod_dds", "prod_dds", CompletionType.schema.name()))); assertTrue(candidates.contains(new InterpreterCompletion("prod_emart", "prod_emart", CompletionType.schema.name()))); } @Test public void testCompleteName_SimpleTable() { String buffer = "prod_dds.fin"; int cursor = 11; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); sqlCompleter.completeName(buffer, cursor, candidates, aliases, false); assertEquals(1, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("financial_account", "financial_account", CompletionType.table.name()))); } @Test public void testCompleteName_SimpleColumn() { String buffer = "prod_dds.financial_account.acc"; int cursor = 30; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); sqlCompleter.completeName(buffer, cursor, candidates, aliases, true); assertEquals(2, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("account_rk", "account_rk", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("account_id", "account_id", CompletionType.column.name()))); } @Test public void testCompleteName_WithAlias() { String buffer = "a.acc"; int cursor = 4; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); aliases.put("a", "prod_dds.financial_account"); sqlCompleter.completeName(buffer, cursor, candidates, aliases, true); assertEquals(2, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("account_rk", "account_rk", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("account_id", "account_id", CompletionType.column.name()))); } @Test public void testCompleteName_WithAliasAndPoint() { String buffer = "a."; int cursor = 2; List<InterpreterCompletion> candidates = new ArrayList<>(); Map<String, String> aliases = new HashMap<>(); aliases.put("a", "prod_dds.financial_account"); sqlCompleter.completeName(buffer, cursor, candidates, aliases, true); assertEquals(2, candidates.size()); assertTrue(candidates.contains(new InterpreterCompletion("account_rk", "account_rk", CompletionType.column.name()))); assertTrue(candidates.contains(new InterpreterCompletion("account_id", "account_id", CompletionType.column.name()))); } @Test public void testSchemaAndTable() { String buffer = "select * from prod_emart.fi"; tester.buffer(buffer).from(19).to(23).expect(newHashSet(new InterpreterCompletion("prod_emart ", "prod_emart ", CompletionType.schema.name()))).test(); tester.buffer(buffer).from(25).to(27).expect(newHashSet(new InterpreterCompletion("financial_account ", "financial_account ", CompletionType.table.name()))).test(); } @Test public void testEdges() { String buffer = " ORDER "; tester.buffer(buffer).from(2).to(6).expect(newHashSet(new InterpreterCompletion("ORDER ", "ORDER ", CompletionType.keyword.name()))).test(); tester.buffer(buffer).from(0).to(1).expect(newHashSet( new InterpreterCompletion("ORDER", "ORDER", CompletionType.keyword.name()), new InterpreterCompletion("SUBCLASS_ORIGIN", "SUBCLASS_ORIGIN", CompletionType.keyword.name()), new InterpreterCompletion("SUBSTRING", "SUBSTRING", CompletionType.keyword.name()), new InterpreterCompletion("prod_emart", "prod_emart", CompletionType.schema.name()), new InterpreterCompletion("LIMIT", "LIMIT", CompletionType.keyword.name()), new InterpreterCompletion("SUM", "SUM", CompletionType.keyword.name()), new InterpreterCompletion("prod_dds", "prod_dds", CompletionType.schema.name()), new InterpreterCompletion("SELECT", "SELECT", CompletionType.keyword.name()), new InterpreterCompletion("FROM", "FROM", CompletionType.keyword.name()), new InterpreterCompletion("financial_account", "financial_account", CompletionType.table.name()), new InterpreterCompletion("customer", "customer", CompletionType.table.name()), new InterpreterCompletion("account_rk", "account_rk", CompletionType.column.name()), new InterpreterCompletion("account_id", "account_id", CompletionType.column.name()), new InterpreterCompletion("customer_rk", "customer_rk", CompletionType.column.name()), new InterpreterCompletion("name", "name", CompletionType.column.name()), new InterpreterCompletion("birth_dt", "birth_dt", CompletionType.column.name()), new InterpreterCompletion("balance_amt", "balance_amt", CompletionType.column.name()) )).test(); } @Test public void testMultipleWords() { String buffer = "SELE FRO LIM"; tester.buffer(buffer).from(1).to(3).expect(newHashSet(new InterpreterCompletion("SELECT ", "SELECT ", CompletionType.keyword.name()))).test(); tester.buffer(buffer).from(6).to(7).expect(newHashSet(new InterpreterCompletion("FROM ", "FROM ", CompletionType.keyword.name()))).test(); tester.buffer(buffer).from(9).to(12).expect(newHashSet(new InterpreterCompletion("LIMIT ", "LIMIT ", CompletionType.keyword.name()))).test(); } @Test public void testMultiLineBuffer() { String buffer = " \n SELE\nFRO"; tester.buffer(buffer).from(4).to(6).expect(newHashSet(new InterpreterCompletion("SELECT ", "SELECT ", CompletionType.keyword.name()))).test(); tester.buffer(buffer).from(9).to(11).expect(newHashSet(new InterpreterCompletion("FROM ", "FROM ", CompletionType.keyword.name()))).test(); } @Test public void testMultipleCompletionSuggestions() { String buffer = "SU"; tester.buffer(buffer).from(1).to(2).expect(newHashSet( new InterpreterCompletion("SUBCLASS_ORIGIN", "SUBCLASS_ORIGIN", CompletionType.keyword.name()), new InterpreterCompletion("SUM", "SUM", CompletionType.keyword.name()), new InterpreterCompletion("SUBSTRING", "SUBSTRING", CompletionType.keyword.name())) ).test(); } @Test public void testSqlDelimiterCharacters() { assertTrue(sqlCompleter.getSqlDelimiter().isDelimiterChar("r,", 1)); assertTrue(sqlCompleter.getSqlDelimiter().isDelimiterChar("SS,", 2)); assertTrue(sqlCompleter.getSqlDelimiter().isDelimiterChar(",", 0)); assertTrue(sqlCompleter.getSqlDelimiter().isDelimiterChar("ttt,", 3)); } }