Skip to content

Commit c699b9e

Browse files
RotRotAlRotRotAl
committed
feat(ticketless-client-tags): client tags in query editor (squashed)
1 parent 68edabe commit c699b9e

File tree

8 files changed

+162
-33
lines changed

8 files changed

+162
-33
lines changed

‎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

‎pkg/trino/datasource-context.go‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package trino
22

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

@@ -12,9 +13,10 @@ import (
1213
)
1314

1415
const (
15-
accessTokenKey = "accessToken"
16-
trinoUserHeader = "X-Trino-User"
17-
bearerPrefix = "Bearer "
16+
accessTokenKey = "accessToken"
17+
trinoUserHeader = "X-Trino-User"
18+
trinoClientTagsKey = "X-Trino-Client-Tags"
19+
bearerPrefix = "Bearer "
1820
)
1921

2022
type SQLDatasourceWithTrinoUserContext struct {
@@ -36,24 +38,35 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req
3638
if user == nil {
3739
return nil, fmt.Errorf("user can't be nil if impersonation is enabled")
3840
}
39-
4041
ctx = context.WithValue(ctx, trinoUserHeader, user)
4142
}
4243

44+
ctx = ds.injectClientTags(ctx, req, settings)
45+
4346
return ds.SQLDatasource.QueryData(ctx, req)
4447
}
4548

46-
func (ds *SQLDatasourceWithTrinoUserContext) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
47-
_, err := ds.SQLDatasource.NewDatasource(settings)
48-
if err != nil {
49-
return nil, err
49+
func (ds *SQLDatasourceWithTrinoUserContext) injectClientTags(contextWithTags context.Context, req *backend.QueryDataRequest, settings models.TrinoDatasourceSettings) context.Context {
50+
type queryClientTag struct {
51+
ClientTags string `json:"clientTags"`
5052
}
51-
return ds, nil
52-
}
5353

54-
func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext {
55-
base := sqlds.NewDatasource(c)
56-
return &SQLDatasourceWithTrinoUserContext{*base}
54+
for i := range req.Queries {
55+
var queryTags queryClientTag
56+
if err := json.Unmarshal(req.Queries[i].JSON, &queryTags); err != nil {
57+
continue
58+
}
59+
if queryTags.ClientTags != "" {
60+
contextWithTags = context.WithValue(contextWithTags, trinoClientTagsKey, queryTags.ClientTags)
61+
return contextWithTags
62+
}
63+
}
64+
65+
if contextWithTags.Value(trinoClientTagsKey) == nil && settings.ClientTags != "" {
66+
contextWithTags = context.WithValue(contextWithTags, trinoClientTagsKey, settings.ClientTags)
67+
}
68+
69+
return contextWithTags
5770
}
5871

5972
func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) context.Context {
@@ -66,3 +79,16 @@ func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) conte
6679

6780
return ctx
6881
}
82+
83+
func (ds *SQLDatasourceWithTrinoUserContext) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
84+
_, err := ds.SQLDatasource.NewDatasource(settings)
85+
if err != nil {
86+
return nil, err
87+
}
88+
return ds, nil
89+
}
90+
91+
func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext {
92+
base := sqlds.NewDatasource(c)
93+
return &SQLDatasourceWithTrinoUserContext{*base}
94+
}

‎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/models/settings.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ type TrinoDatasourceSettings struct {
2020
ClientId string `json:"clientId"`
2121
ClientSecret string `json:"clientSecret"`
2222
ImpersonationUser string `json:"impersonationUser"`
23+
ClientTags string `json:"clientTags"`
2324
}
2425

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

