spring-ai/text-to-sql/src/main/java/com/ankurm/texttosql/SchemaPrompt.java

package com.ankurm.texttosql;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * Builds the schema section of the prompt from the database itself (information_schema plus the
 * COMMENTs on tables and columns), for an explicit list of tables. A table that is not in the list
 * is not described, so the model is never shown that it exists.
 */
public final class SchemaPrompt {

    private SchemaPrompt() {
    }

    public static String describe(Connection c, List<String> tables) throws SQLException {
        StringBuilder sb = new StringBuilder();
        for (String table : tables) {
            String tableComment = scalar(c, "SELECT obj_description(?::regclass, 'pg_class')", table);
            sb.append("TABLE ").append(table);
            if (tableComment != null) {
                sb.append("  -- ").append(tableComment);
            }
            sb.append('\n');
            try (PreparedStatement ps = c.prepareStatement(
                    "SELECT column_name, data_type, col_description(?::regclass, ordinal_position) "
                            + "FROM information_schema.columns WHERE table_schema = 'public' AND table_name = ? "
                            + "ORDER BY ordinal_position")) {
                ps.setString(1, table);
                ps.setString(2, table);
                try (ResultSet rs = ps.executeQuery()) {
                    while (rs.next()) {
                        sb.append("  ").append(rs.getString(1)).append(' ').append(rs.getString(2));
                        if (rs.getString(3) != null) {
                            sb.append("  -- ").append(rs.getString(3));
                        }
                        sb.append('\n');
                    }
                }
            }
        }
        return sb.toString();
    }

    private static String scalar(Connection c, String sql, String arg) throws SQLException {
        try (PreparedStatement ps = c.prepareStatement(sql)) {
            ps.setString(1, arg);
            try (ResultSet rs = ps.executeQuery()) {
                return rs.next() ? rs.getString(1) : null;
            }
        }
    }

    public static List<String> columnNames(Connection c, String table) throws SQLException {
        List<String> out = new ArrayList<>();
        try (PreparedStatement ps = c.prepareStatement(
                "SELECT column_name FROM information_schema.columns WHERE table_schema='public' AND table_name=? ORDER BY ordinal_position")) {
            ps.setString(1, table);
            try (ResultSet rs = ps.executeQuery()) {
                while (rs.next()) {
                    out.add(rs.getString(1));
                }
            }
        }
        return out;
    }
}