Skip to content

Commit 165235a

Browse files
RotRotAlRotRotAl
committed
support client tags in query editor
1 parent 68edabe commit 165235a

File tree

14 files changed

+393
-97
lines changed

14 files changed

+393
-97
lines changed

‎.github/workflows/ci.yml‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,8 @@ jobs:
8686
docker run --rm --detach \
8787
--name trino \
8888
--net trino \
89-
--volume "$(pwd)/test-data/test-trino-config.properties:/etc/trino/config.properties" \
89+
--volume "$(pwd)/test-data/trino/test-trino-config.properties:/etc/trino/config.properties" \
90+
--volume "$(pwd)/test-data/trino/catalog/hive.properties:/etc/trino/catalog/hive.properties" \
9091
trinodb/trino:468
9192
9293
echo "Starting Grafana..."
@@ -97,6 +98,20 @@ jobs:
9798
--volume "$(pwd):/var/lib/grafana/plugins/trino" \
9899
--env "GF_PLUGINS_ALLOW_LOADING_UNSIGNED_PLUGINS=trino-datasource" \
99100
grafana/grafana:11.4.0
101+
102+
echo "Waiting for Trino to be ready..."
103+
while true; do
104+
if docker logs trino 2>&1 | grep -q '======== SERVER STARTED ========'; then
105+
echo "Trino is ready!"
106+
break
107+
fi
108+
echo "Waiting for Trino..."
109+
sleep 5
110+
done
111+
112+
echo "Preconfiguring trino..."
113+
docker exec trino trino --user admin --execute "GRANT admin TO USER grafana IN hive;"
114+
echo "Done."
100115
101116
- name: End to end test
102117
run: |

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ docker run -d -p 3000:3000 \
2626
* OAuth
2727
* Raw SQL editor only, no query builder yet
2828
* Macros
29+
* Client tags support, used to identify resource groups.
2930

3031
## Macros support
3132

‎go.mod‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,12 @@
11
module github.com/trinodb/grafana-trino
22

3-
go 1.23.5
4-
5-
toolchain go1.24.2
3+
go 1.24.7
64

75
require (
86
github.com/grafana/grafana-plugin-sdk-go v0.274.0
97
github.com/grafana/sqlds/v2 v2.7.2
108
github.com/pkg/errors v0.9.1
11-
github.com/trinodb/trino-go-client v0.323.0
9+
github.com/trinodb/trino-go-client v0.333.0
1210
)
1311

