databricks-cli/cmd/lsp/lsp.go

package lsp

import (
	"context"

	"github.com/databricks/cli/cmd/root"
	"github.com/databricks/databricks-sdk-go/httpclient"
	"github.com/spf13/cobra"
	"github.com/tliron/commonlog"
	"github.com/tliron/glsp"
	protocol "github.com/tliron/glsp/protocol_3_16"
	"github.com/tliron/glsp/server"
)

const lsName = "databricks-lsp"

var version string = "0.0.1"
var handler protocol.Handler

var localClient = httpclient.NewApiClient(httpclient.ClientConfig{})

type AnalyseResponse struct {
	Diagnostics []protocol.Diagnostic `json:"diagnostics"`
}

func callUcx(lspctx *glsp.Context, uri protocol.DocumentUri) error {
	var res AnalyseResponse
	err := localClient.Do(context.Background(), "GET", "http://localhost:8000/analyse",
		httpclient.WithRequestData(map[string]any{
			"file_uri": uri,
		}), httpclient.WithResponseUnmarshal(&res))
	if err != nil {
		return err
	}
 	lspctx.Notify(protocol.ServerTextDocumentPublishDiagnostics, &protocol.PublishDiagnosticsParams{
		URI:         uri,
		Diagnostics: res.Diagnostics,
	})
	return nil
}

func startServer(ctx context.Context) error {
	commonlog.Configure(1, nil)

	handler = protocol.Handler{
		Initialize:  initialize,
		Initialized: initialized,
		Shutdown:    shutdown,
		SetTrace:    setTrace,
		TextDocumentCodeAction: func(context *glsp.Context, params *protocol.CodeActionParams) (any, error) {
			foundUcx := false
			var codeRange protocol.Range
			for _, v := range params.Context.Diagnostics {
				if v.Source == nil {
					continue
				}
				if *v.Source == "databricks.labs.ucx" {
					codeRange = v.Range
					foundUcx = true
				}
			}
			if !foundUcx {
				return nil, nil
			}
			quickFix := protocol.CodeActionKindQuickFix
			codeActions := []protocol.CodeAction{
				{
					Title: "Replace table with migrated table",
					Kind:  &quickFix,
					Edit: &protocol.WorkspaceEdit{
						DocumentChanges: []any{
							protocol.TextDocumentEdit{
								TextDocument: protocol.OptionalVersionedTextDocumentIdentifier{
									TextDocumentIdentifier: params.TextDocument,
								},
								Edits: []any{
									protocol.TextEdit{
										Range:   codeRange,
										NewText: "[beep-v3]",
									},
								},
							},
						},
					},
				},
			}
			return codeActions, nil
		},
		CodeActionResolve: func(context *glsp.Context, params *protocol.CodeAction) (*protocol.CodeAction, error) {
			return params, nil
		},
		TextDocumentDidOpen: func(context *glsp.Context, params *protocol.DidOpenTextDocumentParams) error {
			return callUcx(context, params.TextDocument.URI)
		},
		TextDocumentDidChange: func(context *glsp.Context, params *protocol.DidChangeTextDocumentParams) error {
			return callUcx(context, params.TextDocument.URI)
		},
	}

	server := server.NewServer(&handler, lsName, false)
	return server.RunWebSocket("127.0.0.1:12345")
}

func initialize(context *glsp.Context, params *protocol.InitializeParams) (any, error) {
	capabilities := handler.CreateServerCapabilities()
	protocol.SetTraceValue(protocol.TraceValueVerbose)

	return protocol.InitializeResult{
		Capabilities: capabilities,
		ServerInfo: &protocol.InitializeResultServerInfo{
			Name:    lsName,
			Version: &version,
		},
	}, nil
}

func initialized(context *glsp.Context, params *protocol.InitializedParams) error {
	return protocol.Trace(context, protocol.MessageTypeLog, "initialized")
}

func shutdown(context *glsp.Context) error {
	protocol.SetTraceValue(protocol.TraceValueOff)
	return nil
}

func setTrace(context *glsp.Context, params *protocol.SetTraceParams) error {
	protocol.SetTraceValue(params.Value)
	return nil
}

func New() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "lsp",
		Args:  root.NoArgs,
		Short: "Start the databricks language server",
		Annotations: map[string]string{
			"template": "Databricks CLI v{{.Version}}\n",
		},
	}

	cmd.RunE = func(cmd *cobra.Command, args []string) error {
		ctx := cmd.Context()
		return startServer(ctx)
	}

	return cmd
}