Skip to content

Commit 2315ab4

Browse files
01epanineinchnick
01epa
authored and
nineinchnick
committed
Add support for oauth2 client-credentials flow to query from trino.
Token will be retrieved and cached until it is almost expired (1 minute before expiration). After that plugin will automatically retrieve token again.
1 parent 33f273b commit 2315ab4

File tree

10 files changed

+390
-37
lines changed

10 files changed

+390
-37
lines changed

‎.github/workflows/ci.yml‎

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ on:
77
pull_request:
88
branches:
99
- main
10+
1011
jobs:
1112
build:
1213
runs-on: ubuntu-latest
@@ -42,25 +43,63 @@ jobs:
4243
version: latest
4344
args: buildAll
4445

45-
- name: End to end test
46+
- name: Setup services (Trino, Grafana, PostgreSQL, Keycloak)
4647
run: |
4748
docker network create trino
4849
49-
docker run \
50-
--rm --detach \
50+
echo "Starting PostgreSQL..."
51+
docker run --rm --detach \
52+
--name postgres \
53+
--net trino \
54+
--env POSTGRES_USER=keycloak \
55+
--env POSTGRES_PASSWORD=keycloak \
56+
--env POSTGRES_DB=keycloak \
57+
postgres:17.4
58+
59+
echo "Starting Keycloak..."
60+
docker run --rm --detach \
61+
--name keycloak \
62+
--net trino \
63+
--publish 18080:8080 \
64+
--env KC_BOOTSTRAP_ADMIN_USERNAME=admin \
65+
--env KC_BOOTSTRAP_ADMIN_PASSWORD=admin \
66+
--env KC_DB=postgres \
67+
--env KC_DB_URL_HOST=postgres \
68+
--env KC_DB_URL_DATABASE=keycloak \
69+
--env KC_DB_USERNAME=keycloak \
70+
--env KC_DB_PASSWORD=keycloak \
71+
--volume "$(pwd)/test-data/test-keycloak-realm.json:/opt/keycloak/data/import/realm.json" \
72+
quay.io/keycloak/keycloak:26.1.4 \
73+
start-dev --import-realm
74+
75+
echo "Waiting for Keycloak to be ready..."
76+
while true; do
77+
if curl -s http://localhost:18080/realms/master | grep -q "realm"; then
78+
echo "Keycloak is ready!"
79+
break
80+
fi
81+
echo "Waiting for Keycloak..."
82+
sleep 5
83+
done
84+
85+
echo "Starting Trino..."
86+
docker run --rm --detach \
5187
--name trino \
5288
--net trino \
89+
--volume "$(pwd)/test-data/test-trino-config.properties:/etc/trino/config.properties" \
5390
trinodb/trino:468
5491
55-
docker run \
56-
--rm --detach \
92+
echo "Starting Grafana..."
93+
docker run --rm --detach \
5794
--name grafana \
5895
--net trino \
5996
--publish 3000:3000 \
6097
--volume "$(pwd):/var/lib/grafana/plugins/trino" \
6198
--env "GF_PLUGINS_ALLOW_LOADING_UNSIGNED_PLUGINS=trino-datasource" \
62-
grafana/grafana:11.4.0
99+
grafana/grafana:11.4.0
63100
64-
npx tsc -p tsconfig.json --noEmit
65-
npx playwright install
66-
npx playwright test
101+
- name: End to end test
102+
run: |
103+
npx tsc -p tsconfig.json --noEmit
104+
npx playwright install
105+
npx playwright test

