83 lines
2 KiB
Lua
83 lines
2 KiB
Lua
local run_formatter = function(text)
|
|
vim.notify(text)
|
|
vim.notify("----")
|
|
local result = {}
|
|
require("plenary.job"):new({
|
|
command = "sqlfluff",
|
|
cwd = "/usr/bin",
|
|
args = { "format", "-" },
|
|
writer = text,
|
|
on_stdout = function(_, line)
|
|
table.insert(result, line)
|
|
end,
|
|
}):sync()
|
|
|
|
-- add the surrounding r#"..."# back
|
|
for _, line in ipairs(result) do
|
|
vim.notify(line)
|
|
end
|
|
-- table.insert(result, "\"#")
|
|
return result
|
|
end
|
|
|
|
local embedded_sql = vim.treesitter.query.parse(
|
|
"rust",
|
|
[[
|
|
(macro_invocation
|
|
(scoped_identifier
|
|
path: (identifier) @path (#eq? @path "sqlx")
|
|
name: (identifier) @name (#any-of? @name "query" "query_scalar" "query_as"))
|
|
|
|
(token_tree
|
|
(raw_string_literal) @sql)
|
|
(#offset! @sql 0 3 0 -2)
|
|
)
|
|
]]
|
|
)
|
|
|
|
local get_root = function(bufnr)
|
|
local parser = vim.treesitter.get_parser(bufnr, "rust", {})
|
|
local tree = parser:parse()[1]
|
|
return tree:root()
|
|
end
|
|
|
|
local format = function(bufnr)
|
|
bufnr = bufnr or vim.api.nvim_get_current_buf()
|
|
|
|
if vim.api.nvim_get_option_value("filetype", { buf = bufnr }) ~= "rust" then
|
|
vim.notify("SQL format an only be used in Rust")
|
|
return
|
|
end
|
|
|
|
local root = get_root(bufnr)
|
|
|
|
local changes = {}
|
|
for id, node in embedded_sql:iter_captures(root, bufnr, 0, -1) do
|
|
local name = embedded_sql.captures[id]
|
|
if name == "sql" then
|
|
-- {start row, start col, end row, end col }
|
|
local range = { node:range() }
|
|
local indentation = string.rep(" ", range[2])
|
|
-- run the formatter on the node text
|
|
local formatted = run_formatter(vim.treesitter.get_node_text(node, bufnr))
|
|
-- add indentation
|
|
for idx, line in ipairs(formatted) do
|
|
formatted[idx] = indentation .. line
|
|
end
|
|
-- add changes in reverse order
|
|
table.insert(changes, 1, {
|
|
start = range[1],
|
|
final = range[3],
|
|
formatted = formatted,
|
|
})
|
|
end
|
|
end
|
|
|
|
for _, change in ipairs(changes) do
|
|
vim.api.nvim_buf_set_lines(bufnr, change.start, change.final, false, change.formatted)
|
|
end
|
|
end
|
|
|
|
vim.api.nvim_create_user_command("SqlFormat", function()
|
|
format()
|
|
end, {})
|