in tools/unsupervised_dataset/sql_classifier/classifier/src/main/java/Classifier.java [32:71]
public static void main(String[] args) {
if (args.length == 0) {
LOGGER.log(Level.SEVERE, "Please provide a CSV file.");
return;
}
Optional<List<String[]>> allData = readCSV(args[0]);
if (allData.isEmpty()) {
return;
}
SqlParser.Config[] parserConfigs = new SqlParser.Config[4];
parserConfigs[0] = SqlParser.configBuilder().setParserFactory(DefaultDialectParserImpl.FACTORY).build();
parserConfigs[1] = SqlParser.configBuilder().setParserFactory(Dialect1ParserImpl.FACTORY).build();
parserConfigs[2] = SqlParser.configBuilder().setParserFactory(BigQueryParserImpl.FACTORY).build();
parserConfigs[3] = SqlParser.configBuilder().setParserFactory(PostgreSQLParserImpl.FACTORY).build();
CSVWriter[] writers = setupOutput();
for (String[] data : allData.get()) {
Map<SqlParser.Config, Boolean> results = classifyQuery(cleanQuery(data[0]), parserConfigs);
boolean unclassified = true;
String[] nextLine = {data[0], data[1]};
for (int i = 0; i < parserConfigs.length; i++) {
if (results.get(parserConfigs[i])) {
unclassified = false;
writers[i].writeNext(nextLine);
}
}
if (unclassified) {
writers[writers.length-1].writeNext(nextLine);
}
}
try {
for (CSVWriter writer : writers) {
writer.close();
}
} catch (Exception e) {
LOGGER.log(Level.SEVERE, e.getMessage());
}
}