Allow to add any route (#29)

Author: andresmgot

completion.go

@@ -3,6 +3,7 @@ package sqlds
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -95,8 +96,20 @@ func (ds *sqldatasource) getColumns(rw http.ResponseWriter, req *http.Request) {
 	sendResourceResponse(rw, res)
 }
 
-func (ds *sqldatasource) registerRoutes(mux *http.ServeMux) {
-	mux.HandleFunc("/tables", ds.getTables)
-	mux.HandleFunc("/schemas", ds.getSchemas)
-	mux.HandleFunc("/columns", ds.getColumns)
+func (ds *sqldatasource) registerRoutes(mux *http.ServeMux) error {
+	defaultRoutes := map[string]func(http.ResponseWriter, *http.Request){
+		"/tables":  ds.getTables,
+		"/schemas": ds.getSchemas,
+		"/columns": ds.getColumns,
+	}
+	for route, handler := range defaultRoutes {
+		mux.HandleFunc(route, handler)
+	}
+	for route, handler := range ds.CustomRoutes {
+		if _, ok := defaultRoutes[route]; ok {
+			return fmt.Errorf("unable to redefine %s, use the Completable interface instead", route)
+		}
+		mux.HandleFunc(route, handler)
+	}
+	return nil
 }

completion_test.go

@@ -134,3 +134,50 @@ func TestCompletable(t *testing.T) {
 		})
 	}
 }
+
+func Test_registerRoutes(t *testing.T) {
+	t.Run("it should add a new route", func(t *testing.T) {
+		sqlds := &sqldatasource{}
+		sqlds.CustomRoutes = map[string]func(http.ResponseWriter, *http.Request){
+			"/foo": func(w http.ResponseWriter, r *http.Request) {
+				_, err := w.Write([]byte("bar"))
+				if err != nil {
+					t.Fatal((err))
+				}
+			},
+		}
+
+		mux := http.NewServeMux()
+		err := sqlds.registerRoutes(mux)
+		if err != nil {
+			t.Fatalf("unexpected error %v", err)
+		}
+		resp := httptest.NewRecorder()
+		req, err := http.NewRequest("GET", "/foo", nil)
+		if err != nil {
+			t.Fatalf("unexpected error %v", err)
+		}
+		mux.ServeHTTP(resp, req)
+
+		respByte, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			t.Fatalf("unexpected error %v", err)
+		}
+		if string(respByte) != "bar" {
+			t.Errorf("unexpected response %s", string(respByte))
+		}
+	})
+
+	t.Run("it error if tried to add an existing route", func(t *testing.T) {
+		sqlds := &sqldatasource{}
+		sqlds.CustomRoutes = map[string]func(http.ResponseWriter, *http.Request){
+			"/tables": func(w http.ResponseWriter, r *http.Request) {},
+		}
+
+		mux := http.NewServeMux()
+		err := sqlds.registerRoutes(mux)
+		if err == nil || err.Error() != "unable to redefine /tables, use the Completable interface instead" {
+			t.Errorf("unexpected error %v", err)
+		}
+	})
+}

datasource.go

@@ -20,6 +20,7 @@ type sqldatasource struct {
 
 	backend.CallResourceHandler
 	Completable
+	CustomRoutes map[string]func(http.ResponseWriter, *http.Request)
 }
 
 // NewDatasource creates a new `sqldatasource`.
@@ -32,7 +33,10 @@ func (ds *sqldatasource) NewDatasource(settings backend.DataSourceInstanceSettin
 	ds.db = db
 	ds.settings = settings
 	mux := http.NewServeMux()
-	ds.registerRoutes(mux)
+	err = ds.registerRoutes(mux)
+	if err != nil {
+		return nil, err
+	}
 	ds.CallResourceHandler = httpadapter.New(mux)
 
 	return ds, nil