Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix completion in where, order and group clause #27

Merged
merged 3 commits into from
Jul 1, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 63 additions & 24 deletions langserver/internal/source/completion.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"strings"

"cloud.google.com/go/bigquery"
"github.com/goccy/go-zetasql"
"github.com/goccy/go-zetasql/ast"
rast "github.com/goccy/go-zetasql/resolved_ast"
"github.com/goccy/go-zetasql/types"
Expand Down Expand Up @@ -52,40 +53,38 @@ func (p *Project) Complete(ctx context.Context, uri string, position lsp.Positio
termOffset := positionToByteOffset(sql.RawText, position)

parsedFile := p.ParseFile(uri, sql.RawText)
termOffset = parsedFile.fixTermOffsetForNode(termOffset)

// cursor is on table name
if node, ok := searchAstNode[*ast.TablePathExpressionNode](parsedFile.Node, termOffset); ok {
return p.completeTablePath(ctx, node)
tablePathNode, ok := searchAstNode[*ast.TablePathExpressionNode](parsedFile.Node, termOffset)
if ok && tablePathNode.ParseLocationRange().End().ByteOffset() != termOffset {
return p.completeTablePath(ctx, tablePathNode)
}

output, ok := parsedFile.findTargetAnalyzeOutput(termOffset)
output, ok := parsedFile.FindTargetAnalyzeOutput(termOffset)
if !ok {
p.logger.Debug("not found analyze output")
return nil, nil
}
incompleteColumnName := parsedFile.findIncompleteColumnName(position)
incompleteColumnName := parsedFile.FindIncompleteColumnName(position)

node, ok := searchResolvedAstNode[*rast.ProjectScanNode](output, termOffset)
if !ok {
// In some case, *rast.ProjectScanNode.ParseLocationRange() returns nil.
// So, if we cannot find *rast.ProjectScanNode, we search *rast.ProjectScanNode which ParseLocationRange returns nil.
rast.Walk(output.Statement(), func(n rast.Node) error {
sNode, ok := n.(*rast.ProjectScanNode)
if !ok {
return nil
}
lRange := sNode.ParseLocationRange()
if lRange == nil {
node = sNode
}
return nil
})
if node == nil {
return nil, nil
}
node, ok := findScanNode(output, termOffset)
if node == nil {
p.logger.Debug("not found project scan node")
return nil, nil
}

columns := node.InputScan().ColumnList()
if pScanNode, ok := node.(*rast.ProjectScanNode); ok {
node = pScanNode.InputScan()
}
if oScanNode, ok := node.(*rast.OrderByScanNode); ok {
node = oScanNode.InputScan()
}
if aScanNode, ok := node.(*rast.AggregateScanNode); ok {
node = aScanNode.InputScan()
}

columns := node.ColumnList()
for _, column := range columns {
if !strings.HasPrefix(column.Name(), incompleteColumnName) {
continue
Expand All @@ -109,11 +108,51 @@ func (p *Project) Complete(ctx context.Context, uri string, position lsp.Positio
}

// for table alias completion
result = append(result, p.completeScanField(ctx, node.InputScan(), incompleteColumnName)...)
result = append(result, p.completeScanField(ctx, node, incompleteColumnName)...)

return result, nil
}

func findScanNode(output *zetasql.AnalyzerOutput, termOffset int) (node rast.ScanNode, ok bool) {
node, ok = searchResolvedAstNode[*rast.ProjectScanNode](output, termOffset)
if ok {
return node, true
}

node, ok = searchResolvedAstNode[*rast.OrderByScanNode](output, termOffset)
if ok {
return node, true
}

// In some case, *rast.ProjectScanNode.ParseLocationRange() returns nil.
// So, if we cannot find *rast.ProjectScanNode, we search *rast.ProjectScanNode which ParseLocationRange returns nil.
rast.Walk(output.Statement(), func(n rast.Node) error {
if !n.IsScan() {
return nil
}

sNode := n.(rast.ScanNode)

lRange := n.ParseLocationRange()
if lRange == nil {
node = sNode
return nil
}
// if the cursor is on the end of the node, the cursor is out of the node
// So, we need to permit the some offset.
if lRange.End().ByteOffset() < termOffset+5 {
node = sNode
return nil
}
return nil
})

if node == nil {
return nil, false
}
return node, true
}

type tablePathParams struct {
ProjectID string
DatasetID string
Expand Down
102 changes: 100 additions & 2 deletions langserver/internal/source/completion_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/kitagry/bqls/langserver/internal/lsp"
"github.com/kitagry/bqls/langserver/internal/source"
"github.com/kitagry/bqls/langserver/internal/source/helper"
"github.com/sirupsen/logrus"
"google.golang.org/api/cloudresourcemanager/v1"
)

Expand Down Expand Up @@ -397,6 +398,100 @@ func TestProject_CompleteColumn(t *testing.T) {
},
},
},
"Complete column in where clause": {
files: map[string]string{
"file1.sql": "SELECT * FROM `project.dataset.table` WHERE |",
},
bqTableMetadataMap: map[string]*bq.TableMetadata{
"project.dataset.table": {
Schema: bq.Schema{
{
Name: "id",
Description: "id description",
Type: bq.IntegerFieldType,
},
},
},
},
expectCompletionItems: []source.CompletionItem{
{
Kind: lsp.CIKField,
NewText: "id",
Detail: "INTEGER\nid description",
},
},
},
"Complete column in group by clause": {
files: map[string]string{
"file1.sql": "SELECT * FROM `project.dataset.table` GROUP BY |",
},
bqTableMetadataMap: map[string]*bq.TableMetadata{
"project.dataset.table": {
Schema: bq.Schema{
{
Name: "id",
Description: "id description",
Type: bq.IntegerFieldType,
},
},
},
},
expectCompletionItems: []source.CompletionItem{
{
Kind: lsp.CIKField,
NewText: "id",
Detail: "INTEGER\nid description",
},
},
},
"Complete incomplete column in group by clause": {
files: map[string]string{
"file1.sql": "SELECT * FROM `project.dataset.table` GROUP BY i|",
},
bqTableMetadataMap: map[string]*bq.TableMetadata{
"project.dataset.table": {
Schema: bq.Schema{
{
Name: "id",
Description: "id description",
Type: bq.IntegerFieldType,
},
},
},
},
expectCompletionItems: []source.CompletionItem{
{
Kind: lsp.CIKField,
NewText: "id",
Detail: "INTEGER\nid description",
TypedPrefix: "i",
},
},
},
"Complete incomplete column in order by clause": {
files: map[string]string{
"file1.sql": "SELECT * FROM `project.dataset.table` ORDER BY i|",
},
bqTableMetadataMap: map[string]*bq.TableMetadata{
"project.dataset.table": {
Schema: bq.Schema{
{
Name: "id",
Description: "id description",
Type: bq.IntegerFieldType,
},
},
},
},
expectCompletionItems: []source.CompletionItem{
{
Kind: lsp.CIKField,
NewText: "id",
Detail: "INTEGER\nid description",
TypedPrefix: "i",
},
},
},
}

for n, tt := range tests {
Expand All @@ -410,7 +505,10 @@ func TestProject_CompleteColumn(t *testing.T) {
}
bqClient.EXPECT().GetTableMetadata(gomock.Any(), tablePathSplitted[0], tablePathSplitted[1], tablePathSplitted[2]).Return(schema, nil).MinTimes(0)
}
p := source.NewProjectWithBQClient("/", bqClient)
bqClient.EXPECT().ListTables(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).MinTimes(0)
logger := logrus.New()
logger.SetLevel(logrus.DebugLevel)
p := source.NewProjectWithBQClient("/", bqClient, logger)

files, path, position, err := helper.GetLspPosition(tt.files)
if err != nil {
Expand Down Expand Up @@ -585,7 +683,7 @@ func TestProject_CompleteFromClause(t *testing.T) {
for n, tt := range tests {
t.Run(n, func(t *testing.T) {
bqClient := tt.bigqueryClientMockFunc(t)
p := source.NewProjectWithBQClient("/", bqClient)
p := source.NewProjectWithBQClient("/", bqClient, logrus.New())

files, path, position, err := helper.GetLspPosition(tt.files)
if err != nil {
Expand Down
3 changes: 2 additions & 1 deletion langserver/internal/source/document.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ func (p *Project) TermDocument(uri string, position lsp.Position) ([]lsp.MarkedS
parsedFile := p.ParseFile(uri, sql.RawText)

termOffset := positionToByteOffset(sql.RawText, position)
termOffset = parsedFile.fixTermOffsetForNode(termOffset)
targetNode, ok := searchAstNode[*ast.PathExpressionNode](parsedFile.Node, termOffset)
if !ok {
p.logger.Debug("not found target node")
Expand All @@ -37,7 +38,7 @@ func (p *Project) TermDocument(uri string, position lsp.Position) ([]lsp.MarkedS
}
}

output, ok := parsedFile.findTargetAnalyzeOutput(termOffset)
output, ok := parsedFile.FindTargetAnalyzeOutput(termOffset)
if !ok {
return nil, nil
}
Expand Down
3 changes: 2 additions & 1 deletion langserver/internal/source/document_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"github.com/kitagry/bqls/langserver/internal/lsp"
"github.com/kitagry/bqls/langserver/internal/source"
"github.com/kitagry/bqls/langserver/internal/source/helper"
"github.com/sirupsen/logrus"
)

func TestProject_TermDocument(t *testing.T) {
Expand Down Expand Up @@ -347,7 +348,7 @@ json description`,
ctrl := gomock.NewController(t)
bqClient := mock_bigquery.NewMockClient(ctrl)
bqClient.EXPECT().GetTableMetadata(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(tt.bqTableMetadata, nil).MinTimes(0)
p := source.NewProjectWithBQClient("/", bqClient)
p := source.NewProjectWithBQClient("/", bqClient, logrus.New())

files, path, position, err := helper.GetLspPosition(tt.files)
if err != nil {
Expand Down
Loading