Sql commands for neovim

sql_commands.lua

-- DB commands
--
-- Comandos para ejecutar SQL usando vim-dadbod (:h dadbod).
-- Cada comando escribe en el buffer, debajo del cursor (o de la selección),
-- el query ejecutado y su resultado. Así el buffer queda como una "historia"
-- que se puede copiar completa como contexto para un prompt.
--
-- Si no hay una conexión seleccionada, se pregunta con vim.ui.select.
-- Si el comando necesita una tabla y no se indica, también se pregunta.
--
-- Mappings (normal: línea actual, visual: selección):
--   <C-S> / <leader>S   :DBExecute
--
-- Comandos:
--   :DBExecute             ejecuta la línea actual o el rango
--   :DBDescribe            elige la tabla con vim.ui.select
--   :DBDescribeIndexes     idem
--   :DBShowCreateTable     idem
--   :DBTables              lista las tablas
--   :DBShowSchema          SHOW CREATE TABLE de todas las tablas (menos las ignoradas)
--   :DBToCsv [archivo]     exporta la línea actual o el rango a CSV (default: export.csv)
--   :DBNewQuery            buffer SQL nuevo ligado a la conexión actual
--
-- Para conexiones con vim-dotenv ver plugin/61-sql-gui.lua

local M = {}

-- Tablas que DBShowSchema no muestra
local IGNORED_TABLES = { "sqlite_sequence", "adonis_schema", "adonis_schema_versions" }

-- Queries por motor de base de datos (esquema de la url).
-- Cada entrada recibe el nombre de la tabla y devuelve el statement.
-- Los comandos que empiezan con "." (sqlite) o "\" (psql) son meta-comandos del
-- cliente: se envían tal cual, sin ";".
-- vertical = true (mysql) usa -E -r: una columna por línea y sin escapar saltos de línea,
-- ideal para CREATE TABLE. No se usa \G porque el cliente mysql de Homebrew lo rechaza
-- ("Unknown command '\G'").
-- Refs: https://www.sqlite.org/cli.html  https://www.postgresql.org/docs/current/app-psql.html
local MYSQL = {
  list_tables = "SHOW TABLES",
  describe = function(t) return string.format("DESCRIBE `%s`", t) end,
  indexes = function(t) return string.format("SHOW INDEXES FROM `%s`", t) end,
  create_table = function(t) return string.format("SHOW CREATE TABLE `%s`", t) end,
  vertical = true,
}

local POSTGRES = {
  list_tables = "SELECT schemaname, tablename FROM pg_tables WHERE schemaname NOT IN ('pg_catalog','information_schema') ORDER BY 1, 2",
  describe = function(t) return string.format("\\d %s", t) end,
  indexes = function(t) return string.format("SELECT indexname, indexdef FROM pg_indexes WHERE tablename = '%s'", t) end,
  create_table = function(t) return string.format("\\d+ %s", t) end,
}

local SQLITE = {
  list_tables = "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name",
  describe = function(t) return string.format("PRAGMA table_info('%s')", t) end,
  indexes = function(t) return string.format("PRAGMA index_list('%s')", t) end,
  create_table = function(t) return string.format(".schema %s", t) end,
}

local DIALECTS = {
  mysql = MYSQL,
  mariadb = MYSQL,
  postgresql = POSTGRES,
  postgres = POSTGRES,
  sqlite = SQLITE,
}

-- Devuelve el dialecto de la conexión (nil y avisa si no está soportado)
local function get_dialect(db)
  local scheme = tostring(db):match("^(%w+)")
  local dialect = DIALECTS[scheme]
  if not dialect then
    vim.notify("Esquema no soportado: " .. tostring(scheme), vim.log.levels.ERROR)
  end
  return dialect
end

----------------------
--- CONEXIÓN ---------
----------------------

-- Devuelve las conexiones que conoce DBUI (:h db_ui)
local function get_connections()
  local ok, conns = pcall(vim.fn["db_ui#connections_list"])
  return ok and conns or {}
end

-- Resuelve la conexión y llama a `cb(db)`.
-- Orden: b:db -> g:dbui_current_db -> preguntar (y guardar en global)
function M.with_db(cb)
  local db = vim.b.db
  if db and db ~= "" then
    return cb(db)
  end

  db = vim.g.dbui_current_db
  if db and db ~= "" then
    return cb(db)
  end

  local conns = get_connections()
  if #conns == 0 then
    vim.notify("No hay conexiones definidas en DBUI.", vim.log.levels.WARN)
    return
  end

  vim.ui.select(conns, {
    prompt = "Selecciona una conexión:",
    format_item = function(c)
      return c.name .. " (" .. c.source .. ")"
    end,
  }, function(choice)
    if not choice then
      return
    end
    vim.g.dbui_current_db = choice.url -- queda guardada
    cb(choice.url)
  end)
end