‎pkg/trino/client/client.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package client
2+
3+
import (
4+
"encoding/json"
5+
"errors"
6+
"fmt"
7+
"github.com/grafana/grafana-plugin-sdk-go/backend/log"
8+
"io"
9+
"net/http"
10+
"net/url"
11+
"sync"
12+
"time"
13+
)
14+
15+
var (
16+
lock = &sync.Mutex{}
17+
token *Token
18+
)
19+
20+
type Client struct {
21+
*http.Client
22+
ClientId string
23+
ClientSecret string
24+
Url string
25+
ImpersonationUser string
26+
}
27+
28+
func (c *Client) Do(req *http.Request) (*http.Response, error) {
29+
token, err := c.getToken()
30+
if err != nil {
31+
return nil, err
32+
}
33+
req.Header.Set("Authorization", "Bearer "+token.AccessToken)
34+
if c.ImpersonationUser != "" {
35+
req.Header.Set("X-Trino-User", c.ImpersonationUser)
36+
}
37+
return c.Client.Do(req)
38+
}
39+
40+
func (c *Client) getToken() (*Token, error) {
41+
if token != nil && !token.isAlmostExpired() {
42+
return token, nil
43+
}
44+
lock.Lock()
45+
defer lock.Unlock()
46+
if token != nil && !token.isAlmostExpired() {
47+
return token, nil
48+
}
49+
newToken, err := c.retrieveToken()
50+
if err != nil {
51+
return nil, err
52+
}
53+
token = newToken
54+
return token, nil
55+
}
56+
57+
func (c *Client) retrieveToken() (*Token, error) {
58+
log.DefaultLogger.Debug("Try retrieve token")
59+
values := url.Values{
60+
"client_id": []string{c.ClientId},
61+
"client_secret": []string{c.ClientSecret},
62+
"grant_type": []string{"client_credentials"},
63+
}
64+
65+
token := &Token{}
66+
response, err := c.PostForm(c.Url, values)
67+
if err != nil {
68+
return nil, fmt.Errorf("failed to request the token response: %w", err)
69+
}
70+
defer response.Body.Close()
71+
if response.StatusCode != 200 {
72+
return nil, errors.New("Cannot obtain token from IDP. Status code=" + response.Status)
73+
}
74+
var jsonResponse []byte
75+
if jsonResponse, err = io.ReadAll(response.Body); err != nil {
76+
return nil, fmt.Errorf("failed to read the token response: %w", err)
77+
}
78+
err = json.Unmarshal(jsonResponse, token)
79+
if err != nil {
80+
return nil, fmt.Errorf("failed to decode the token response: %w", err)
81+
}
82+
token.ExpiresAt = time.Now().Add(time.Second * time.Duration(token.ExpiresIn))
83+
log.DefaultLogger.Debug("Token will expire at:", "date", token.ExpiresAt.Format(time.RFC1123Z))
84+
return token, nil
85+
}

‎pkg/trino/client/token.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package client
2+
3+
import (
4+
"time"
5+
)
6+
7+
// UntilExpirationInSeconds time has to be a greater than HTTP request timeout
8+
// In most HTTP clients HTTP request timeout is 30 seconds.
9+
const UntilExpirationInSeconds = 60
10+
11+
type Token struct {
12+
AccessToken string `json:"access_token"`
13+
ExpiresIn int `json:"expires_in"`
14+
ExpiresAt time.Time
15+
}
16+
17+
func (token *Token) isAlmostExpired() bool {
18+
if token.AccessToken == "" {
19+
return true
20+
} else {
21+
if time.Now().Add(time.Second * time.Duration(UntilExpirationInSeconds)).After(token.ExpiresAt) {
22+
return true
23+
}
24+
}
25+
return false
26+
}

