Java Code Examples for org.apache.calcite.sql.parser.SqlParserUtil#StringAndPos

public void assertExceptionIsThrown(String sql, String expectedMsgPattern) {
  final SqlValidator validator;
  final SqlNode sqlNode;
  final SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
  try {
    sqlNode = parseQuery(sap.sql);
    validator = getValidator();
  } catch (Throwable e) {
    checkParseEx(e, expectedMsgPattern, sap.sql);
    return;
  }

  Throwable thrown = null;
  try {
    validator.validate(sqlNode);
  } catch (Throwable ex) {
    thrown = ex;
  }

  SqlTests.checkEx(thrown, expectedMsgPattern, sap, SqlTests.Stage.VALIDATE);
}
 

@Test
public void testSuggestFromSchemaSeparator() throws Exception {
  final String partialQuery = "Select * from mysrc.^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, new ArrayList<String>());

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(3, returnedSuggestions.getSuggestions().size());
  for (int i = 0; i < 3; i++) {
    SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(i);
    assertEquals("TABLE", suggestion.getType());
    assertEquals(String.format("mysrc.ds%d", i + 1), suggestion.getName());
  }
}
 

/**
 * Tests that a given SQL which may be invalid or incomplete simplifies
 * itself and yields the salesTables set of completion hints. This is an
 * integration test of {@link #assertHint} and {@link #assertSimplify}.
 *
 * @param sql             SQL statement
 * @param expectedResults Expected list of hints
 * @param expectedWord    Word that we expect to be replaced, or null if we
 *                        don't care
 */
protected void assertComplete(
    String sql,
    String expectedResults,
    String expectedWord,
    Map<String, String> replacements) {
  SqlAdvisor advisor = tester.getFactory().createAdvisor();

  SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
  final String[] replaced = {null};
  List<SqlMoniker> results =
      advisor.getCompletionHints(sap.sql, sap.cursor, replaced);
  Assertions.assertEquals(expectedResults, convertCompletionHints(results),
      () -> "Completion hints for " + sql);
  if (expectedWord != null) {
    Assertions.assertEquals(expectedWord, replaced[0], "replaced[0] for " + sql);
  } else {
    assertNotNull(replaced[0]);
  }
  assertReplacements(sql, replacements, advisor, replaced[0], results);
}
 

