some multiplexing

cmd/lsp/lsp.go

@@ -2,6 +2,8 @@ package lsp
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/databricks-sdk-go/httpclient"
@@ -17,37 +19,171 @@ const lsName = "databricks-lsp"
 var version string = "0.0.1"
 var handler protocol.Handler
 
-var localClient = httpclient.NewApiClient(httpclient.ClientConfig{})
+type LspThingy interface {
-
+	Match(uri protocol.DocumentUri) bool
-type AnalyseResponse struct {
-	Diagnostics []protocol.Diagnostic `json:"diagnostics"`
 }
 
-func callUcx(lspctx *glsp.Context, uri protocol.DocumentUri) error {
+type Linter interface {
-	var res AnalyseResponse
+	Lint(ctx context.Context, uri protocol.DocumentUri) ([]protocol.Diagnostic, error)
-	err := localClient.Do(context.Background(), "GET", "http://localhost:8000/analyse",
+}
+
+type QuickFixer interface {
+	QuickFix(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error)
+}
+
+type LspMultiplexer struct {
+	things []LspThingy
+}
+
+func (m *LspMultiplexer) Lint(ctx context.Context, uri protocol.DocumentUri) ([]protocol.Diagnostic, error) {
+	diags := []protocol.Diagnostic{}
+	for _, thing := range m.things {
+		if !thing.Match(uri) {
+			continue
+		}
+		linter, ok := thing.(Linter)
+		if !ok {
+			continue
+		}
+		problems, err := linter.Lint(ctx, uri)
+		if err != nil {
+			return nil, fmt.Errorf("linter: %w", err)
+		}
+		diags = append(diags, problems...)
+	}
+	return diags, nil
+}
+
+func (m *LspMultiplexer) QuickFix(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
+	actions := []protocol.CodeAction{}
+	for _, thing := range m.things {
+		if !thing.Match(params.TextDocument.URI) {
+			continue
+		}
+		fixer, ok := thing.(QuickFixer)
+		if !ok {
+			continue
+		}
+		fixes, err := fixer.QuickFix(ctx, params)
+		if err != nil {
+			return nil, fmt.Errorf("quick fixer: %w", err)
+		}
+		actions = append(actions, fixes...)
+	}
+	return actions, nil
+}
+
+type LocalLspProxy struct {
+	host       string
+	source     string
+	extensions []string
+	client     *httpclient.ApiClient
+}
+
+func (p *LocalLspProxy) Match(uri protocol.DocumentUri) bool {
+	for _, ext := range p.extensions {
+		if strings.HasSuffix(string(uri), ext) {
+			return true
+		}
+	}
+	return false
+}
+
+func (p *LocalLspProxy) Lint(ctx context.Context, uri protocol.DocumentUri) ([]protocol.Diagnostic, error) {
+	var res struct {
+		Diagnostics []protocol.Diagnostic `json:"diagnostics"`
+	}
+	err := p.client.Do(ctx, "GET", fmt.Sprintf("%s/lint", p.host),
 		httpclient.WithRequestData(map[string]any{
 			"file_uri": uri,
 		}), httpclient.WithResponseUnmarshal(&res))
 	if err != nil {
-		return err
+		return nil, err
+	}
+	return res.Diagnostics, nil
+}
+
+type FixMe struct {
+	Range protocol.Range `json:"range"`
+	Code  string         `json:"code"`
+
+	resolves protocol.Diagnostic `json:"-"`
+}
+
+// match diagnostics produced by a given source
+func (p *LocalLspProxy) matchDiagnostic(diagnostics []protocol.Diagnostic) *FixMe {
+	for _, v := range diagnostics {
+		if v.Source == nil {
+			continue
+		}
+		if *v.Source != p.source {
+			continue
+		}
+		if v.Code == nil {
+			continue
+		}
+		return &FixMe{
+			Range:    v.Range,
+			Code:     fmt.Sprint(v.Code.Value),
+			resolves: v,
+		}
 	}
- 	lspctx.Notify(protocol.ServerTextDocumentPublishDiagnostics, &protocol.PublishDiagnosticsParams{
-		URI:         uri,
-		Diagnostics: res.Diagnostics,
-	})
 	return nil
 }
 
+func (p *LocalLspProxy) QuickFix(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
+	fixMe := p.matchDiagnostic(params.Context.Diagnostics)
+	if fixMe == nil {
+		return nil, nil
+	}
+	var res struct {
+		Actions []protocol.CodeAction `json:"actions"`
+	}
+	err := p.client.Do(ctx, "POST", fmt.Sprintf("%s/quickfix", p.host),
+		httpclient.WithRequestData(map[string]any{
+			"file_uri": params.TextDocument.URI,
+			"code":     fixMe.Code,
+			"range":    params.Range,
+		}), httpclient.WithResponseUnmarshal(&res))
+	if err != nil {
+		return nil, err
+	}
+	// protocol.CodeActionKindSource has to be handled by a separate method, not QuickFix(...) - e.g reformatting
+	quickFixKind := protocol.CodeActionKindQuickFix
+	for i := range res.Actions {
+		res.Actions[i].Diagnostics = []protocol.Diagnostic{fixMe.resolves}
+		res.Actions[i].Kind = &quickFixKind
+	}
+	return res.Actions, nil
+}
+
 func startServer(ctx context.Context) error {
 	commonlog.Configure(1, nil)
 
+	// in production, we'll launch Databricks Labs command proxy, that
+	// will return a JSON on stdout with the following structure:
+	// {
+	//   "host": "http://localhost:<random-port>",
+	//   "source": "databricks.labs.<project-name>",
+	//   "extensions": [".py", <other-extensions>]
+	// }
+	ucx := &LocalLspProxy{
+		host:       "http://localhost:8000",
+		source:     "databricks.labs.ucx",
+		extensions: []string{".py", ".sql"},
+		client:     httpclient.NewApiClient(httpclient.ClientConfig{}),
+	}
+	// and here we'll add DABs, DLT, linters, more SQL introspection, etc
+	multiplexer := &LspMultiplexer{
+		things: []LspThingy{ucx},
+	}
 	handler = protocol.Handler{
 		Initialize:  initialize,
 		Initialized: initialized,
 		Shutdown:    shutdown,
 		SetTrace:    setTrace,
 		TextDocumentCodeAction: func(context *glsp.Context, params *protocol.CodeActionParams) (any, error) {
+			return multiplexer.QuickFix(ctx, params)
 			foundUcx := false
 			var codeRange protocol.Range
 			for _, v := range params.Context.Diagnostics {
@@ -89,11 +225,33 @@ func startServer(ctx context.Context) error {
 		CodeActionResolve: func(context *glsp.Context, params *protocol.CodeAction) (*protocol.CodeAction, error) {
 			return params, nil
 		},
-		TextDocumentDidOpen: func(context *glsp.Context, params *protocol.DidOpenTextDocumentParams) error {
+		TextDocumentDidOpen: func(lsp *glsp.Context, params *protocol.DidOpenTextDocumentParams) error {
-			return callUcx(context, params.TextDocument.URI)
+			problems, err := multiplexer.Lint(ctx, params.TextDocument.URI)
+			if err != nil {
+				return err
+			}
+			if len(problems) == 0 {
+				return nil
+			}
+			lsp.Notify(protocol.ServerTextDocumentPublishDiagnostics, &protocol.PublishDiagnosticsParams{
+				URI:         params.TextDocument.URI,
+				Diagnostics: problems,
+			})
+			return nil
 		},
-		TextDocumentDidChange: func(context *glsp.Context, params *protocol.DidChangeTextDocumentParams) error {
+		TextDocumentDidChange: func(lsp *glsp.Context, params *protocol.DidChangeTextDocumentParams) error {
-			return callUcx(context, params.TextDocument.URI)
+			problems, err := multiplexer.Lint(ctx, params.TextDocument.URI)
+			if err != nil {
+				return err
+			}
+			if len(problems) == 0 {
+				return nil
+			}
+			lsp.Notify(protocol.ServerTextDocumentPublishDiagnostics, &protocol.PublishDiagnosticsParams{
+				URI:         params.TextDocument.URI,
+				Diagnostics: problems,
+			})
+			return nil
 		},
 	}
 

go.mod

@@ -35,6 +35,9 @@ require (
 	gopkg.in/yaml.v3 v3.0.1
 )
 
+// until https://github.com/tliron/glsp/pull/26 gets merged
+replace github.com/tliron/glsp => github.com/nfx/glsp v0.2.3-0.20240319102602-78d58dda21a7
+
 require (
 	cloud.google.com/go/compute v1.23.4 // indirect
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect

go.sum

@@ -135,6 +135,8 @@ github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWV
 github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo=
 github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8=
+github.com/nfx/glsp v0.2.3-0.20240319102602-78d58dda21a7 h1:pfkYSCgBf0hNDDwuqk9G9J6+bFkw2/vJZzCMZm96HUA=
+github.com/nfx/glsp v0.2.3-0.20240319102602-78d58dda21a7/go.mod h1:GMVWDNeODxHzmDPvYbYTCs7yHVaEATfYtXiYJ9w1nBg=
 github.com/nwidger/jsoncolor v0.3.2 h1:rVJJlwAWDJShnbTYOQ5RM7yTA20INyKXlJ/fg4JMhHQ=
 github.com/nwidger/jsoncolor v0.3.2/go.mod h1:Cs34umxLbJvgBMnVNVqhji9BhoT/N/KinHqZptQ7cf4=
 github.com/petermattis/goid v0.0.0-20180202154549-b0b1615b78e5 h1:q2e307iGHPdTGp0hoxKjt1H5pDo6utceo3dQVK3I5XQ=
@@ -178,8 +180,6 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/tliron/commonlog v0.2.8 h1:vpKrEsZX4nlneC9673pXpeKqv3cFLxwpzNEZF1qiaQQ=
 github.com/tliron/commonlog v0.2.8/go.mod h1:HgQZrJEuiKLLRvUixtPWGcmTmWWtKkCtywF6x9X5Spw=
-github.com/tliron/glsp v0.2.1 h1:QS1c22YO1EiY0YZmJXca4Eq13gP0HX2vgy2t38gLNJ0=
-github.com/tliron/glsp v0.2.1/go.mod h1:GMVWDNeODxHzmDPvYbYTCs7yHVaEATfYtXiYJ9w1nBg=
 github.com/tliron/kutil v0.3.11 h1:kongR0dhrrn9FR/3QRFoUfQe27t78/xQvrU9aXIy5bk=
 github.com/tliron/kutil v0.3.11/go.mod h1:4IqOAAdpJuDxYbJxMv4nL8LSH0mPofSrdwIv8u99PDc=
 github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM=