Add Support for Per-Query-User-Information / Impersonation

go.mod

@@ -4,7 +4,7 @@ go 1.18
 
 require (
  github.com/grafana/grafana-plugin-sdk-go v0.162.0
- github.com/grafana/sqlds/v2 v2.5.1
+ github.com/grafana/sqlds/v2 v2.7.2
  github.com/pkg/errors v0.9.1
  github.com/trinodb/trino-go-client v0.312.0
 )

go.sum

@@ -206,8 +206,8 @@ github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
 github.com/grafana/grafana-plugin-sdk-go v0.162.0 h1:ij2ARWohf0IoK9yCVC1Wup4Gp6zwBq2AueVXRYsv/to=
 github.com/grafana/grafana-plugin-sdk-go v0.162.0/go.mod h1:dPhljkVno3Bg/ZYafMrR/BfYjtCRJD2hU2719Nl3QzM=
-github.com/grafana/sqlds/v2 v2.5.1 h1:tZiKWn0eGjCTZNj8CfoW3ls1f3LO/b5bp54OWH/PazQ=
-github.com/grafana/sqlds/v2 v2.5.1/go.mod h1:7lDTt2p0LU82nQ9n1K95WE8YkWATDMbPbl9ZB7r8rO4=
+github.com/grafana/sqlds/v2 v2.7.2 h1:5bkY9QO5Nc4uAWm3aHF2fHWcOSWhqJmfdJfPRHJVJoo=
+github.com/grafana/sqlds/v2 v2.7.2/go.mod h1:u5FkkJfuL6EwqesXWjV6cNw7aS5F51sCo/0Op57C8rs=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0/go.mod h1:g5qyo/la0ALbONm6Vbp88Yd8NsDy6rZz+RcrMPxvld8=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92BcuyuQ/YW4NSIpoGtfXNho=

pkg/main.go

@@ -5,7 +5,6 @@ import (
 
  "github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
  "github.com/grafana/grafana-plugin-sdk-go/backend/log"
- "github.com/grafana/sqlds/v2"
  "github.com/trinodb/grafana-trino/pkg/trino"
 )
 
@@ -19,7 +18,7 @@ func main() {
  // ID). When datasource configuration changed Dispose method will be called and
  // new datasource instance created using New factory.
  s := &trino.TrinoDatasource{}
- ds := sqlds.NewDatasource(s)
+ ds := trino.NewDatasource(s)
  ds.Completable = s
  if err := datasource.Manage("trino-datasource", ds.NewDatasource, datasource.ManageOpts{}); err != nil {
  log.DefaultLogger.Error(err.Error())

pkg/trino/datasource-context.go

@@ -0,0 +1,47 @@
+package trino
+
+import (
+ "context"
+ "fmt"
+ "github.com/grafana/grafana-plugin-sdk-go/backend"
+ "github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
+ "github.com/grafana/sqlds/v2"
+ "github.com/trinodb/grafana-trino/pkg/trino/models"
+)
+
+type SQLDatasourceWithTrinoUserContext struct {
+ sqlds.SQLDatasource
+}
+
+func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+ config := req.PluginContext.DataSourceInstanceSettings
+ settings := models.TrinoDatasourceSettings{}
+ err := settings.Load(*config)
+ if err != nil {
+ return nil, fmt.Errorf("error reading settings: %s", err.Error())
+ }
+
+ if settings.EnableImpersonation {
+ user := req.PluginContext.User
+ if user == nil {
+ return nil, fmt.Errorf("user can't be nil if impersonation is enabled")
+ }
+
+ ctx = context.WithValue(ctx, "X-Trino-User", user)
+ }
+
+ return ds.SQLDatasource.QueryData(ctx, req)
+}
+
+func (ds *SQLDatasourceWithTrinoUserContext) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+ _, err := ds.SQLDatasource.NewDatasource(settings)
+ if err != nil {
+ return nil, err
+ }
+ return ds, nil
+}
+
+func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext {
+ base := sqlds.NewDatasource(c)
+ return &SQLDatasourceWithTrinoUserContext{*base}
+}

pkg/trino/datasource.go

@@ -5,6 +5,7 @@ import (
  "database/sql"
  "encoding/json"
  "fmt"
+ "net/http"
  "regexp"
 
  "github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -21,8 +22,9 @@ type TrinoDatasource struct {
 }
 
 var (
- _ sqlds.Driver = (*TrinoDatasource)(nil)
- _ sqlds.Completable = (*TrinoDatasource)(nil)
+ _ sqlds.Driver = (*TrinoDatasource)(nil)
+ _ sqlds.QueryArgSetter = (*TrinoDatasource)(nil)
+ _ sqlds.Completable = (*TrinoDatasource)(nil)
 )
 
 func New() *TrinoDatasource {
@@ -76,6 +78,14 @@ func (s *TrinoDatasource) Converters() (sc []sqlutil.Converter) {
  }
 }
 
+func (s *TrinoDatasource) SetQueryArgs(ctx context.Context, headers http.Header) []interface{} {
+ user := ctx.Value("X-Trino-User")
+ if user == nil {
+ return []interface{}{}
+ }
+ return []interface{}{sql.Named("X-Trino-User", string(user.(*backend.User).Login))}
+}
+
 func (s *TrinoDatasource) Schemas(ctx context.Context, options sqlds.Options) ([]string, error) {
  // TBD
  return []string{}, nil

pkg/trino/models/settings.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+ "encoding/json"
  "errors"
  "net/url"
 
@@ -10,8 +11,9 @@ import (
 )
 
 type TrinoDatasourceSettings struct {
- URL *url.URL
- Opts httpclient.Options
+ URL *url.URL `json:"-"`
+ Opts httpclient.Options `json:"-"`
+ EnableImpersonation bool `json:"enableImpersonation"`
 }
 
 func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings) error {
@@ -37,5 +39,9 @@ func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings
  s.URL.User = url.User("grafana")
  }
  s.Opts = opts
+ err = json.Unmarshal(config.JSONData, &s)
+ if err != nil {
+ return err
+ }
  return nil
 }

src/ConfigEditor.tsx

@@ -1,5 +1,5 @@
-import React, { PureComponent } from 'react';
-import { DataSourceHttpSettings } from '@grafana/ui';
+import React, { ChangeEvent, PureComponent } from 'react';
+import { DataSourceHttpSettings, InlineField, InlineSwitch } from '@grafana/ui';
 import { DataSourcePluginOptionsEditorProps } from '@grafana/data';
 import { TrinoDataSourceOptions } from './types';
 
@@ -10,13 +10,32 @@ interface State {}
 export class ConfigEditor extends PureComponent<Props, State> {
  render() {
  const { options, onOptionsChange } = this.props;
+ const onEnableImpersonationChange = (event: ChangeEvent<HTMLInputElement>) => {
+ onOptionsChange({...options, jsonData: {...options.jsonData, enableImpersonation: event.target.checked}})
+ }
  return (
  <div className="gf-form-group">
  <DataSourceHttpSettings
  defaultUrl="http://localhost:8080"
  dataSourceConfig={options}
  onChange={onOptionsChange}
  />
+
+ <h3 className="page-heading">Trino</h3>
+ <div className="gf-form-group">
+ <div className="gf-form-inline">
+ <InlineField
+ label="Impersonate logged in user"
+ tooltip="If enabled, set X-Trino-User to the current Grafana user"
+ >
+ <InlineSwitch
+ id="trino-settings-enable-impersonation"
+ value={options.jsonData?.enableImpersonation ?? false}
+ onChange={onEnableImpersonationChange}
+ />
+ </InlineField>
+ </div>
+ </div>
  </div>
  );
  }

src/types.ts

@@ -46,7 +46,9 @@ ORDER BY
  * These are options configured for each DataSource instance.
  */
 
-export interface TrinoDataSourceOptions extends DataSourceJsonData {}
+export interface TrinoDataSourceOptions extends DataSourceJsonData {
+ enableImpersonation?: boolean;
+}
 /**
  * Value that is used in the backend, but never sent over HTTP to the frontend
  */