1412
require (
@@ -52,7 +50,7 @@ require (
5250
github.com/josharian/intern v1.0.0 // indirect
5351
github.com/json-iterator/go v1.1.12 // indirect
5452
github.com/jszwedko/go-datemath v0.1.1-0.20230526204004-640a500621d6 // indirect
55-
github.com/klauspost/compress v1.18.0 // indirect
53+
github.com/klauspost/compress v1.18.1 // indirect
5654
github.com/klauspost/cpuid/v2 v2.2.10 // indirect
5755
github.com/magefile/mage v1.15.0 // indirect
5856
github.com/mailru/easyjson v0.9.0 // indirect
@@ -75,6 +73,7 @@ require (
7573
github.com/oklog/run v1.1.0 // indirect
7674
github.com/olekukonko/tablewriter v0.0.5 // indirect
7775
github.com/perimeterx/marshmallow v1.1.5 // indirect
76+
github.com/pierrec/lz4 v2.6.1+incompatible // indirect
7877
github.com/pierrec/lz4/v4 v4.1.22 // indirect
7978
github.com/prometheus/client_golang v1.22.0 // indirect
8079
github.com/prometheus/client_model v0.6.1 // indirect
@@ -99,15 +98,16 @@ require (
9998
go.opentelemetry.io/otel/sdk v1.35.0 // indirect
10099
go.opentelemetry.io/otel/trace v1.35.0 // indirect
101100
go.opentelemetry.io/proto/otlp v1.5.0 // indirect
102-
golang.org/x/crypto v0.37.0 // indirect
101+
golang.org/x/crypto v0.45.0 // indirect
103102
golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 // indirect
104-
golang.org/x/mod v0.24.0 // indirect
105-
golang.org/x/net v0.39.0 // indirect
106-
golang.org/x/sync v0.13.0 // indirect
107-
golang.org/x/sys v0.32.0 // indirect
108-
golang.org/x/term v0.31.0 // indirect
109-
golang.org/x/text v0.24.0 // indirect
110-
golang.org/x/tools v0.32.0 // indirect
103+
golang.org/x/mod v0.29.0 // indirect
104+
golang.org/x/net v0.47.0 // indirect
105+
golang.org/x/sync v0.18.0 // indirect
106+
golang.org/x/sys v0.38.0 // indirect
107+
golang.org/x/telemetry v0.0.0-20251008203120-078029d740a8 // indirect
108+
golang.org/x/term v0.37.0 // indirect
109+
golang.org/x/text v0.31.0 // indirect
110+
golang.org/x/tools v0.38.0 // indirect
111111
golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
112112
google.golang.org/genproto/googleapis/api v0.0.0-20250407143221-ac9807e6c755 // indirect
113113
google.golang.org/genproto/googleapis/rpc v0.0.0-20250407143221-ac9807e6c755 // indirect

‎go.sum‎

Lines changed: 92 additions & 47 deletions
Large diffs are not rendered by default.

‎pkg/trino/datasource-context.go‎

Lines changed: 61 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,22 @@ package trino
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"strings"
78

89
"github.com/grafana/grafana-plugin-sdk-go/backend"
910
"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
11+
"github.com/grafana/grafana-plugin-sdk-go/backend/log"
1012
"github.com/grafana/sqlds/v2"
1113
"github.com/trinodb/grafana-trino/pkg/trino/models"
1214
)
1315

1416
const (
15-
accessTokenKey = "accessToken"
16-
trinoUserHeader = "X-Trino-User"
17-
bearerPrefix = "Bearer "
17+
accessTokenKey = "accessToken"
18+
trinoUserHeader = "X-Trino-User"
19+
trinoClientTagsKey = "X-Trino-Client-Tags"
20+
bearerPrefix = "Bearer "
1821
)
1922

2023
type SQLDatasourceWithTrinoUserContext struct {
@@ -36,13 +39,67 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req
3639
if user == nil {
3740
return nil, fmt.Errorf("user can't be nil if impersonation is enabled")
3841
}
39-
4042
ctx = context.WithValue(ctx, trinoUserHeader, user)
4143
}
4244

45+
ctx = injectClientTags(ctx, req, settings)
46+
4347
return ds.SQLDatasource.QueryData(ctx, req)
4448
}
4549

50+
func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) context.Context {
51+
header := req.GetHTTPHeader(backend.OAuthIdentityTokenHeaderName)
52+
53+
if strings.HasPrefix(header, bearerPrefix) {
54+
token := strings.TrimPrefix(header, bearerPrefix)
55+
return context.WithValue(ctx, accessTokenKey, token)
56+
}
57+
58+
return ctx
59+
}
60+
61+
func injectClientTags(ctx context.Context, req *backend.QueryDataRequest, settings models.TrinoDatasourceSettings) context.Context {
62+
type queryClientTag struct {
63+
ClientTags string `json:"clientTags"`
64+
}
65+
66+
tagSet := make(map[string]struct{})
67+
68+
for i := range req.Queries {
69+
var queryTags queryClientTag
70+
if err := json.Unmarshal(req.Queries[i].JSON, &queryTags); err != nil {
71+
log.DefaultLogger.Warn(
72+
"failed to unmarshal query client tags",
73+
"queryIndex", i,
74+
"error", err,
75+
)
76+
continue
77+
}
78+
79+
for _, tag := range strings.Split(queryTags.ClientTags, ",") {
80+
tag = strings.TrimSpace(tag)
81+
if tag != "" {
82+
tagSet[tag] = struct{}{}
83+
}
84+
}
85+
}
86+
87+
if len(tagSet) > 0 {
88+
tags := make([]string, 0, len(tagSet))
89+
for tag := range tagSet {
90+
tags = append(tags, tag)
91+
}
92+
93+
return context.WithValue(ctx, trinoClientTagsKey, strings.Join(tags, ","))
94+
}
95+
96+
if settings.ClientTags != "" {
97+
return context.WithValue(ctx, trinoClientTagsKey, settings.ClientTags)
98+
}
99+
100+
return ctx
101+
}
102+
46103
func (ds *SQLDatasourceWithTrinoUserContext) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
47104
_, err := ds.SQLDatasource.NewDatasource(settings)
48105
if err != nil {
@@ -55,14 +112,3 @@ func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext {
55112
base := sqlds.NewDatasource(c)
56113
return &SQLDatasourceWithTrinoUserContext{*base}
57114
}
58-
59-
func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) context.Context {
60-
header := req.GetHTTPHeader(backend.OAuthIdentityTokenHeaderName)
61-
62-
if strings.HasPrefix(header, bearerPrefix) {
63-
token := strings.TrimPrefix(header, bearerPrefix)
64-
return context.WithValue(ctx, accessTokenKey, token)
65-
}
66-
67-
return ctx
68-
}

‎pkg/trino/datasource.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ func (s *TrinoDatasource) SetQueryArgs(ctx context.Context, headers http.Header)
8383

8484
user := ctx.Value(trinoUserHeader)
8585
accessToken := ctx.Value(accessTokenKey)
86+
clientTags := ctx.Value(trinoClientTagsKey)
8687

8788
if user != nil {
8889
args = append(args, sql.Named(trinoUserHeader, string(user.(*backend.User).Login)))
@@ -92,6 +93,10 @@ func (s *TrinoDatasource) SetQueryArgs(ctx context.Context, headers http.Header)
9293
args = append(args, sql.Named(accessTokenKey, accessToken.(string)))
9394
}
9495

96+
if clientTags != nil {
97+
args = append(args, sql.Named(trinoClientTagsKey, clientTags.(string)))
98+
}
99+
95100
return args
96101
}
97102

