Add Client tags to the Trino Grafana plugin

Author: laserninja

README.md

@@ -26,6 +26,7 @@ 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

@@ -12,9 +12,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 {
@@ -40,6 +41,10 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req
 		ctx = context.WithValue(ctx, trinoUserHeader, user)
 	}
 
+	if settings.ClientTags != "" {
+		ctx = context.WithValue(ctx, trinoClientTagsKey, settings.ClientTags)
+	}
+
 	return ds.SQLDatasource.QueryData(ctx, req)
 }
 

pkg/trino/datasource.go

@@ -83,6 +83,7 @@ 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 {
 		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)
 		args = append(args, sql.Named(accessTokenKey, accessToken.(string)))
 	}
 
+	if clientTags != nil {
+		args = append(args, sql.Named(trinoClientTagsKey, clientTags.(string)))
+	}
+
 	return args
 }
 

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 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/e2e.test.ts

@@ -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('div').filter({hasText: /^Impersonate logged in user$/}).getByLabel('Toggle switch').click();
+    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,10 @@ 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);
+});

src/types.ts

@@ -55,7 +55,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