summaryrefslogtreecommitdiff
path: root/plugin/promql.lua
blob: e44d1c3f31053dde4b673dbbab4c13dfb2731fd8 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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 = "%" })