‎pkg/trino/driver/driver.go‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@ import (
66
"database/sql"
77
"errors"
88
"fmt"
9-
trinoClient "github.com/trinodb/grafana-trino/pkg/trino/client"
109
"net/http"
1110
"strings"
1211

12+
trinoClient "github.com/trinodb/grafana-trino/pkg/trino/client"
13+
1314
"github.com/trinodb/grafana-trino/pkg/trino/models"
1415
"github.com/trinodb/trino-go-client/trino"
1516
_ "github.com/trinodb/trino-go-client/trino"
@@ -94,12 +95,19 @@ func Open(settings models.TrinoDatasourceSettings) (*sql.DB, error) {
9495
if err != nil {
9596
return nil, err
9697
}
98+
99+
roles, err := parseRoles(settings.Roles)
100+
if err != nil {
101+
return nil, err
102+
}
103+
97104
config := trino.Config{
98105
ServerURI: settings.URL.String(),
99106
Source: "grafana",
100107
CustomClientName: "grafana",
101108
ForwardAuthorizationHeader: true,
102109
AccessToken: settings.AccessToken,
110+
Roles: roles,
103111
}
104112

105113
dsn, err := config.FormatDSN()
@@ -108,3 +116,23 @@ func Open(settings models.TrinoDatasourceSettings) (*sql.DB, error) {
108116
}
109117
return sql.Open(DriverName, dsn)
110118
}
119+
120+
func parseRoles(roleStr string) (map[string]string, error) {
121+
roles := make(map[string]string)
122+
if strings.TrimSpace(roleStr) == "" {
123+
return roles, nil
124+
}
125+
pairs := strings.Split(roleStr, ";")
126+
for _, pair := range pairs {
127+
parts := strings.SplitN(pair, ":", 2)
128+
if len(parts) != 2 {
129+
return nil, fmt.Errorf("Invalid role format. expected catalog:role, got '%s'", pair)
130+
}
131+
catalog := strings.TrimSpace(parts[0])
132+
role := strings.TrimSpace(parts[1])
133+
if catalog != "" && role != "" {
134+
roles[catalog] = role
135+
}
136+
}
137+
return roles, nil
138+
}

‎pkg/trino/models/settings.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ type TrinoDatasourceSettings struct {
2020
ClientId string `json:"clientId"`
2121
ClientSecret string `json:"clientSecret"`
2222
ImpersonationUser string `json:"impersonationUser"`
23+
Roles string `json:"roles"`
24+
ClientTags string `json:"clientTags"`
2325
}
2426

2527
func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings) error {

‎src/ConfigEditor.tsx‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,12 @@ export class ConfigEditor extends PureComponent<Props, State> {
3434
const onImpersonationUserChange = (event: ChangeEvent<HTMLInputElement>) => {
3535
onOptionsChange({...options, jsonData: {...options.jsonData, impersonationUser: event.target.value}})
3636
};
37+
const onRolesChange = (event: ChangeEvent<HTMLInputElement>) => {
38+
onOptionsChange({...options, jsonData: {...options.jsonData, roles: event.target.value}})
39+
};
40+
const onClientTagsChange = (event: ChangeEvent<HTMLInputElement>) => {
41+
onOptionsChange({...options, jsonData: {...options.jsonData, clientTags: event.target.value}})
42+
};
3743
return (
3844
<div className="gf-form-group">
3945
<DataSourceHttpSettings
@@ -72,6 +78,33 @@ export class ConfigEditor extends PureComponent<Props, State> {
7278
/>
7379
</InlineField>
7480
</div>
81+
<div className="gf-form-inline">
82+
<InlineField
83+
label="Roles"
84+
tooltip="Authorization roles to use for catalogs, specified as a list of key-value pairs for the catalog and role. For example, system:roleS;catalog1:roleA;catalog2:roleB"
85+
labelWidth={26}
86+
>
87+
<Input
88+
value={options.jsonData?.roles ?? ''}
89+
onChange={onRolesChange}
90+
width={40}
91+
/>
92+
</InlineField>
93+
</div>
94+
<div className="gf-form-inline">
95+
<InlineField
96+
label="Client Tags"
97+
tooltip="A comma-separated list of strings, used to identify Trino resource groups."
98+
labelWidth={26}
99+
>
100+
<Input
101+
value={options.jsonData?.clientTags ?? ''}
102+
onChange={onClientTagsChange}
103+
width={60}
104+
placeholder="tag1,tag2,tag3"
105+
/>
106+
</InlineField>
107+
</div>
75108
</div>
76109

77110
<h3 className="page-heading">OAuth Trino Authentication</h3>

0 commit comments

Comments
 (0)