From 5ed657277dc42cf580f8f007cd2e8170692f626a Mon Sep 17 00:00:00 2001
From: Felipe Contreras Salinas <felipe@bstr.cl>
Date: Mon, 20 Nov 2023 23:44:28 -0300
Subject: [PATCH] treesitter: add playground and first attempt at sql in rust

---
 .gitmodules                   |  3 ++
 after/plugin/rust_sql.lua     | 83 +++++++++++++++++++++++++++++++++++
 init.lua                      | 24 ++++++++++
 pack/general/start/playground |  1 +
 queries/rust/injections.scm   | 10 +++++
 5 files changed, 121 insertions(+)
 create mode 100644 after/plugin/rust_sql.lua
 create mode 160000 pack/general/start/playground
 create mode 100644 queries/rust/injections.scm

diff --git a/.gitmodules b/.gitmodules
index e706fed..74d94fb 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -146,3 +146,6 @@
 [submodule "pack/general/opt/core.nvim"]
 	path = pack/general/opt/core.nvim
 	url = https://github.com/niuiic/core.nvim
+[submodule "pack/general/start/playground"]
+	path = pack/general/start/playground
+	url = https://github.com/nvim-treesitter/playground
diff --git a/after/plugin/rust_sql.lua b/after/plugin/rust_sql.lua
new file mode 100644
index 0000000..144d93e
--- /dev/null
+++ b/after/plugin/rust_sql.lua
@@ -0,0 +1,83 @@
+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, {})
diff --git a/init.lua b/init.lua
index 3272e61..9c79a41 100644
--- a/init.lua
+++ b/init.lua
@@ -159,6 +159,7 @@ require("nvim-treesitter.configs").setup({
 		"php",
 		"proto",
 		"python",
+		"query",
 		"regex",
 		"ron",
 		"rst",
@@ -174,6 +175,29 @@ require("nvim-treesitter.configs").setup({
 		"yaml",
 	},
 	highlight = { enable = true, indent = true },
+	playground = {
+		enable = true,
+		disable = {},
+		updatetime = 25,       -- Debounced time for highlighting nodes in the playground from source code
+		persist_queries = false, -- Whether the query persists across vim sessions
+		keybindings = {
+			toggle_query_editor = 'o',
+			toggle_hl_groups = 'i',
+			toggle_injected_languages = 't',
+			toggle_anonymous_nodes = 'a',
+			toggle_language_display = 'I',
+			focus_language = 'f',
+			unfocus_language = 'F',
+			update = 'R',
+			goto_node = '<cr>',
+			show_help = '?',
+		},
+	},
+	query_linter = {
+		enable = true,
+		use_virtual_text = true,
+		lint_events = { "BufWrite", "CursorHold" },
+	},
 })
 
 ---- Filetypes ---
diff --git a/pack/general/start/playground b/pack/general/start/playground
new file mode 160000
index 0000000..ba48c6a
--- /dev/null
+++ b/pack/general/start/playground
@@ -0,0 +1 @@
+Subproject commit ba48c6a62a280eefb7c85725b0915e021a1a0749
diff --git a/queries/rust/injections.scm b/queries/rust/injections.scm
new file mode 100644
index 0000000..1bfd35b
--- /dev/null
+++ b/queries/rust/injections.scm
@@ -0,0 +1,10 @@
+; Inject into sqlx::query!(r#"..."#, ...) as sql
+(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)
+)