Feature/ticketless client tags in query editor

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}
+}

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<DataSource, TrinoQuery, TrinoDataSourceOptions>;
 
@@ -12,26 +13,43 @@ export function QueryEditor(props: Props) {
     ...props.query,
   };
 
+  const onClientTagsChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    props.onChange({
+      ...props.query,
+      clientTags: event.target.value,
+    });
+  };
+
   return (
     <>
-        <div className="gf-form-group">
-          <h6>Frames</h6>
-          <FormatSelect
-            query={props.query}
-            options={SelectableFormatOptions}
-            onChange={props.onChange}
-            onRunQuery={props.onRunQuery}
-          />
-        </div>
-        <div style={{ minWidth: '400px', marginLeft: '10px', flex: 1 }}>
-          <QueryCodeEditor
-            language="sql"
-            query={queryWithDefaults}
-            onChange={props.onChange}
-            onRunQuery={props.onRunQuery}
-            getSuggestions={() => []}
-          />
-        </div>
+      <div className="gf-form-group">
+        <h6>Frames</h6>
+        <FormatSelect
+          query={props.query}
+          options={SelectableFormatOptions}
+          onChange={props.onChange}
+          onRunQuery={props.onRunQuery}
+        />
+      </div>
+
+      <div className="gf-form-group">
+        <h6>Client Tags</h6>
+        <Input
+          value={queryWithDefaults.clientTags || ''}
+          placeholder="e.g. tag1,tag2,tag3"
+          onChange={onClientTagsChange}
+        />
+      </div>
+
+      <div style={{ minWidth: '400px', marginLeft: '10px', flex: 1 }}>
+        <QueryCodeEditor
+          language="sql"
+          query={queryWithDefaults}
+          onChange={props.onChange}
+          onRunQuery={props.onRunQuery}
+          getSuggestions={() => []}
+        />
+      </div>
     </>
   );
 }

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<import('@playwright/test').Request > {
+    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');
+  });
+

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<SelectableValue<FormatOptions>> = [