From 07d16d29fa90f263a7d8de92e36a7c96d6c2ae3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Samson?= <lukaszsamson@gmail.com>
Date: Thu, 12 Jun 2025 00:45:06 +0200
Subject: [PATCH] fix query bindings extraction

---
 .../providers/plugins/ecto/query.ex           | 16 ++--
 .../plugins/ecto/query_bindings_test.exs      | 80 +++++++++++++++++++
 2 files changed, 87 insertions(+), 9 deletions(-)
 create mode 100644 apps/language_server/test/providers/plugins/ecto/query_bindings_test.exs

diff --git a/apps/language_server/lib/language_server/providers/plugins/ecto/query.ex b/apps/language_server/lib/language_server/providers/plugins/ecto/query.ex
index 29ca7a515..4a343c3eb 100644
--- a/apps/language_server/lib/language_server/providers/plugins/ecto/query.ex
+++ b/apps/language_server/lib/language_server/providers/plugins/ecto/query.ex
@@ -240,15 +240,13 @@ defmodule ElixirLS.LanguageServer.Plugins.Ecto.Query do
 
     from_matches = Regex.scan(~r/^.+\(?\s*(#{@binding_r})/u, func_code)
 
-    # TODO this code is broken
-    # depends on join positions that we are unable to get from AST
-    # line and col was previously assigned to each option in Source.which_func
-    join_matches =
-      for join when join in @joins <- func_info.options_so_far,
-          code = Source.text_after(prefix, line, col),
-          match <- Regex.scan(~r/^#{Regex.escape(join)}\:\s*(#{@binding_r})/u, code) do
-        match
-      end
+    joins_pattern =
+      @joins
+      |> Enum.map(&Regex.escape(to_string(&1)))
+      |> Enum.join("|")
+
+    join_regex = ~r/(?:^|\s)(?:#{joins_pattern})\s*:\s*(#{@binding_r})/u
+    join_matches = Regex.scan(join_regex, func_code)
 
     matches = from_matches ++ join_matches
 
diff --git a/apps/language_server/test/providers/plugins/ecto/query_bindings_test.exs b/apps/language_server/test/providers/plugins/ecto/query_bindings_test.exs
new file mode 100644
index 000000000..99835d7cc
--- /dev/null
+++ b/apps/language_server/test/providers/plugins/ecto/query_bindings_test.exs
@@ -0,0 +1,80 @@
+defmodule ElixirLS.LanguageServer.Plugins.Ecto.QueryBindingsTest do
+  use ExUnit.Case
+
+  alias ElixirLS.LanguageServer.Plugins.Ecto.Query
+  alias ElixirSense.Core.{Source, Parser, Metadata, Binding}
+
+  defp cursor(text) do
+    {_, cursors} =
+      Source.walk_text(text, {false, []}, fn
+        "#", rest, _, _, {_comment?, cursors} -> {rest, {true, cursors}}
+        "\n", rest, _, _, {_comment?, cursors} -> {rest, {false, cursors}}
+        "^", rest, line, col, {true, cursors} -> {rest, {true, [%{line: line - 1, col: col} | cursors]}}
+        _, rest, _, _, acc -> {rest, acc}
+      end)
+
+    List.first(Enum.reverse(cursors))
+  end
+
+  defp env_and_meta(buffer, {line, col}) do
+    metadata = Parser.parse_string(buffer, true, false, {line, col})
+    {prefix, suffix} = Source.prefix_suffix(buffer, line, col)
+
+    surround =
+      case {prefix, suffix} do
+        {"", ""} -> nil
+        _ -> {{line, col - String.length(prefix)}, {line, col + String.length(suffix)}}
+      end
+
+    env = Metadata.get_cursor_env(metadata, {line, col}, surround)
+    {env, metadata}
+  end
+
+  defp extract_bindings(buffer) do
+    cur = cursor(buffer)
+    {env, meta} = env_and_meta(buffer, {cur.line, cur.col})
+    prefix = Source.text_before(buffer, cur.line, cur.col)
+    binding_env = Binding.from_env(env, meta, {cur.line, cur.col})
+    func_info = Source.which_func(prefix, binding_env)
+    Query.extract_bindings(prefix, func_info, env, meta)
+  end
+
+  test "extract binding from from clause" do
+    buffer = """
+    import Ecto.Query
+    alias ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Post
+
+    from p in Post,
+      where: true
+      #       ^
+    """
+
+    assert %{"p" => %{type: ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Post}} =
+             extract_bindings(buffer)
+  end
+
+  test "extract bindings from join clauses" do
+    buffer = """
+    import Ecto.Query
+    alias ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Post
+    alias ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Comment
+    alias ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.User
+
+    from(
+      p in Post,
+      join: c in Comment,
+      left_join: u in assoc(p, :user),
+      where: true
+      #       ^
+    )
+    """
+
+    result = extract_bindings(buffer)
+
+    assert %{
+             "p" => %{type: ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Post},
+             "c" => %{type: ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.Comment},
+             "u" => %{type: ElixirLS.LanguageServer.Plugins.Ecto.FakeSchemas.User}
+           } = result
+  end
+end