Feature/ticketless client tags in query editor

README.md

@@ -26,6 +26,8 @@ docker run -d -p 3000:3000 \
   * OAuth
 * Raw SQL editor only, no query builder yet
 * Macros
+* Client tags support, used to identify resource groups.
+* ```
 
 ## Macros support
 

pkg/trino/datasource-context.go

@@ -2,6 +2,7 @@ package trino
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -12,9 +13,10 @@ import (
 )
 
 const (
-	accessTokenKey  = "accessToken"
-	trinoUserHeader = "X-Trino-User"
-	bearerPrefix    = "Bearer "
+	accessTokenKey     = "accessToken"
+	trinoUserHeader    = "X-Trino-User"
+	trinoClientTagsKey = "X-Trino-Client-Tags"
+	bearerPrefix       = "Bearer "
 )
 
 type SQLDatasourceWithTrinoUserContext struct {
@@ -36,13 +38,42 @@ 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)
 	}
 
+	ctx = ds.injectClientTags(ctx, req, settings)
+
 	return ds.SQLDatasource.QueryData(ctx, req)
 }
 
+func (ds *SQLDatasourceWithTrinoUserContext) injectClientTags(contextWithTags context.Context, req *backend.QueryDataRequest, settings models.TrinoDatasourceSettings) context.Context {
+	type queryClientTag struct {
+		ClientTags string `json:"clientTags"`
+	}
+
+	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)
+			fmt.Printf("Using client tag from query editor: %s\n", queryTags.ClientTags)
+			return contextWithTags
+		}
+	}
+
+	if contextWithTags.Value(trinoClientTagsKey) == nil && settings.ClientTags != "" {
+		contextWithTags = context.WithValue(contextWithTags, trinoClientTagsKey, settings.ClientTags)
+		fmt.Printf("Using client tag from datasource settings: %s\n", settings.ClientTags)
+	}
+
+	return contextWithTags
+}
+
+
+
 func (ds *SQLDatasourceWithTrinoUserContext) NewDatasource(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
 	_, err := ds.SQLDatasource.NewDatasource(settings)
 	if err != nil {

pkg/trino/datasource.go

@@ -83,18 +83,27 @@ func (s *TrinoDatasource) SetQueryArgs(ctx context.Context, headers http.Header)
 
 	user := ctx.Value(trinoUserHeader)
 	accessToken := ctx.Value(accessTokenKey)
+	clientTags := ctx.Value(trinoClientTagsKey)
 
 	if user != nil {
+		headers.Set("X-Trino-User", string(user.(*backend.User).Login))
 		args = append(args, sql.Named(trinoUserHeader, string(user.(*backend.User).Login)))
 	}
 
 	if accessToken != nil {
+		headers.Set("X-Trino-Access-Token", accessToken.(string))
 		args = append(args, sql.Named(accessTokenKey, accessToken.(string)))
 	}
 
+	if clientTags != nil {
+		headers.Set("X-Trino-Client-Tags", clientTags.(string)) 
+		args = append(args, sql.Named(trinoClientTagsKey, clientTags.(string)))
+	}
+
 	return args
 }
 
+
 func (s *TrinoDatasource) Schemas(ctx context.Context, options sqlds.Options) ([]string, error) {
 	// TBD
 	return []string{}, nil

pkg/trino/models/settings.go

@@ -20,6 +20,7 @@ type TrinoDatasourceSettings struct {
 	ClientId            string             `json:"clientId"`
 	ClientSecret        string             `json:"clientSecret"`
 	ImpersonationUser   string             `json:"impersonationUser"`
+	ClientTags          string             `json:"clientTags"`
 }
 
 func (s *TrinoDatasourceSettings) Load(config backend.DataSourceInstanceSettings) error {

src/ConfigEditor.tsx

@@ -34,6 +34,9 @@ export class ConfigEditor extends PureComponent<Props, State> {
     const onImpersonationUserChange = (event: ChangeEvent<HTMLInputElement>) => {
       onOptionsChange({...options, jsonData: {...options.jsonData, impersonationUser: event.target.value}})
     };
+    const onClientTagsChange = (event: ChangeEvent<HTMLInputElement>) => {
+      onOptionsChange({...options, jsonData: {...options.jsonData, clientTags: event.target.value}})
+    };
     return (
       <div className="gf-form-group">
         <DataSourceHttpSettings
@@ -72,6 +75,20 @@ export class ConfigEditor extends PureComponent<Props, State> {
                 />
             </InlineField>
           </div>
+          <div className="gf-form-inline">
+            <InlineField
+              label="Client Tags"
+              tooltip="A comma-separated list of “tag” strings, used to identify Trino resource groups."
+              labelWidth={26}
+            >
+              <Input
+                value={options.jsonData?.clientTags ?? ''}
+                onChange={onClientTagsChange}
+                width={60}
+                placeholder="tag1,tag2,tag3"
+              />
+            </InlineField>
+          </div>
         </div>
 
         <h3 className="page-heading">OAuth Trino Authentication</h3>

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

@@ -21,7 +21,7 @@ async function goToTrinoSettings(page: Page) {
 
 async function setupDataSourceWithAccessToken(page: Page) {
     await page.getByTestId('data-testid Datasource HTTP settings url').fill('http://trino:8080');
-    await page.locator('div').filter({hasText: /^Impersonate logged in userAccess token$/}).getByLabel('Toggle switch').click();
+    await page.locator('#trino-settings-enable-impersonation').click({ force: true });
     await page.locator('div').filter({hasText: /^Access token$/}).locator('input[type="password"]').fill('aaa');
     await page.getByTestId('data-testid Data source settings page Save and Test button').click();
 }
@@ -35,6 +35,14 @@ async function setupDataSourceWithClientCredentials(page: Page, clientId: string
     await page.getByTestId('data-testid Data source settings page Save and Test button').click();
 }
 
+async function setupDataSourceWithClientTags(page: Page, clientTags: string) {
+    await page.getByTestId('data-testid Datasource HTTP settings url').fill('http://trino:8080');
+    await page.locator('#trino-settings-enable-impersonation').click({ force: true });
+    await page.locator('div').filter({hasText: /^Access token$/}).locator('input[type="password"]').fill('aaa');
+    await page.locator('div').filter({hasText: /^Client Tags$/}).locator('input').fill(clientTags);
+    await page.getByTestId('data-testid Data source settings page Save and Test button').click();
+}
+
 async function runQueryAndCheckResults(page: Page) {
     await page.getByLabel(EXPORT_DATA).click();
     await page.getByTestId('data-testid TimePicker Open Button').click();
@@ -76,3 +84,37 @@ test('test client credentials flow with configured access token', async ({ page
     await setupDataSourceWithClientCredentials(page, GRAFANA_CLIENT);
     await expect(page.getByLabel(EXPORT_DATA)).toHaveCount(0);
 });
+
+test('test with client tags', async ({ page }) => {
+    await login(page);
+    await goToTrinoSettings(page);
+    await setupDataSourceWithClientTags(page, 'tag1,tag2,tag3');
+    await runQueryAndCheckResults(page);
+});
+
+test('query editor client tags override datasource-level tags', async ({ page }) => {
+    await page.goto('http://localhost:3000/login');
+    await page.fill('input[name="user"]', 'admin');
+    await page.fill('input[name="password"]', 'admin');
+    await page.click('button[type="submit"]');
+
+    await page.click('[data-testid="toggle-menu"]');
+    await page.click('text=Connections');
+    await page.click('text=Trino');
+
+    await page.click('text=Add new data source');
+    await page.fill('input[placeholder="URL"]', 'http://trino:8080');
+    await page.fill('input[placeholder="Client Tags"]', 'datasourceTag');
+    await page.click('button:has-text("Save & Test")');
+
+    await page.click('text=Explore');
+
+    await page.fill('textarea[aria-label="Query editor"]', 'SHOW CATALOGS');
+    await page.fill('input[placeholder="Client Tags"]', 'editorTag');
+
+    await page.click('button[aria-label="Run query"]');
+
+    const request = await page.waitForRequest(req => req.url().includes('/api/ds/query') && req.method() === 'POST');
+    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>> = [
@@ -55,7 +56,8 @@ export interface TrinoDataSourceOptions extends DataSourceJsonData {
   enableImpersonation?: boolean;
   tokenUrl?: string;
   clientId?: string;
-  impersonationUser?: string
+  impersonationUser?: string;
+  clientTags?: string;
 }
 /**
  * Value that is used in the backend, but never sent over HTTP to the frontend