1 files changed, 130 insertions, 0 deletions
diff --git a/plugin/promql.lua b/plugin/promql.lua
new file mode 100644
index 0000000..e44d1c3
--- /dev/null
+++ b/plugin/promql.lua
@@ -0,0 +1,130 @@
+-- How it started
+--vim.api.nvim_create_user_command(
+--  "Promqlfmt",
+--  '<line1>,<line2>!promtool promql format --experimental "$(cat /dev/stdin)"',
+--  { nargs = 0, range = "%" }
+--)
+
+-- How it's going
+if vim.g.promqlfmt == 1 then
+  return
+end
+
+vim.g.promqlfmt = 1
+
+--- Format a PromQL query using promtool
+--- @param query string @query to format
+--- @param padding number|nil @number of spaces to add to each line, default 0
+--- @param padchar string|nil @character for padding, default " "
+--- @return string[]|nil @formatted query, one element per line or nil
+local promql_format = function(query, padding, padchar)
+  padding = padding or 0
+  padchar = padchar or " "
+
+  local success, result = pcall(function()
+    return vim
+      .system(
+        { "promtool", "promql", "format", "--experimental", query },
+        { text = true }
+      )
+      :wait()
+  end)
+
+  if success and result and result.code == 0 then
+    local lines = vim.fn.split(result.stdout, "\n")
+    if padding == 0 then
+      return lines
+    end
+
+    for i, line in ipairs(lines or {}) do
+      lines[i] = (padchar):rep(padding) .. line
+    end
+    return lines
+  else
+    print("woops error occurred while formatting query")
+    print(vim.inspect(result))
+  end
+end
+
+--- Get the last column of a line, ie the position of the last character
+--- @param bufnr number
+--- @param linenr number
+local last_col_line = function(bufnr, linenr)
+  local line = vim.api.nvim_buf_get_lines(bufnr, linenr, linenr + 1, true)[1]
+  return #line
+end
+
+--- Get the selected text in a buffer
+--- @param bufnr number|nil
+--- @return table
+local buf_get_selected_text = function(bufnr)
+  local r = {
+    text = "",
+    start_line = nil,
+    start_row = nil,
+    end_line = nil,
+    end_row = nil,
+  }
+  local region = vim.region(bufnr or 0, "'<", "'>", vim.fn.visualmode(), true)
+  local maxcol = vim.v.maxcol
+  for line, cols in vim.spairs(region) do
+    if r.start_line == nil then
+      r.start_line = line
+      r.start_row = cols[1]
+    end
+
+    local endcol = cols[2] == maxcol and -1 or cols[2]
+    local chunk =
+      vim.api.nvim_buf_get_text(0, line, cols[1], line, endcol, {})[1]
+    r.text = ("%s%s\n"):format(r.text, chunk)
+
+    r.end_line = line
+    r.end_row = endcol
+  end
+
+  r.end_row = r.end_row > last_col_line(0, r.end_line) and -1 or r.end_row
+  return r
+end
+
+vim.api.nvim_create_user_command("Promqlfmt", function(opts)
+  -- it was not called from a visual selection
+  if opts.range == 0 then
+    -- get all lines
+    local buf_text = vim.api.nvim_buf_get_lines(0, 0, -1, true)
+    -- format query
+    local text_formatted = promql_format(table.concat(buf_text, "\n"))
+    if text_formatted == nil then
+      return
+    end
+    -- replace all text with the formatted version
+    vim.api.nvim_buf_set_lines(0, 0, -1, true, text_formatted)
+    return
+  end
+
+  local s = buf_get_selected_text()
+  -- use the position of the first non-whitespace character of the first line
+  -- as padding
+  local padding = string.find(
+    vim.api.nvim_buf_get_lines(0, s.start_line, s.start_line + 1, true)[1],
+    "%S"
+  ) - 1
+  local text_formatted = promql_format(s.text, padding)
+
+  if text_formatted == nil then
+    return
+  end
+
+  -- remove padding from first line
+  if s.start_row ~= 0 then
+    text_formatted[1] = text_formatted[1]:sub(padding + 1)
+  end
+
+  vim.api.nvim_buf_set_text(
+    0,
+    s.start_line,
+    s.start_row,
+    s.end_line,
+    s.end_row,
+    text_formatted
+  )
+end, { nargs = 0, range = "%" })