feat: add virtualTextDocument

langserver/execute_command.go

@@ -120,7 +120,7 @@ func (h *Handler) commandExecuteQuery(ctx context.Context, params lsp.ExecuteCom
 
 	return &ExecuteQueryResult{
 		TextDocument: lsp.TextDocumentIdentifier{
-			URI: lsp.DocumentURI(fmt.Sprintf("bqls://job/%s", job.ID())),
+			URI: newJobVirtualTextDocumentURI(h.project.BigQueryProjectID, job.ID()),
 		},
 		Result: QueryResult{
 			Columns: columns,

langserver/internal/bigquery/bigquery.go

@@ -27,6 +27,9 @@ type Client interface {
 	// GetTableMetadata returns the metadata of the specified table.
 	GetTableMetadata(ctx context.Context, projectID, datasetID, tableID string) (*bigquery.TableMetadata, error)
 
+	// GetTableRecord returns the row of the specified table.
+	GetTableRecord(ctx context.Context, projectID, datasetID, tableID string) (*bigquery.RowIterator, error)
+
 	// Run runs the specified query.
 	Run(ctx context.Context, q string, dryrun bool) (BigqueryJob, error)
 }
@@ -141,6 +144,10 @@ func (c *client) GetTableMetadata(ctx context.Context, projectID, datasetID, tab
 	return md, nil
 }
 
+func (c *client) GetTableRecord(ctx context.Context, projectID, datasetID, tableID string) (*bigquery.RowIterator, error) {
+	return c.bqClient.DatasetInProject(projectID, datasetID).Table(tableID).Read(ctx), nil
+}
+
 type BigqueryJob interface {
 	ID() string
 	Read(context.Context) (*bigquery.RowIterator, error)

langserver/internal/bigquery/cache.go

@@ -197,6 +197,10 @@ func (c *cache) GetTableMetadata(ctx context.Context, projectID, datasetID, tabl
 	return result, nil
 }
 
+func (c *cache) GetTableRecord(ctx context.Context, projectID, datasetID, tableID string) (*bigquery.RowIterator, error) {
+	return c.bqClient.GetTableRecord(ctx, projectID, datasetID, tableID)
+}
+
 func (c *cache) Run(ctx context.Context, q string, dryrun bool) (BigqueryJob, error) {
 	return c.bqClient.Run(ctx, q, dryrun)
 }

langserver/internal/source/project.go

@@ -1,24 +1,29 @@
 package source
 
 import (
+	"bytes"
 	"context"
+	"encoding/csv"
 	"fmt"
 	"os/exec"
 	"strings"
 
 	bq "cloud.google.com/go/bigquery"
 	"github.com/kitagry/bqls/langserver/internal/bigquery"
 	"github.com/kitagry/bqls/langserver/internal/cache"
+	"github.com/kitagry/bqls/langserver/internal/lsp"
 	"github.com/kitagry/bqls/langserver/internal/source/file"
 	"github.com/sirupsen/logrus"
+	"google.golang.org/api/iterator"
 )
 
 type Project struct {
-	rootPath string
-	logger   *logrus.Logger
-	cache    *cache.GlobalCache
-	bqClient bigquery.Client
-	analyzer *file.Analyzer
+	BigQueryProjectID string
+	rootPath          string
+	logger            *logrus.Logger
+	cache             *cache.GlobalCache
+	bqClient          bigquery.Client
+	analyzer          *file.Analyzer
 }
 
 type File struct {
@@ -50,11 +55,12 @@ func NewProject(ctx context.Context, rootPath string, projectID string, logger *
 	analyzer := file.NewAnalyzer(logger, bqClient)
 
 	return &Project{
-		rootPath: rootPath,
-		logger:   logger,
-		cache:    cache,
-		bqClient: bqClient,
-		analyzer: analyzer,
+		BigQueryProjectID: projectID,
+		rootPath:          rootPath,
+		logger:            logger,
+		cache:             cache,
+		bqClient:          bqClient,
+		analyzer:          analyzer,
 	}, nil
 }
 
@@ -143,3 +149,55 @@ func (p *Project) ListDatasets(ctx context.Context, projectID string) ([]*bq.Dat
 func (p *Project) ListTables(ctx context.Context, projectID, datasetID string) ([]*bq.Table, error) {
 	return p.bqClient.ListTables(ctx, projectID, datasetID, true)
 }
+
+func (p *Project) GetTableInfo(ctx context.Context, projectID, datasetID, tableID string) ([]lsp.MarkedString, error) {
+	tableMetadata, err := p.bqClient.GetTableMetadata(ctx, projectID, datasetID, tableID)
+	if err != nil {
+		return nil, err
+	}
+
+	result, err := buildBigQueryTableMetadataMarkedString(tableMetadata)
+	if err != nil {
+		return nil, err
+	}
+
+	it, err := p.bqClient.GetTableRecord(ctx, projectID, datasetID, tableID)
+	if err != nil {
+		return result, err
+	}
+
+	data := make([][]string, 0)
+	for i := 0; i < 100; i++ {
+		var values []bq.Value
+		err := it.Next(&values)
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return result, err
+		}
+
+		strs := make([]string, len(values))
+		for i, v := range values {
+			strs[i] = fmt.Sprint(v)
+		}
+		data = append(data, strs)
+	}
+
+	columns := make([]string, 0)
+	for _, f := range it.Schema {
+		columns = append(columns, f.Name)
+	}
+
+	var buf bytes.Buffer
+	cw := csv.NewWriter(&buf)
+	cw.Write(columns)
+	cw.WriteAll(data)
+
+	result = append(result, lsp.MarkedString{
+		Language: "csv",
+		Value:    buf.String(),
+	})
+
+	return result, nil
+}

langserver/langserver.go

@@ -93,6 +93,8 @@ func (h *Handler) handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2
 		return h.handleTextDocumentCodeAction(ctx, conn, req)
 	case "workspace/executeCommand":
 		return h.handleWorkspaceExecuteCommand(ctx, conn, req)
+	case "bqls/virtualTextDocument":
+		return h.handleVirtualTextDocument(ctx, conn, req)
 	}
 	return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeMethodNotFound, Message: fmt.Sprintf("method not supported: %s", req.Method)}
 }

langserver/virtual_text_document.go

@@ -0,0 +1,121 @@
+package langserver
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+	"github.com/sourcegraph/jsonrpc2"
+)
+
+type VirtualTextDocumentParams struct {
+	TextDocument lsp.TextDocumentIdentifier `json:"textDocument"`
+}
+
+type VirtualTextDocument struct {
+	Contents []lsp.MarkedString `json:"contents"`
+}
+
+type VirtualTextDocumentInfo struct {
+	ProjectID string
+	DatasetID string
+	TableID   string
+	JobID     string
+}
+
+func ParseVirtualTextDocument(textDocument lsp.DocumentURI) (VirtualTextDocumentInfo, error) {
+	suffix, ok := strings.CutPrefix(string(textDocument), "bqls://")
+	if !ok {
+		return VirtualTextDocumentInfo{}, errors.New("invalid text document URI")
+	}
+
+	result := VirtualTextDocumentInfo{}
+	for suffix != "" {
+		for _, prefix := range []string{"project/", "dataset/", "table/", "job/"} {
+			after, ok := strings.CutPrefix(suffix, prefix)
+			if !ok {
+				continue
+			}
+
+			ind := strings.Index(after, "/")
+			var val string
+			if ind == -1 {
+				val = after
+				suffix = ""
+			} else {
+				val = after[:ind]
+				suffix = after[ind+1:]
+			}
+
+			if prefix == "project/" {
+				result.ProjectID = val
+			} else if prefix == "dataset/" {
+				result.DatasetID = val
+			} else if prefix == "table/" {
+				result.TableID = val
+			} else if prefix == "job/" {
+				result.JobID = val
+			}
+		}
+	}
+
+	if err := result.validate(); err != nil {
+		return VirtualTextDocumentInfo{}, err
+	}
+
+	return result, nil
+}
+
+func (v VirtualTextDocumentInfo) validate() error {
+	if v.ProjectID == "" {
+		return errors.New("project ID is required")
+	}
+
+	if v.DatasetID != "" && v.TableID != "" {
+		return nil
+	}
+
+	if v.JobID != "" {
+		return nil
+	}
+
+	return fmt.Errorf("either dataset ID and table ID or job ID is required")
+}
+
+func newJobVirtualTextDocumentURI(projectID, jobID string) lsp.DocumentURI {
+	return lsp.DocumentURI(fmt.Sprintf("bqls://project/%s/job/%s", projectID, jobID))
+}
+
+func newTableVirtualTextDocumentURI(projectID, datasetID, tableID string) lsp.DocumentURI {
+	return lsp.DocumentURI(fmt.Sprintf("bqls://project/%s/dataset/%s/table/%s", projectID, datasetID, tableID))
+}
+
+func (h *Handler) handleVirtualTextDocument(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (result any, err error) {
+	if req.Params == nil {
+		return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeInvalidParams}
+	}
+
+	var params VirtualTextDocumentParams
+	if err := json.Unmarshal(*req.Params, &params); err != nil {
+		return nil, err
+	}
+
+	virtualTextDocument, err := ParseVirtualTextDocument(params.TextDocument.URI)
+	if err != nil {
+		return nil, err
+	}
+
+	if virtualTextDocument.TableID != "" {
+		marks, err := h.project.GetTableInfo(ctx, virtualTextDocument.ProjectID, virtualTextDocument.DatasetID, virtualTextDocument.TableID)
+		if err != nil {
+			return nil, err
+		}
+
+		return VirtualTextDocument{Contents: marks}, nil
+	}
+
+	return nil, nil
+}

langserver/virtual_text_document_test.go

@@ -0,0 +1,46 @@
+package langserver_test
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/kitagry/bqls/langserver"
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+)
+
+func TestParseVirtualTextDocument(t *testing.T) {
+	tests := map[string]struct {
+		uri         string
+		expected    langserver.VirtualTextDocumentInfo
+		expectedErr error
+	}{
+		"Parse project/dataset/table": {
+			uri: "bqls://project/p/dataset/d/table/t",
+			expected: langserver.VirtualTextDocumentInfo{
+				ProjectID: "p",
+				DatasetID: "d",
+				TableID:   "t",
+			},
+		},
+		"Parse project job": {
+			uri: "bqls://project/p/job/j",
+			expected: langserver.VirtualTextDocumentInfo{
+				ProjectID: "p",
+				JobID:     "j",
+			},
+		},
+	}
+
+	for n, tt := range tests {
+		t.Run(n, func(t *testing.T) {
+			got, err := langserver.ParseVirtualTextDocument(lsp.DocumentURI(tt.uri))
+			if err != tt.expectedErr {
+				t.Fatalf("ParseVirtualTextDocument error expected %v, got %v", tt.expectedErr, err)
+			}
+
+			if diff := cmp.Diff(tt.expected, got); diff != "" {
+				t.Errorf("ParseVirtualTextDocument result diff (-expect, +got)\n%s", diff)
+			}
+		})
+	}
+}