From 8a6420bf92e56580cab4877c99956f6d2f8de0ad Mon Sep 17 00:00:00 2001
From: Felipe Contreras Salinas <felipe@bstr.cl>
Date: Tue, 21 Nov 2023 11:39:42 -0300
Subject: [PATCH] sqlformat: fix text insertion and job handling

---
 after/plugin/rust_sql.lua | 51 ++++++++++++++++++++++++---------------
 1 file changed, 32 insertions(+), 19 deletions(-)

diff --git a/after/plugin/rust_sql.lua b/after/plugin/rust_sql.lua
index 144d93e..3837083 100644
--- a/after/plugin/rust_sql.lua
+++ b/after/plugin/rust_sql.lua
@@ -1,22 +1,23 @@
 local run_formatter = function(text)
-	vim.notify(text)
-	vim.notify("----")
-	local result = {}
-	require("plenary.job"):new({
+	local inner_text = string.sub(text, 4, -3)
+	local job_result = require("plenary.job"):new({
 		command = "sqlfluff",
-		cwd = "/usr/bin",
 		args = { "format", "-" },
-		writer = text,
-		on_stdout = function(_, line)
-			table.insert(result, line)
-		end,
+		writer = inner_text,
 	}):sync()
 
+	if job_result == nil then
+		return nil
+	end
 	-- add the surrounding r#"..."# back
-	for _, line in ipairs(result) do
-		vim.notify(line)
+	if #job_result == 1 then
+		return { "r#\"" .. job_result[1] .. "#\"" }
 	end
-	-- table.insert(result, "\"#")
+	local result = { "r#\"" }
+	for _, line in ipairs(job_result) do
+		table.insert(result, line)
+	end
+	table.insert(result, "\"#")
 	return result
 end
 
@@ -28,9 +29,7 @@ local embedded_sql = vim.treesitter.query.parse(
 				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)
+			(token_tree (raw_string_literal) @sql)
 		)
 	]]
 )
@@ -60,21 +59,35 @@ local format = function(bufnr)
 			local indentation = string.rep(" ", range[2])
 			-- run the formatter on the node text
 			local formatted = run_formatter(vim.treesitter.get_node_text(node, bufnr))
+			if formatted == nil then
+				goto continue
+			end
 			-- add indentation
 			for idx, line in ipairs(formatted) do
-				formatted[idx] = indentation .. line
+				if idx > 1 then
+					formatted[idx] = indentation .. line
+				end
 			end
 			-- add changes in reverse order
 			table.insert(changes, 1, {
-				start = range[1],
-				final = range[3],
+				start_row = range[1],
+				start_col = range[2],
+				end_row = range[3],
+				end_col = range[4],
 				formatted = formatted,
 			})
+			::continue::
 		end
 	end
 
 	for _, change in ipairs(changes) do
-		vim.api.nvim_buf_set_lines(bufnr, change.start, change.final, false, change.formatted)
+		vim.api.nvim_buf_set_text(
+			bufnr,
+			change.start_row,
+			change.start_col,
+			change.end_row,
+			change.end_col,
+			change.formatted)
 	end
 end