fix: Fix definition

Author: lighttiger2505

internal/handler/definition.go

@@ -52,7 +52,7 @@ func definition(url, text string, params lsp.DefinitionParams, dbCache *database
 		return nil, nil
 	}
 
-	aliases := parseutil.ExtractAliasedIdentifer(parsed)
+	aliases := parseutil.ExtractAliased(parsed)
 	if len(aliases) == 0 {
 		return nil, nil
 	}
@@ -62,9 +62,14 @@ func definition(url, text string, params lsp.DefinitionParams, dbCache *database
 		alias, _ := v.(*ast.Aliased)
 		if alias.AliasedName.String() == currentVariable.String() {
 			define = alias.AliasedName
+			break
 		}
 	}
 
+	if define == nil {
+		return nil, nil
+	}
+
 	res := []lsp.Location{
 		{
 			URI: url,

internal/handler/definition_test.go

@@ -10,33 +10,79 @@ import (
 )
 
 var definitionTestCases = []struct {
-	name   string
-	input  string
-	output lsp.Definition
-	pos    lsp.Position
+	name  string
+	input string
+	pos   lsp.Position
+	want  lsp.Definition
 }{
 	{
-		name:  "",
-		input: "SELECT ci.ID, ci.Name FROM city as ci",
-		output: []lsp.Location{
+		name:  "subquery",
+		input: "SELECT it.ID, it.Name FROM (SELECT ci.ID, ci.Name, ci.CountryCode, ci.District, ci.Population FROM city AS ci) as it",
+		pos: lsp.Position{
+			Line:      0,
+			Character: 8,
+		},
+		want: []lsp.Location{
 			{
 				URI: testFileURI,
 				Range: lsp.Range{
 					Start: lsp.Position{
 						Line:      0,
-						Character: 35,
+						Character: 114,
 					},
 					End: lsp.Position{
 						Line:      0,
-						Character: 37,
+						Character: 116,
 					},
 				},
 			},
 		},
+	},
+	{
+		name:  "inner subquery",
+		input: "SELECT it.ID, it.Name FROM (SELECT ci.ID, ci.Name, ci.CountryCode, ci.District, ci.Population FROM city AS ci) as it",
+		pos: lsp.Position{
+			Line:      0,
+			Character: 36,
+		},
+		want: []lsp.Location{
+			{
+				URI: testFileURI,
+				Range: lsp.Range{
+					Start: lsp.Position{
+						Line:      0,
+						Character: 107,
+					},
+					End: lsp.Position{
+						Line:      0,
+						Character: 109,
+					},
+				},
+			},
+		},
+	},
+	{
+		name:  "alias",
+		input: "SELECT ci.ID, ci.Name FROM city as ci",
 		pos: lsp.Position{
 			Line:      0,
 			Character: 8,
 		},
+		want: []lsp.Location{
+			{
+				URI: testFileURI,
+				Range: lsp.Range{
+					Start: lsp.Position{
+						Line:      0,
+						Character: 35,
+					},
+					End: lsp.Position{
+						Line:      0,
+						Character: 37,
+					},
+				},
+			},
+		},
 	},
 }
 
@@ -71,7 +117,7 @@ func TestDefinition(t *testing.T) {
 				return
 			}
 
-			if diff := cmp.Diff(tt.output, got); diff != "" {
+			if diff := cmp.Diff(tt.want, got); diff != "" {
 				t.Errorf("unmatch hover contents (- want, + got):\n%s", diff)
 			}
 		})
@@ -109,7 +155,7 @@ func TestTypeDefinition(t *testing.T) {
 				return
 			}
 
-			if diff := cmp.Diff(tt.output, got); diff != "" {
+			if diff := cmp.Diff(tt.want, got); diff != "" {
 				t.Errorf("unmatch hover contents (- want, + got):\n%s", diff)
 			}
 		})

parser/parseutil/extract.go

@@ -119,6 +119,13 @@ func ExtractWhereCondition(parsed ast.TokenList) []ast.Node {
 	return filterPrefixGroup(astutil.NewNodeReader(parsed), prefixMatcher, peekMatcher)
 }
 
+func ExtractAliased(parsed ast.TokenList) []ast.Node {
+	reader := astutil.NewNodeReader(parsed)
+	matcher := astutil.NodeMatcher{NodeTypes: []ast.NodeType{ast.TypeAliased}}
+	aliases := reader.FindRecursive(matcher)
+	return aliases
+}
+
 func ExtractAliasedIdentifer(parsed ast.TokenList) []ast.Node {
 	reader := astutil.NewNodeReader(parsed)
 	matcher := astutil.NodeMatcher{NodeTypes: []ast.NodeType{ast.TypeAliased}}