sqlx formatting, kinda jank so far

lazy-lock.json

@@ -1,5 +1,5 @@
 {
-  "LuaSnip": { "branch": "master", "commit": "a658ae2906344a1d2b9c507738e585cf68e685c7" },
+  "LuaSnip": { "branch": "master", "commit": "105b5f7f72c13e682a3aa5d29eac2408ae513b22" },
   "auto-save.nvim": { "branch": "main", "commit": "979b6c82f60cfa80f4cf437d77446d0ded0addf0" },
   "autoclose.nvim": { "branch": "main", "commit": "469782b0456f0b4f764378ffda94c18599544e09" },
   "ccc.nvim": { "branch": "main", "commit": "4a0ddaf787cc82796e84ab8a7f70d086f250aeb6" },
@@ -9,30 +9,30 @@
   "cmp_luasnip": { "branch": "master", "commit": "18095520391186d634a0045dacaa346291096566" },
   "crates.nvim": { "branch": "main", "commit": "1dffccc0a95f656ebe00cacb4de282473430c5a1" },
   "glow.nvim": { "branch": "advanced_window", "commit": "f1157d4cb7e46e830c72004e7e1adb81a1f9b04c" },
-  "indent-blankline.nvim": { "branch": "master", "commit": "4541d690816cb99a7fc248f1486aa87f3abce91c" },
+  "indent-blankline.nvim": { "branch": "master", "commit": "7075d7861f7a6bbf0de0298c83f8a13195e6ec01" },
-  "lazy.nvim": { "branch": "main", "commit": "5af331ea65418dc9361769891921fdee4bcc837a" },
+  "lazy.nvim": { "branch": "main", "commit": "0d61488b89a570415177f75a36ef93616aac6c77" },
   "lazygit.nvim": { "branch": "main", "commit": "3466e48439601445e26c65635421625886f2d0c0" },
   "lsp-inlayhints.nvim": { "branch": "main", "commit": "d981f65c9ae0b6062176f0accb9c151daeda6f16" },
   "lualine.nvim": { "branch": "master", "commit": "05d78e9fd0cdfb4545974a5aa14b1be95a86e9c9" },
-  "mason-lspconfig.nvim": { "branch": "main", "commit": "82685fdd0f67a694b244cb058b83761f54664d68" },
+  "mason-lspconfig.nvim": { "branch": "main", "commit": "4f1c72767bec31397d59554f84096909b2887195" },
-  "mason-null-ls.nvim": { "branch": "main", "commit": "ae0c5fa57468ac65617f1bf821ba0c3a1e251f0c" },
+  "mason-null-ls.nvim": { "branch": "main", "commit": "73c68abdf65279e41526eb152876511a8ae84ea2" },
-  "mason.nvim": { "branch": "main", "commit": "5ad3e113b0c3fde3caba8630599373046f6197e8" },
+  "mason.nvim": { "branch": "main", "commit": "b68d3be4b664671002221d43c82e74a0f1006b26" },
-  "material.nvim": { "branch": "main", "commit": "c7631d373d3ae02ef502ec6b8620a8ff36ab922e" },
+  "material.nvim": { "branch": "main", "commit": "f3783f2a8e55904887d3ea7cae8d7ab57d7f3e8e" },
   "neo-tree.nvim": { "branch": "v2.x", "commit": "20c2f2f5ba083bbb1e37b8bc3d590621434f31e9" },
-  "neodev.nvim": { "branch": "main", "commit": "0705c72c281626f45c78d0916151d3e2bfb53fae" },
+  "neodev.nvim": { "branch": "main", "commit": "a2b1d8fb9fa4daa35d3fd9123bccccccbd4a3520" },
   "nerdcommenter": { "branch": "master", "commit": "ab2ae4d502a26bc591db78a8548823ddd04bbc9c" },
   "nui.nvim": { "branch": "main", "commit": "d146966a423e60699b084eeb28489fe3b6427599" },
-  "null-ls.nvim": { "branch": "main", "commit": "db09b6c691def0038c456551e4e2772186449f35" },
+  "null-ls.nvim": { "branch": "main", "commit": "aac27a1fa550de3d0b2c651168167cc0d5366a9a" },
   "nvim-cmp": { "branch": "main", "commit": "2743dd989e9b932e1b4813a4927d7b84272a14e2" },
-  "nvim-lspconfig": { "branch": "master", "commit": "51739a8bc531542079698c58512feb68469f5d27" },
+  "nvim-lspconfig": { "branch": "master", "commit": "0011c435282f043a018e23393cae06ed926c3f4a" },
-  "nvim-treesitter": { "branch": "master", "commit": "1ef286e5b0cfd17f56586a8445cd83d61647f851" },
+  "nvim-treesitter": { "branch": "master", "commit": "0522bbe90d6f9bb7080621fe70f8ab565e072b98" },
   "nvim-ts-autotag": { "branch": "main", "commit": "6be1192965df35f94b8ea6d323354f7dc7a557e4" },
   "nvim-ts-context-commentstring": { "branch": "main", "commit": "7f625207f225eea97ef7a6abe7611e556c396d2f" },
   "nvim-web-devicons": { "branch": "master", "commit": "9ab9b0b894b2388a9dbcdee5f00ce72e25d85bf9" },
   "openingh.nvim": { "branch": "main", "commit": "2719e5759ecf4b9a2d492fbf52d03d2e6fc6126a" },
   "playground": { "branch": "master", "commit": "2b81a018a49f8e476341dfcb228b7b808baba68b" },
-  "plenary.nvim": { "branch": "master", "commit": "bda256fab0eb66a15e8190937e417e6a14ee5d72" },
+  "plenary.nvim": { "branch": "master", "commit": "102c02903c74b93c705406bf362049383abc87c8" },
-  "rest.nvim": { "branch": "main", "commit": "0d4b6176ce7fee2d33a716402d47d65a35acf173" },
+  "rest.nvim": { "branch": "main", "commit": "aea7c64bdff1073beed9bd7fddb60cce7796d7ff" },
   "rust-tools.nvim": { "branch": "master", "commit": "71d2cf67b5ed120a0e31b2c8adb210dd2834242f" },
   "telescope-fzf-native.nvim": { "branch": "main", "commit": "9bc8237565ded606e6c366a71c64c0af25cd7a50" },
   "telescope.nvim": { "branch": "master", "commit": "c1a2af0af69e80e14e6b226d3957a064cd080805" },

lua/keymaps.lua

@@ -38,6 +38,7 @@ require("util").keymaps({
     -- reformat LSP
     ["<leader>lf"] = {
       function()
+        vim.cmd "SqlMagic"
         vim.lsp.buf.format()
       end,
       desc = "Reformat file",

lua/plugins/lsp.lua

@@ -25,8 +25,17 @@ local servers = {
   lua_ls = {
     -- lua
     Lua = {
-      workspace = { checkThirdParty = false },
+      runtime = {
+        version = "LuaJIT",
+      },
+      workspace = { checkThirdParty = false, library = vim.api.nvim_get_runtime_file("", true) },
       telemetry = { enable = false },
+      diagnostics = {
+        globals = {
+          "vim",
+          "require"
+        },
+      },
     },
   },
   bashls = {
@@ -138,6 +147,7 @@ return {
     build = ":MasonUpdate",
     opts = {},
   },
+  { "folke/neodev.nvim",  opts = {} }, -- lua stuff
   {
     "williamboman/mason-lspconfig.nvim",
   },
@@ -238,7 +248,6 @@ return {
       })
     end,
   },
-  { "folke/neodev.nvim",  opts = {} }, -- lua stuff
   { -- Rust tools
     "simrat39/rust-tools.nvim",
     build = prereqs,

lua/plugins/null-ls.lua

@@ -1,19 +1,137 @@
 function prereqs()
-  local output = vim.fn.system({
+  local output_cspell = vim.fn.system({
     "which",
     "cspell",
   })
-  if output == nil or output == "" then
+  if output_cspell == nil or output_cspell == "" or output_cspell == "cspell not found" then
     print("Installing cspell globally with npm")
     vim.fn.system({
       "npm",
       "install",
       "-g",
-      "cspell@latest",
+      "cspell",
     })
   end
 end
 
+-- https://github.com/darold/pgFormatter
+
+local rust_sqlx_f = function()
+  return vim.treesitter.query.parse(
+    "rust",
+    [[
+; query macro
+(macro_invocation
+   (scoped_identifier
+     path: (identifier) @_path (#eq? @_path "sqlx")
+     name: (identifier) @_name (#any-of? @_name "query" "query_scalar"))
+
+ (token_tree
+   . (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+ )
+
+; query_as macro
+(macro_invocation
+   (scoped_identifier
+     path: (identifier) @_path (#eq? @_path "sqlx")
+     name: (identifier) @_name (#eq? @_name "query_as"))
+
+ (token_tree
+   (_) . (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+
+ )
+
+; query and query_as function
+(call_expression
+  (scoped_identifier
+    path: (identifier) @_path (#eq? @_path "sqlx")
+    name: (identifier) @_name (#contains? @_name "query"))
+
+ (arguments
+   (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+  )
+  ]]
+  )
+end
+
+local get_root = function(bufnr)
+  local parser = vim.treesitter.get_parser(bufnr, "rust", {})
+  local tree = parser:parse()[1]
+  return tree:root()
+end
+
+local format_dat_sql = function(bufnr)
+  local rust_sqlx = rust_sqlx_f()
+  bufnr = bufnr or vim.api.nvim_get_current_buf()
+
+  if vim.bo[bufnr].filetype ~= "rust" then
+    vim.notify "can only be used in rust"
+    return
+  end
+
+  local root = get_root(bufnr)
+
+  local changes = {}
+  for id, node in rust_sqlx:iter_captures(root, bufnr, 0, -1) do
+    local name = rust_sqlx.captures[id]
+    if name == "sql" then
+      -- range: { row_start [1], col_start [2], row_end [3], col_end }
+      local range = { node:range() }
+      local indentation = string.rep(" ", range[2])
+
+      local node_text = vim.treesitter.get_node_text(node, bufnr)
+      if node_text:match('r#"\n(.*\n)+?%s*"#$') == nil then
+        -- text is invalid because it does not have newlines in raw string literal
+        goto continue
+      end
+
+      local text = node_text:sub(4, -3) -- get just the inside of the raw string literal
+      -- TODO get current spaces for indentation and run pg_format with that indendation amount
+      local formatted = vim.fn.system({ "pg_format", "-L" }, text)
+      local lines = {}
+      for line in formatted:gmatch("[^\r\n]+") do
+        lines[#lines + 1] = line
+      end
+
+      for idx, line in ipairs(lines) do
+        lines[idx] = indentation .. line
+      end
+
+      -- Fix indentation of end of raw string literal as well
+      local last_line = vim.api.nvim_buf_get_lines(bufnr, range[3], range[3] + 1, true)[1]:match("^%s*(.*)")
+      -- print("LAST LINE", last_line)
+      lines[#lines + 1] = indentation .. last_line
+
+      table.insert(changes, 1, {
+        start = range[1] + 1,
+        final = range[3] + 1,
+        formatted = lines,
+      })
+
+      -- changes[#changes + 1] = {
+      --   col = 0,
+      --   row = range[1] + 2,
+      --   end_col = -1,
+      --   end_row = range[3] + 1,
+      --   text = table.concat(lines, '\n'),
+      -- }
+
+      ::continue::
+    end
+  end
+
+  -- return changes
+
+  print("Change", vim.inspect(changes))
+  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("SqlMagic", function()
+  format_dat_sql()
+end, {})
+
 return {
   {
     "jose-elias-alvarez/null-ls.nvim",
@@ -47,8 +165,22 @@ return {
         }),
       }
 
+      -- local rust_formatter_sqlx = {
+      --   name = "rust_formatter_sqlx",
+      --   method = null_ls.methods.FORMATTING,
+      --   filetypes = { "rust" },
+      --   generator = {
+      --     fn = function(params)
+      --       local changes = format_dat_sql(params.bufnr)
+      --       -- print("CHANGES:\n", vim.inspect(changes))
+      --       -- return changes
+      --     end
+      --   },
+      -- }
+
       null_ls.register(rust_formatter_genemichaels)
       null_ls.register(rust_formatter_rustfmt)
+      -- null_ls.register(rust_formatter_sqlx)
 
       -- Check supported formatters and linters
       -- https://github.com/jose-elias-alvarez/null-ls.nvim/tree/main/lua/null-ls/builtins/formatting
@@ -59,6 +191,7 @@ return {
         --null_ls.builtins.formatting.rustfmt, -- rust
         rust_formatter_genemichaels,          -- order matters, we run genemichaels first then rustfmt
         rust_formatter_rustfmt,
+        rust_formatter_sqlx,
         null_ls.builtins.formatting.black, -- python
         -- null_ls.builtins.code_actions.proselint, -- TODO looks interesting
         null_ls.builtins.code_actions.cspell.with({

queries/rust/injections.scm

@@ -0,0 +1,35 @@
+; ==== SQLX syntax highlighting
+; query macro
+(macro_invocation
+   (scoped_identifier
+     path: (identifier) @_path (#eq? @_path "sqlx")
+     name: (identifier) @_name (#any-of? @_name "query" "query_scalar"))
+
+ (token_tree
+   ; . [(raw_string_literal) @sql (string_literal) @sql])
+   . (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+ )
+
+; query_as macro
+(macro_invocation
+   (scoped_identifier
+     path: (identifier) @_path (#eq? @_path "sqlx")
+     name: (identifier) @_name (#eq? @_name "query_as"))
+
+ (token_tree
+   ; (_) . [(raw_string_literal) @sql (string_literal) @sql])
+   (_) . (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+   
+ )
+
+; query and query_as function
+(call_expression
+  (scoped_identifier
+    path: (identifier) @_path (#eq? @_path "sqlx")
+    name: (identifier) @_name (#contains? @_name "query"))
+
+ (arguments
+   ; [(raw_string_literal) @sql  (string_literal) @sql])
+   (raw_string_literal) @sql (#offset! @sql 1 0 -1 0))
+  )
+