diff --git a/langserver/execute_command.go b/langserver/execute_command.go index 0ddd474..044204b 100644 --- a/langserver/execute_command.go +++ b/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, diff --git a/langserver/internal/bigquery/bigquery.go b/langserver/internal/bigquery/bigquery.go index 9945218..832057d 100644 --- a/langserver/internal/bigquery/bigquery.go +++ b/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) diff --git a/langserver/internal/bigquery/cache.go b/langserver/internal/bigquery/cache.go index e0f71d8..c73df27 100644 --- a/langserver/internal/bigquery/cache.go +++ b/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) } diff --git a/langserver/internal/source/project.go b/langserver/internal/source/project.go index f28be3a..7d7a83a 100644 --- a/langserver/internal/source/project.go +++ b/langserver/internal/source/project.go @@ -1,7 +1,9 @@ package source import ( + "bytes" "context" + "encoding/csv" "fmt" "os/exec" "strings" @@ -9,16 +11,19 @@ import ( 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 +} diff --git a/langserver/langserver.go b/langserver/langserver.go index 17aba7e..23a1c61 100644 --- a/langserver/langserver.go +++ b/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)} } diff --git a/langserver/virtual_text_document.go b/langserver/virtual_text_document.go new file mode 100644 index 0000000..b6ce422 --- /dev/null +++ b/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, ¶ms); 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 +} diff --git a/langserver/virtual_text_document_test.go b/langserver/virtual_text_document_test.go new file mode 100644 index 0000000..94e42d6 --- /dev/null +++ b/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) + } + }) + } +}