SQLParser { public static Collection<PrimaryKeyDefinition> parsePrimaryKeys(String sqlFile) { Collection<PrimaryKeyDefinition> primaryKeys = new LinkedList<>(); List<String> sqlStatements = loadStatements(sqlFile); for (String sqlStatement : sqlStatements) { SQLiteLexer lexer = new SQLiteLexer(new ANTLRInputStream(sqlStatement)); SQLiteParser parser = new SQLiteParser(new CommonTokenStream(lexer)); parser.removeErrorListeners(); ParseTree tree = parser.sql_stmt(); ParseTreeWalker.DEFAULT.walk(new SQLiteBaseListener() { String tableName; String columnName; @Override public void enterCreate_table_stmt(SQLiteParser.Create_table_stmtContext ctx) { super.enterCreate_table_stmt(ctx); if (ctx.table_name() != null) { this.tableName = cleanIdentifier(ctx.table_name().getText()); } } @Override public void enterAlter_table_stmt(SQLiteParser.Alter_table_stmtContext ctx) { super.enterAlter_table_stmt(ctx); if (ctx.table_name() != null) { this.tableName = cleanIdentifier(ctx.table_name().getText()); } } @Override public void enterColumn_def(SQLiteParser.Column_defContext ctx) { super.enterColumn_def(ctx); this.columnName = cleanIdentifier(ctx.column_name().getText()); } @Override public void enterColumn_constraint(SQLiteParser.Column_constraintContext ctx) { super.enterColumn_constraint(ctx); if (ctx.K_PRIMARY() == null) return; if (this.tableName == null || this.columnName == null) return; primaryKeys.add(new PrimaryKeyDefinition(this.tableName, this.columnName)); } @Override public void enterTable_constraint(SQLiteParser.Table_constraintContext ctx) { super.enterTable_constraint(ctx); if (ctx.K_PRIMARY() == null) return; if (this.tableName == null) return; List<String> columns = new ArrayList<>(); for (SQLiteParser.Indexed_columnContext columnCtx : ctx.indexed_column()) { columns.add(cleanIdentifier(columnCtx.column_name().getText())); } primaryKeys.add(new PrimaryKeyDefinition(this.tableName, columns)); } }, tree); } return primaryKeys; }  static Map<String, List<String>> parseSchema(String sqlFile); static Collection<PrimaryKeyDefinition> parsePrimaryKeys(String sqlFile); static Collection<ForeignKeyDefinition> parseForeignKeys(String sqlFile); static Collection<NotNullDefinition> parseNotNullConstraints(String sqlFile); static Collection<DataTypeDefinition> parseDataTypes(String sqlFile);  }
@Test public void testParsePrimaryKeyDefinitions() { String inputFile = Thread.currentThread().getContextClassLoader().getResource("primary-keys.sql").getFile(); Collection<PrimaryKeyDefinition> primaryKeyDefinitions = SQLParser.parsePrimaryKeys(inputFile); Assert.assertEquals( new HashSet<>(Arrays.asList( new PrimaryKeyDefinition("table1", "id"), new PrimaryKeyDefinition("table2", "id"), new PrimaryKeyDefinition("table3", "first_name", "last_name"), new PrimaryKeyDefinition("table4", "added_id"), new PrimaryKeyDefinition("table5", "added_fn", "added_ln") )), new HashSet<>(primaryKeyDefinitions) ); }