codex.rb

require "openai"
require "pg"
require "mini_sql"

API_KEY = ENV["OPEN_AI_API_KEY"]
MAX_TABLE_DEF_LEN = 6000

pg_conn = PG.connect(dbname: "discourse_development")
conn = MiniSql::Connection.get(pg_conn)

schema = []

table_name = nil
columns = nil

priority_tables = %w[posts topics notifications users user_actions]

conn
  .query(
    "select table_name, column_name from information_schema.columns order by case when table_name in (?) then 0 else 1 end asc, table_name ",
    priority_tables
  )
  .each do |row|
    if table_name != row.table_name
      schema << "#{table_name}(#{columns.join(",")})" if columns
      table_name = row.table_name
      columns = []
    end
    columns << row.column_name
  end

schema << "#{table_name}(#{columns.join(",")})"

puts "What query would you like to write?"
question = gets.chomp.gsub("\n", "")

schema_comment = +""
schema.each do |table|
  schema_comment << "# #{table}\n"
  break if schema_comment.length > MAX_TABLE_DEF_LEN
end

client = OpenAI::Client.new(access_token: API_KEY)
response =
  client.completions(
    parameters: {
      model: "code-davinci-002",
      prompt:
        "### Postgres SQL running against discourse/discourse repo\n### SQL for #{question}\nSELECT",
      max_tokens: 150,
      temperature: 0,
      top_p: 1.0,
      stop: %w[# ;]
    }
  )

puts "SELECT" + response.dig("choices", 0, "text")