diff --git a/pkg/trino/datasource-context.go b/pkg/trino/datasource-context.go index 72c723d..c179139 100644 --- a/pkg/trino/datasource-context.go +++ b/pkg/trino/datasource-context.go @@ -2,6 +2,7 @@ package trino import ( "context" + "encoding/json" "fmt" "strings" @@ -37,28 +38,35 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req if user == nil { return nil, fmt.Errorf("user can't be nil if impersonation is enabled") } - ctx = context.WithValue(ctx, trinoUserHeader, user) } - if settings.ClientTags != "" { - ctx = context.WithValue(ctx, trinoClientTagsKey, settings.ClientTags) - } + ctx = ds.injectClientTags(ctx, req, settings) 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 +func (ds *SQLDatasourceWithTrinoUserContext) injectClientTags(contextWithTags context.Context, req *backend.QueryDataRequest, settings models.TrinoDatasourceSettings) context.Context { + type queryClientTag struct { + ClientTags string `json:"clientTags"` } - return ds, nil -} -func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext { - base := sqlds.NewDatasource(c) - return &SQLDatasourceWithTrinoUserContext{*base} + for i := range req.Queries { + var queryTags queryClientTag + if err := json.Unmarshal(req.Queries[i].JSON, &queryTags); err != nil { + continue + } + if queryTags.ClientTags != "" { + contextWithTags = context.WithValue(contextWithTags, trinoClientTagsKey, queryTags.ClientTags) + return contextWithTags + } + } + + if contextWithTags.Value(trinoClientTagsKey) == nil && settings.ClientTags != "" { + contextWithTags = context.WithValue(contextWithTags, trinoClientTagsKey, settings.ClientTags) + } + + return contextWithTags } func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) context.Context { @@ -71,3 +79,16 @@ func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) conte return ctx } + +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} +} diff --git a/src/QueryEditor.tsx b/src/QueryEditor.tsx index b7bce81..c321cb5 100644 --- a/src/QueryEditor.tsx +++ b/src/QueryEditor.tsx @@ -3,6 +3,7 @@ import { QueryEditorProps } from '@grafana/data'; import { DataSource } from './datasource'; import { TrinoDataSourceOptions, TrinoQuery, defaultQuery, SelectableFormatOptions } from './types'; import { FormatSelect, QueryCodeEditor } from '@grafana/aws-sdk'; +import { Input } from '@grafana/ui'; // <-- ADD THIS type Props = QueryEditorProps; @@ -12,26 +13,43 @@ export function QueryEditor(props: Props) { ...props.query, }; + const onClientTagsChange = (event: React.ChangeEvent) => { + props.onChange({ + ...props.query, + clientTags: event.target.value, + }); + }; + return ( <> -
-
Frames
- -
-
- []} - /> -
+
+
Frames
+ +
+ +
+
Client Tags
+ +
+ +
+ []} + /> +
); } diff --git a/src/e2e.test.ts b/src/e2e.test.ts index 0058714..4ada275 100644 --- a/src/e2e.test.ts +++ b/src/e2e.test.ts @@ -56,6 +56,37 @@ async function runQueryAndCheckResults(page: Page) { await expect(page.getByTestId('data-testid table body')).toContainText(/.*1995-01-19 0.:00:005703857F.*/); } +async function runQueryAndRetrunRequset( + page: Page, + clientTag: string + ): Promise { + await page.getByLabel(EXPORT_DATA).click(); + await page.getByTestId('data-testid TimePicker Open Button').click(); + await page.getByTestId('data-testid Time Range from field').fill('1995-01-01'); + await page.getByTestId('data-testid Time Range to field').fill('1995-12-31'); + await page.getByTestId('data-testid TimePicker submit button').click(); + await page.locator('div').filter({ hasText: /^Format asChoose$/ }).locator('svg').click(); + await page.getByRole('option', { name: 'Table' }).click(); + await page.getByTestId('data-testid Code editor container').click(); + + await page.locator('div').filter({hasText: /^Client Tags$/}).locator('input').fill(clientTag); + // Commit the input change + await page.keyboard.press('Tab'); + + const [response] = await Promise.all([ + page.waitForResponse( + res => res.url().includes('/api/ds/query') && res.request().method() === 'POST', + { timeout: 30000 } + ), + page.getByTestId('data-testid RefreshPicker run button').click() + ]); + + await expect(page.getByTestId('data-testid table body')).toContainText(/.*1995-01-19.*/); + + return response.request(); + } + + test('test with access token', async ({ page }) => { await login(page); await goToTrinoSettings(page); @@ -91,3 +122,16 @@ test('test with client tags', async ({ page }) => { await setupDataSourceWithClientTags(page, 'tag1,tag2,tag3'); await runQueryAndCheckResults(page); }); + +test('query editor client tags override datasource-level tags', async ({ page }) => { + await login(page); + await goToTrinoSettings(page); + await setupDataSourceWithClientTags(page, 'datasourceTag'); + + const request = await runQueryAndRetrunRequset(page, 'editorTag'); + + expect(request).toBeDefined(); + const body = JSON.parse(request.postData() || '{}'); + expect(body.queries?.[0]?.clientTags).toBe('editorTag'); + }); + diff --git a/src/types.ts b/src/types.ts index 528a5da..dea20e3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -8,6 +8,7 @@ export enum FormatOptions { export interface TrinoQuery extends DataQuery { rawSQL?: string; format?: FormatOptions; + clientTags?: string; } export const SelectableFormatOptions: Array> = [