‎pkg/trino/driver/driver.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@ import (
66
"database/sql"
77
"errors"
88
"fmt"
9+
trinoClient "github.com/trinodb/grafana-trino/pkg/trino/client"
910
"net/http"
11+
"strings"
1012

1113
"github.com/trinodb/grafana-trino/pkg/trino/models"
1214
"github.com/trinodb/trino-go-client/trino"
@@ -15,6 +17,17 @@ import (
1517

1618
const DriverName string = "trino"
1719

20+
// just compile time assertion
21+
var _ http.RoundTripper = &customTransport{}
22+
23+
type customTransport struct {
24+
client *trinoClient.Client
25+
}
26+
27+
func (t *customTransport) RoundTrip(req *http.Request) (*http.Response, error) {
28+
return t.client.Do(req)
29+
}
30+
1831
// Open registers a new driver with a unique name
1932
func Open(settings models.TrinoDatasourceSettings) (*sql.DB, error) {
2033
skipVerify := false
@@ -48,6 +61,35 @@ func Open(settings models.TrinoDatasourceSettings) (*sql.DB, error) {
4861
},
4962
},
5063
}
64+
if settings.TokenUrl != "" || settings.ClientId != "" || settings.ClientSecret != "" {
65+
if settings.AccessToken != "" {
66+
return nil, errors.New("access token must not be set within 'OAuth Trino Authentication' settings")
67+
}
68+
var missingParams []string
69+
if settings.TokenUrl == "" {
70+
missingParams = append(missingParams, "Token URL")
71+
}
72+
if settings.ClientId == "" {
73+
missingParams = append(missingParams, "Client id")
74+
}
75+
if settings.ClientSecret == "" {
76+
missingParams = append(missingParams, "Client secret")
77+
}
78+
if len(missingParams) > 0 {
79+
return nil, fmt.Errorf("missing parameters for 'OAuth Trino Authentication': %v", strings.Join(missingParams, ", "))
80+
}
81+
client = &http.Client{
82+
Transport: &customTransport{
83+
client: &trinoClient.Client{
84+
Client: client,
85+
ClientId: settings.ClientId,
86+
ClientSecret: settings.ClientSecret,
87+
Url: settings.TokenUrl,
88+
ImpersonationUser: settings.ImpersonationUser,
89+
},
90+
},
91+
}
92+
}
5193
err := trino.RegisterCustomClient("grafana", client)
5294
if err != nil {
5395
return nil, err

‎pkg/trino/models/settings.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ type TrinoDatasourceSettings struct {
1616
Opts httpclient.Options `json:"-"`
1717
EnableImpersonation bool `json:"enableImpersonation"`
1818
AccessToken string `json:"accessToken"`
19+
TokenUrl string `json:"tokenUrl"`
20+
ClientId string `json:"clientId"`
21+
ClientSecret string `json:"clientSecret"`
22+
ImpersonationUser string `json:"impersonationUser"`
1923
}
2024

2125
func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings) error {
@@ -48,5 +52,8 @@ func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings
4852
if token, ok := config.DecryptedSecureJSONData["accessToken"]; ok {
4953
s.AccessToken = token
5054
}
55+
if clientSecret, ok := config.DecryptedSecureJSONData["clientSecret"]; ok {
56+
s.ClientSecret = clientSecret
57+
}
5158
return nil
5259
}

‎src/ConfigEditor.tsx‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import React, { ChangeEvent, PureComponent } from 'react';
2-
import { DataSourceHttpSettings, InlineField, InlineSwitch, SecretInput } from '@grafana/ui';
2+
import { DataSourceHttpSettings, InlineField, InlineSwitch, SecretInput, Input } from '@grafana/ui';
33
import { DataSourcePluginOptionsEditorProps } from '@grafana/data';
44
import {TrinoDataSourceOptions, TrinoSecureJsonData} from './types';
55

@@ -19,6 +19,21 @@ export class ConfigEditor extends PureComponent<Props, State> {
1919
const onResetToken = () => {
2020
onOptionsChange({...options, secureJsonFields: {...options.secureJsonFields, accessToken: false }, secureJsonData: {...options.secureJsonData, accessToken: '' }});
2121
};
22+
const onTokenUrlChange = (event: ChangeEvent<HTMLInputElement>) => {
23+
onOptionsChange({...options, jsonData: {...options.jsonData, tokenUrl: event.target.value}})
24+
};
25+
const onClientIdChange = (event: ChangeEvent<HTMLInputElement>) => {
26+
onOptionsChange({...options, jsonData: {...options.jsonData, clientId: event.target.value}})
27+
};
28+
const onClientSecretChange = (event: ChangeEvent<HTMLInputElement>) => {
29+
onOptionsChange({...options, secureJsonData: {...options.secureJsonData, clientSecret: event.target.value}})
30+
};
31+
const onResetClientSecret = () => {
32+
onOptionsChange({...options, secureJsonFields: {...options.secureJsonFields, clientSecret: false}, secureJsonData: {...options.secureJsonData, clientSecret: ''}});
33+
};
34+
const onImpersonationUserChange = (event: ChangeEvent<HTMLInputElement>) => {
35+
onOptionsChange({...options, jsonData: {...options.jsonData, impersonationUser: event.target.value}})
36+
};
2237
return (
2338
<div className="gf-form-group">
2439
<DataSourceHttpSettings
@@ -58,6 +73,64 @@ export class ConfigEditor extends PureComponent<Props, State> {
5873
</InlineField>
5974
</div>
6075
</div>
76+
77+
<h3 className="page-heading">OAuth Trino Authentication</h3>
78+
<div className="gf-form-group">
79+
<div className="gf-form-inline">
80+
<InlineField
81+
label="Token URL"
82+
tooltip="If set, token is retrieved by client credentials flow before request to Trino is sent"
83+
labelWidth={26}
84+
>
85+
<Input
86+
value={options.jsonData?.tokenUrl ?? ''}
87+
onChange={onTokenUrlChange}
88+
width={60}
89+
/>
90+
</InlineField>
91+
</div>
92+
<div className="gf-form-inline">
93+
<InlineField
94+
label="Client id"
95+
tooltip="Required if Token URL is set"
96+
labelWidth={26}
97+
>
98+
<Input
99+
value={options.jsonData?.clientId ?? ''}
100+
onChange={onClientIdChange}
101+
width={60}
102+
/>
103+
</InlineField>
104+
</div>
105+
<div className="gf-form-inline">
106+
<InlineField
107+
label="Client secret"
108+
tooltip="Required if Token URL is set"
109+
labelWidth={26}
110+
>
111+
<SecretInput
112+
value={options.secureJsonData?.clientSecret ?? ''}
113+
isConfigured={options.secureJsonFields?.clientSecret}
114+
onChange={onClientSecretChange}
115+
width={60}
116+
onReset={onResetClientSecret}
117+
/>
118+
</InlineField>
119+
</div>
120+
<div className="gf-form-inline">
121+
<InlineField
122+
label="Impersonation user"
123+
tooltip="If set, this user will be used for impersonation in Trino"
124+
labelWidth={26}
125+
>
126+
<Input
127+
value={options.jsonData?.impersonationUser ?? ''}
128+
onChange={onImpersonationUserChange}
129+
width={60}
130+
/>
131+
</InlineField>
132+
</div>
133+
</div>
61134
</div>
62135
);
63136
}

0 commit comments

Comments
 (0)