‎src/ConfigEditor.tsx‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ 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 onClientTagsChange = (event: ChangeEvent<HTMLInputElement>) => {
38+
onOptionsChange({...options, jsonData: {...options.jsonData, clientTags: event.target.value}})
39+
};
3740
return (
3841
<div className="gf-form-group">
3942
<DataSourceHttpSettings
@@ -72,6 +75,20 @@ export class ConfigEditor extends PureComponent<Props, State> {
7275
/>
7376
</InlineField>
7477
</div>
78+
<div className="gf-form-inline">
79+
<InlineField
80+
label="Client Tags"
81+
tooltip="A comma-separated list of strings, used to identify Trino resource groups."
82+
labelWidth={26}
83+
>
84+
<Input
85+
value={options.jsonData?.clientTags ?? ''}
86+
onChange={onClientTagsChange}
87+
width={60}
88+
placeholder="tag1,tag2,tag3"
89+
/>
90+
</InlineField>
91+
</div>
7592
</div>
7693

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

‎src/QueryEditor.tsx‎

Lines changed: 36 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { QueryEditorProps } from '@grafana/data';
33
import { DataSource } from './datasource';
44
import { TrinoDataSourceOptions, TrinoQuery, defaultQuery, SelectableFormatOptions } from './types';
55
import { FormatSelect, QueryCodeEditor } from '@grafana/aws-sdk';
6+
import { Input } from '@grafana/ui'; // <-- ADD THIS
67

78
type Props = QueryEditorProps<DataSource, TrinoQuery, TrinoDataSourceOptions>;
89

@@ -12,26 +13,43 @@ export function QueryEditor(props: Props) {
1213
...props.query,
1314
};
1415

16+
const onClientTagsChange = (event: React.ChangeEvent<HTMLInputElement>) => {
17+
props.onChange({
18+
...props.query,
19+
clientTags: event.target.value,
20+
});
21+
};
22+
1523
return (
1624
<>
17-
<div className="gf-form-group">
18-
<h6>Frames</h6>
19-
<FormatSelect
20-
query={props.query}
21-
options={SelectableFormatOptions}
22-
onChange={props.onChange}
23-
onRunQuery={props.onRunQuery}
24-
/>
25-
</div>
26-
<div style={{ minWidth: '400px', marginLeft: '10px', flex: 1 }}>
27-
<QueryCodeEditor
28-
language="sql"
29-
query={queryWithDefaults}
30-
onChange={props.onChange}
31-
onRunQuery={props.onRunQuery}
32-
getSuggestions={() => []}
33-
/>
34-
</div>
25+
<div className="gf-form-group">
26+
<h6>Frames</h6>
27+
<FormatSelect
28+
query={props.query}
29+
options={SelectableFormatOptions}
30+
onChange={props.onChange}
31+
onRunQuery={props.onRunQuery}
32+
/>
33+
</div>
34+
35+
<div className="gf-form-group">
36+
<h6>Client Tags</h6>
37+
<Input
38+
value={queryWithDefaults.clientTags || ''}
39+
placeholder="e.g. tag1,tag2,tag3"
40+
onChange={onClientTagsChange}
41+
/>
42+
</div>
43+
44+
<div style={{ minWidth: '400px', marginLeft: '10px', flex: 1 }}>
45+
<QueryCodeEditor
46+
language="sql"
47+
query={queryWithDefaults}
48+
onChange={props.onChange}
49+
onRunQuery={props.onRunQuery}
50+
getSuggestions={() => []}
51+
/>
52+
</div>
3553
</>
3654
);
3755
}

‎src/e2e.test.ts‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ async function goToTrinoSettings(page: Page) {
2121

2222
async function setupDataSourceWithAccessToken(page: Page) {
2323
await page.getByTestId('data-testid Datasource HTTP settings url').fill('http://trino:8080');
24-
await page.locator('div').filter({hasText: /^Impersonate logged in userAccess token$/}).getByLabel('Toggle switch').click();
24+
await page.locator('div').filter({hasText: /^Impersonate logged in user$/}).getByLabel('Toggle switch').click();
2525
await page.locator('div').filter({hasText: /^Access token$/}).locator('input[type="password"]').fill('aaa');
2626
await page.getByTestId('data-testid Data source settings page Save and Test button').click();
2727
}
@@ -35,6 +35,14 @@ async function setupDataSourceWithClientCredentials(page: Page, clientId: string
3535
await page.getByTestId('data-testid Data source settings page Save and Test button').click();
3636
}
3737

38+
async function setupDataSourceWithClientTags(page: Page, clientTags: string) {
39+
await page.getByTestId('data-testid Datasource HTTP settings url').fill('http://trino:8080');
40+
await page.locator('div').filter({hasText: /^Impersonate logged in user$/}).getByLabel('Toggle switch').click();
41+
await page.locator('div').filter({hasText: /^Access token$/}).locator('input[type="password"]').fill('aaa');
42+
await page.locator('div').filter({hasText: /^Client Tags$/}).locator('input').fill(clientTags);
43+
await page.getByTestId('data-testid Data source settings page Save and Test button').click();
44+
}
45+
3846
async function runQueryAndCheckResults(page: Page) {
3947
await page.getByLabel(EXPORT_DATA).click();
4048
await page.getByTestId('data-testid TimePicker Open Button').click();
@@ -48,6 +56,37 @@ async function runQueryAndCheckResults(page: Page) {
4856
await expect(page.getByTestId('data-testid table body')).toContainText(/.*1995-01-19 0.:00:005703857F.*/);
4957
}
5058

59+
async function runQueryAndRetrunRequset(
60+
page: Page,
61+
clientTag: string
62+
): Promise<import('@playwright/test').Request > {
63+
await page.getByLabel(EXPORT_DATA).click();
64+
await page.getByTestId('data-testid TimePicker Open Button').click();
65+
await page.getByTestId('data-testid Time Range from field').fill('1995-01-01');
66+
await page.getByTestId('data-testid Time Range to field').fill('1995-12-31');
67+
await page.getByTestId('data-testid TimePicker submit button').click();
68+
await page.locator('div').filter({ hasText: /^Format asChoose$/ }).locator('svg').click();
69+
await page.getByRole('option', { name: 'Table' }).click();
70+
await page.getByTestId('data-testid Code editor container').click();
71+
72+
await page.locator('div').filter({hasText: /^Client Tags$/}).locator('input').fill(clientTag);
73+
// Commit the input change
74+
await page.keyboard.press('Tab');
75+
76+
const [response] = await Promise.all([
77+
page.waitForResponse(
78+
res => res.url().includes('/api/ds/query') && res.request().method() === 'POST',
79+
{ timeout: 30000 }
80+
),
81+
page.getByTestId('data-testid RefreshPicker run button').click()
82+
]);
83+
84+
await expect(page.getByTestId('data-testid table body')).toContainText(/.*1995-01-19.*/);
85+
86+
return response.request();
87+
}
88+
89+
5190
test('test with access token', async ({ page }) => {
5291
await login(page);
5392
await goToTrinoSettings(page);
@@ -76,3 +115,23 @@ test('test client credentials flow with configured access token', async ({ page
76115
await setupDataSourceWithClientCredentials(page, GRAFANA_CLIENT);
77116
await expect(page.getByLabel(EXPORT_DATA)).toHaveCount(0);
78117
});
118+
119+
test('test with client tags', async ({ page }) => {
120+
await login(page);
121+
await goToTrinoSettings(page);
122+
await setupDataSourceWithClientTags(page, 'tag1,tag2,tag3');
123+
await runQueryAndCheckResults(page);
124+
});
125+
126+
test('query editor client tags override datasource-level tags', async ({ page }) => {
127+
await login(page);
128+
await goToTrinoSettings(page);
129+
await setupDataSourceWithClientTags(page, 'datasourceTag');
130+
131+
const request = await runQueryAndRetrunRequset(page, 'editorTag');
132+
133+
expect(request).toBeDefined();
134+
const body = JSON.parse(request.postData() || '{}');
135+
expect(body.queries?.[0]?.clientTags).toBe('editorTag');
136+
});
137+

‎src/types.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ export enum FormatOptions {
88
export interface TrinoQuery extends DataQuery {
99
rawSQL?: string;
1010
format?: FormatOptions;
11+
clientTags?: string;
1112
}
1213

1314
export const SelectableFormatOptions: Array<SelectableValue<FormatOptions>> = [
@@ -55,7 +56,8 @@ export interface TrinoDataSourceOptions extends DataSourceJsonData {
5556
enableImpersonation?: boolean;
5657
tokenUrl?: string;
5758
clientId?: string;
58-
impersonationUser?: string
59+
impersonationUser?: string;
60+
clientTags?: string;
5961
}
6062
/**
6163
* Value that is used in the backend, but never sent over HTTP to the frontend

0 commit comments

Comments
 (0)