@Test
public void testPartialColumnCompletionWithAlias() throws Exception {
  final String partialQuery = "SELECT t1.s^ from testSpace.supplier t1";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  ArrayList<String> expectedColumns = Lists.newArrayList();
  expectedColumns.addAll(
    asList(
      "s_name",
      "s_phone"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(2, returnedSuggestions.getSuggestions().size());
  for (int i = 0; i < 2; i++) {
    SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(i);
    if (suggestion.getType().equals("COLUMN")) {
      assertTrue(expectedColumns.contains(suggestion.getName()));
    }
  }
}
 

public void checkFails(
	String sql,
	String expectedMsgPattern) {
	SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
	Throwable thrown = null;
	try {
		final SqlNode sqlNode;
		sqlNode = factory.createParser(sap.sql).parseStmt();
		Util.discard(sqlNode);
	} catch (Throwable ex) {
		thrown = ex;
	}

	checkEx(expectedMsgPattern, sap, thrown);
}
 

public void checkFails(
	String sql,
	String expectedMsgPattern) {
	SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
	Throwable thrown = null;
	try {
		final SqlNode sqlNode;
		sqlNode = getSqlParser(sap.sql).parseStmt();
		Util.discard(sqlNode);
	} catch (Throwable ex) {
		thrown = ex;
	}

	checkEx(expectedMsgPattern, sap, thrown);
}
 

@Test // Could improve to suggest Dremio specific keywords
public void testSuggestSelectList() throws Exception {
  final String partialQuery = "Select ^ from mysrc.ds1";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  for (int i = 0; i < returnedSuggestions.getSuggestions().size(); i++) {
    assertEquals("KEYWORD", returnedSuggestions.getSuggestions().get(i).getType());
  }
}
 

@Test
public void testErrorUnrecognizedColumn() throws Exception {
  final String partialQuery = "SELECT testCol^ FROM testSpace.supplier";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final ValidationResponse returnedSuggestions = testValidateSQL(stringAndPos.sql, stringAndPos.cursor, asList("cp"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getErrors());
  QueryError firstError =  returnedSuggestions.getErrors().get(0);
  assertEquals("Column 'TESTCOL' not found in any table", firstError.getMessage());
  assertEquals(new QueryError.Range(1,8,2,15), firstError.getRange());
}
 

private void adviseSql(int apiVersion, String sql, Consumer<ResultSet> checker)
    throws SQLException {
  Properties info = new Properties();
  if (apiVersion == 1) {
    info.put("lex", "JAVA");
    info.put("quoting", "DOUBLE_QUOTE");
  } else if (apiVersion == 2) {
    info.put("lex", "SQL_SERVER");
    info.put("quoting", "BRACKET");
  }
  Connection connection =
      DriverManager.getConnection("jdbc:calcite:", info);
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("hr", new ReflectiveSchema(new JdbcTest.HrSchema()));
  SchemaPlus schema = rootSchema.add("s", new AbstractSchema());
  calciteConnection.setSchema("hr");
  final TableFunction getHints =
      apiVersion == 1 ? new SqlAdvisorGetHintsFunction() : new SqlAdvisorGetHintsFunction2();
  schema.add("get_hints", getHints);
  String getHintsSql;
  if (apiVersion == 1) {
    getHintsSql = "select id, names, type from table(\"s\".\"get_hints\"(?, ?)) as t";
  } else {
    getHintsSql = "select id, names, type, replacement from table([s].[get_hints](?, ?)) as t";
  }

  PreparedStatement ps = connection.prepareStatement(getHintsSql);
  SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
  ps.setString(1, sap.sql);
  ps.setInt(2, sap.cursor);
  final ResultSet resultSet = ps.executeQuery();
  checker.accept(resultSet);
  resultSet.close();
  connection.close();
}
 

@Test
public void testSuggestFromPartialSchema() throws Exception {
  final String partialQuery = "Select * from m^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  ArrayList<String> expectedTables = Lists.newArrayList();
  expectedTables.addAll(
    asList(
      "mysrc.ds1",
      "mysrc.ds2",
      "mysrc.ds3",
      "sys.materializations",
      "sys.memory",
      "\"sys.cache\".mount_points"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals((expectedTables.size() + 1), returnedSuggestions.getSuggestions().size());
  for (int i = 0; i < (expectedTables.size() + 1); i++) {
    SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(i);
    if (suggestion.getType().equals("TABLE")) {
      assertTrue(expectedTables.contains(suggestion.getName()));
    } else if (suggestion.getType().equals("SCHEMA")) {
      assertEquals("mysrc", suggestion.getName());
    }
  }
}
 

@Test
public void testSQLAnalyzeSuggestInfoSchema() throws Exception {
  final String partialQuery = "select * from i^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("@dremio"));

  ArrayList<String> expectedTables = Lists.newArrayList();
  expectedTables.addAll(
    asList(
    "INFORMATION_SCHEMA.CATALOGS",
    "INFORMATION_SCHEMA.COLUMNS",
    "INFORMATION_SCHEMA.SCHEMATA",
    "INFORMATION_SCHEMA.\"TABLES\"",
    "INFORMATION_SCHEMA.VIEWS"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(6, returnedSuggestions.getSuggestions().size());

  for (SuggestionResponse.Suggestion suggestion : returnedSuggestions.getSuggestions()) {
    if (suggestion.getType().equals("TABLE")) {
      assertTrue(expectedTables.contains(suggestion.getName()));
    } else if (suggestion.getType().equals("SCHEMA")) {
      assertEquals("INFORMATION_SCHEMA", suggestion.getName());
    }
  }
}
 

@Test // Error message identical to current. (unrecognized * intead of missing keyword FROM) Can be improved.
public void testErrorIncompleteFrom() throws Exception {
  final String partialQuery = "Select * fro^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final ValidationResponse returnedSuggestions = testValidateSQL(stringAndPos.sql, stringAndPos.cursor, asList("@dremio"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getErrors());
  QueryError firstError =  returnedSuggestions.getErrors().get(0);
  assertEquals("Unknown identifier '*'", firstError.getMessage());
  assertEquals(new QueryError.Range(1,8,2,9), firstError.getRange());
}
 

@Test
public void testSpaceVDSPartialDatasetCompletion() throws Exception {
  final String partialQuery = "SELECT * from testSpace.sup^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(1, returnedSuggestions.getSuggestions().size());
  SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(0);
  assertEquals("TABLE", suggestion.getType());
  assertEquals("testSpace.supplier", suggestion.getName());
}
 

@Test
public void testPDSFullDatasetCompletionWPeriod() throws Exception {
  final String partialQuery = "SELECT * from cp.^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(1, returnedSuggestions.getSuggestions().size());
  SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(0);
  assertEquals("TABLE", suggestion.getType());
  assertEquals("cp.\"tpch/supplier.parquet\"", suggestion.getName());
}
 

@Test
public void testSpaceVDSFullDatasetCompletionWPeriod() throws Exception {
  final String partialQuery = "SELECT * from testSpace.^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals(1, returnedSuggestions.getSuggestions().size());
  SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(0);
  assertEquals("TABLE", suggestion.getType());
  assertEquals("testSpace.supplier", suggestion.getName());
}
 

@Test
public void spacePDSPartialSchemaCompletion() throws Exception {
  final String partialQuery = "SELECT * from c^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final SuggestionResponse returnedSuggestions = testSuggestSQL(stringAndPos.sql, stringAndPos.cursor, asList("mysrc"));

  ArrayList<String> expectedTables = Lists.newArrayList();
  expectedTables.addAll(
    asList(
      "INFORMATION_SCHEMA.CATALOGS",
      "INFORMATION_SCHEMA.COLUMNS",
      "cp.\"tpch/supplier.parquet\""
      ));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getSuggestions());
  assertEquals((expectedTables.size() + 1), returnedSuggestions.getSuggestions().size());
  for (int i = 0; i < (expectedTables.size() + 1); i++) {
    SuggestionResponse.Suggestion suggestion = returnedSuggestions.getSuggestions().get(i);
    if (suggestion.getType().equals("TABLE")) {
      assertTrue(expectedTables.contains(suggestion.getName()));
    } else if (suggestion.getType().equals("SCHEMA")) {
      assertEquals("cp", suggestion.getName());
    }
  }
}
 

@Test // Current error-handling wraps this error in a generic parse error.
public void testErrorIncompleteSelect() throws Exception {
  final String partialQuery = "Sel^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final ValidationResponse returnedSuggestions = testValidateSQL(stringAndPos.sql, stringAndPos.cursor, asList("@dremio"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getErrors());
  QueryError firstError =  returnedSuggestions.getErrors().get(0);
  assertEquals("Non-query expression encountered in illegal context", firstError.getMessage());
  assertEquals(new QueryError.Range(1,1,2,4), firstError.getRange());
}
 

@Test // Range can be improved
public void testErrorUnrecognizedTable() throws Exception {
  final String partialQuery = "Select * from m^";
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(partialQuery);
  final ValidationResponse returnedSuggestions = testValidateSQL(stringAndPos.sql, stringAndPos.cursor, asList("@dremio"));

  logAdvisorResponse(returnedSuggestions);
  assertNotNull(returnedSuggestions);
  assertNotNull(returnedSuggestions.getErrors());
  QueryError firstError =  returnedSuggestions.getErrors().get(0);
  assertEquals("Table 'M' not found", firstError.getMessage());
  assertEquals(new QueryError.Range(1,15,2,16), firstError.getRange());
}
 

@Test
public void testSuggestion() {
  final SqlParserUtil.StringAndPos stringAndPos = SqlParserUtil.findPos(sql);
  List<SqlMoniker> suggestions = sqlAnalyzer.suggest(stringAndPos.sql, stringAndPos.cursor);
  assertEquals(expectedSuggestionCount, suggestions.size());
  if (checkSuggestions) {
    assertSuggestions(suggestions);
  }
}
 

/**
 * Tests that a given SQL statement simplifies to the salesTables result.
 *
 * @param sql      SQL statement to simplify. The SQL statement must contain
 *                 precisely one caret '^', which marks the location where
 *                 completion is to occur.
 * @param expected Expected result after simplification.
 */
protected void assertSimplify(String sql, String expected) {
  SqlAdvisor advisor = tester.getFactory().createAdvisor();

  SqlParserUtil.StringAndPos sap = SqlParserUtil.findPos(sql);
  String actual = advisor.simplifySql(sap.sql, sap.cursor);
  Assertions.assertEquals(expected, actual);
}