-- Quita de la lista de tablas el ruido que a veces trae el cliente: el encabezado
-- "Tables_in_<db>" y el aviso "mysql: [Warning] Using a password..." (stderr se mezcla)
local function clean_tables(tables)
  return vim.tbl_filter(function(name)
    return not (name:match("^Tables_in_") or name:match("^mysql: ") or name:match("^mariadb: "))
  end, tables)
end

-- Devuelve la lista de tablas de la conexión usando el adaptador de dadbod
-- Ref: autoload/db/adapter/*.vim (función `tables`)
local function get_tables(db)
  local ok, tables = pcall(vim.fn["db#adapter#call"], db, "tables", { db }, {})
  return ok and clean_tables(tables) or {}
end

-- Resuelve la tabla: usa `name` si viene, si no pregunta con vim.ui.select.
-- Llama a `cb(db, table_name)`
local function with_table(name, cb)
  M.with_db(function(db)
    if name and name ~= "" then
      return cb(db, name)
    end

    local tables = get_tables(db)
    if #tables == 0 then
      vim.notify("No se encontraron tablas.", vim.log.levels.WARN)
      return
    end

    vim.ui.select(tables, { prompt = "Selecciona una tabla:" }, function(choice)
      if choice then
        cb(db, choice)
      end
    end)
  end)
end

----------------------
--- EJECUCIÓN --------
----------------------

-- Guarda dónde escribir el resultado (el selector es asíncrono y el cursor puede moverse)
local function get_target(line)
  return { buf = vim.api.nvim_get_current_buf(), line = line or vim.api.nvim_win_get_cursor(0)[1] }
end

-- Flags del cliente mysql para salida vertical (ver --vertical y --raw en `mysql --help`)
local VERTICAL_FLAGS = { "-E", "-r" }

-- Ejecuta `sql` con el cliente de la conexión y devuelve las líneas de salida
-- Usa el comando `filter` del adaptador de dadbod (ej. mysql -t); si el adaptador
-- no lo tiene (ej. sqlite) usa `interactive`. Ref: autoload/db/adapter/*.vim
local function run_sql(db, sql, vertical)
  local kind = vim.fn["db#adapter#supports"](db, "filter") == 1 and "filter" or "interactive"
  local cmd = vim.fn["db#adapter#dispatch"](db, kind)
  if vertical then
    vim.list_extend(cmd, VERTICAL_FLAGS)
  end
  local res = vim.system(cmd, { stdin = sql, text = true }):wait()
  -- stderr solo si falló (mysql avisa del password en la línea de comandos aunque funcione)
  local output = res.code == 0 and res.stdout or (res.stdout .. res.stderr)
  return vim.split(output, "\n", { trimempty = true })
end

-- Escribe debajo de target: línea en blanco, `shown` (lo ejecutado), línea en blanco y resultado
-- Si `shown` está vacío solo escribe el resultado
local function insert_result(target, shown, result)
  local lines = { "" }
  if #shown > 0 then
    vim.list_extend(lines, shown)
    table.insert(lines, "")
  end
  vim.list_extend(lines, result)
  vim.api.nvim_buf_set_lines(target.buf, target.line, target.line, false, lines)
end

-- Agrega ";" al statement, salvo en meta-comandos (".schema", "\d")
local function with_terminator(statement)
  if statement:match("^[.\\]") then
    return statement
  end
  return statement .. ";"
end

-- Ejecuta un statement y escribe el statement + resultado en el buffer.
local function run_statement(db, target, statement, vertical)
  local sql = with_terminator(statement)
  insert_result(target, { sql }, run_sql(db, sql, vertical))
end

-- Ejecuta el statement `kind` del dialecto para una tabla y lo escribe en el buffer.
-- `kind` es la entrada del dialecto: "describe", "indexes" o "create_table"
-- `table_name` vacío o nil -> pregunta la tabla con vim.ui.select
local function run_table_statement(kind, table_name)
  local target = get_target()
  with_table(table_name, function(db, name)
    local dialect = get_dialect(db)
    if dialect then
      local vertical = dialect.vertical and kind == "create_table"
      run_statement(db, target, dialect[kind](name), vertical)
    end
  end)
end

-- Crea un comando sin argumentos que elige la tabla con vim.ui.select
local function table_command(name, kind, desc)
  vim.api.nvim_create_user_command(name, function()
    run_table_statement(kind, nil)
  end, { desc = desc })
end

-- :DBExecute -> línea actual (o selección con :'<,'>DBExecute)
-- El query ya está en el buffer, solo se agrega el resultado debajo
vim.api.nvim_create_user_command("DBExecute", function(opts)
  local lines = vim.api.nvim_buf_get_lines(0, opts.line1 - 1, opts.line2, false)
  local sql = table.concat(lines, "\n")
  local target = get_target(opts.line2)
  M.with_db(function(db)
    insert_result(target, {}, run_sql(db, sql))
  end)
end, { range = true, desc = "Ejecutar SQL de la línea actual o selección" })

table_command("DBDescribe", "describe", "Describir una tabla")
table_command("DBDescribeIndexes", "indexes", "Mostrar índices de una tabla")
table_command("DBShowCreateTable", "create_table", "Mostrar CREATE TABLE de una tabla")

-- :DBTables -> lista todas las tablas de la conexión actual
vim.api.nvim_create_user_command("DBTables", function()
  local target = get_target()
  M.with_db(function(db)
    local dialect = get_dialect(db)
    if dialect then
      run_statement(db, target, dialect.list_tables)
    end
  end)
end, { desc = "Listar tablas de la conexión actual" })

-- :DBShowSchema -> CREATE TABLE de todas las tablas, menos las ignoradas
vim.api.nvim_create_user_command("DBShowSchema", function()
  local target = get_target()
  M.with_db(function(db)
    local dialect = get_dialect(db)
    if not dialect then
      return
    end
    local statements = {}
    for _, name in ipairs(get_tables(db)) do
      if not vim.tbl_contains(IGNORED_TABLES, name) then
        table.insert(statements, with_terminator(dialect.create_table(name)))
      end
    end
    local shown = { "-- CREATE TABLE de todas las tablas" }
    insert_result(target, shown, run_sql(db, table.concat(statements, "\n"), dialect.vertical))
  end)
end, { desc = "CREATE TABLE de todas las tablas" })

-- :DBNewQuery -> buffer SQL nuevo ligado a la conexión actual
vim.api.nvim_create_user_command("DBNewQuery", function()
  M.with_db(function(db)
    vim.cmd("enew")
    vim.bo.filetype = "sql"
    vim.b.db = db
  end)
end, { desc = "Nueva query en la conexión actual" })

----------------------
--- MAPPINGS ---------
----------------------

-- Normal: ejecuta la línea actual. Visual: ":" agrega solo el rango '<,'>
for _, lhs in ipairs({ "<C-S>", "<leader>S" }) do
  vim.keymap.set("n", lhs, "<Cmd>DBExecute<CR>", { desc = "SQL: ejecutar línea" })
  vim.keymap.set("x", lhs, ":DBExecute<CR>", { desc = "SQL: ejecutar selección", silent = true })
end

----------------------
--- CSV EXPORTING ---
----------------------

-- Convierte los escapes del modo batch de mariadb (-B) a texto real
local function unescape_batch(field)
  if field == "\\N" then
    return nil -- NULL
  end
  local escapes = { n = "\n", t = "\t", r = "\r", ["0"] = "\0" }
  return (field:gsub("\\(.)", function(c)
    return escapes[c] or c -- \\ y cualquier otro
  end))
end

-- Escapa un campo para CSV (RFC 4180)
local function csv_field(value)
  if value == nil then
    return ""
  end
  if value:find('[",\r\n]') then
    return '"' .. value:gsub('"', '""') .. '"'
  end
  return value
end

-- Convierte las filas TSV de mariadb -B en líneas CSV
local function tsv_to_csv(rows)
  local csv = {}
  for _, row in ipairs(rows) do
    local fields = {}
    for field in (row .. "\t"):gmatch("(.-)\t") do -- el TSV nunca trae tabs reales dentro de un campo
      table.insert(fields, csv_field(unescape_batch(field)))
    end
    table.insert(csv, table.concat(fields, ","))
  end
  return csv
end

-- Ejecuta sql con el cliente de la conexión en modo batch (-B) y devuelve las filas
-- Usa el comando interactivo del adaptador (mysql/mariadb) de dadbod.
local function query_batch(db, sql)
  local cmd = vim.fn["db#adapter#dispatch"](db, "interactive")
  vim.list_extend(cmd, { "-B", "--default-character-set=utf8mb4" })
  return vim.fn["db#systemlist"](cmd, sql)
end

-- :DBToCsv [archivo] -> exporta la línea actual o el rango a un CSV
vim.api.nvim_create_user_command("DBToCsv", function(opts)
  local lines = vim.api.nvim_buf_get_lines(0, opts.line1 - 1, opts.line2, false)
  local sql = vim.trim(table.concat(lines, "\n"))
  if sql == "" then
    return -- query vacío: no hacer nada
  end

  M.with_db(function(db)
    local rows = query_batch(db, sql)
    if #rows == 0 then
      vim.notify("El query no devolvió resultados (o falló)", vim.log.levels.WARN)
      return
    end

    local path = vim.fn.fnamemodify(opts.args ~= "" and opts.args or "export.csv", ":p")
    vim.fn.mkdir(vim.fn.fnamemodify(path, ":h"), "p")

    local csv = tsv_to_csv(rows)
    if vim.fn.writefile(csv, path) ~= 0 then
      vim.notify("No se pudo escribir " .. path, vim.log.levels.ERROR)
      return
    end
    vim.notify(string.format("Exportadas %d líneas a %s", #csv, path), vim.log.levels.INFO)
  end)
end, { range = true, nargs = "?", complete = "file", desc = "Exportar query a CSV" })

return M