Support X-Trino-Role header

Author: ssheikin

.github/workflows/ci.yml

@@ -87,6 +87,7 @@ jobs:
             --name trino \
             --net trino \
             --volume "$(pwd)/test-data/trino/test-trino-config.properties:/etc/trino/config.properties" \
+            --volume "$(pwd)/test-data/trino/catalog/hive.properties:/etc/trino/catalog/hive.properties" \
             trinodb/trino:468
 
           echo "Starting Grafana..."
@@ -97,6 +98,20 @@ jobs:
             --volume "$(pwd):/var/lib/grafana/plugins/trino" \
             --env "GF_PLUGINS_ALLOW_LOADING_UNSIGNED_PLUGINS=trino-datasource" \
             grafana/grafana:11.4.0
+          
+          echo "Waiting for Trino to be ready..."
+          while true; do
+            if docker logs trino 2>&1 | grep -q '======== SERVER STARTED ========'; then
+              echo "Trino is ready!"
+              break
+            fi
+            echo "Waiting for Trino..."
+            sleep 5
+          done
+
+          echo "Preconfiguring trino..."
+          docker exec trino trino --user admin --execute "GRANT admin TO USER grafana IN hive;"
+          echo "Done."
 
       - name: End to end test
         run: |

pkg/trino/datasource-context.go

@@ -14,6 +14,7 @@ import (
 const (
 	accessTokenKey     = "accessToken"
 	trinoUserHeader    = "X-Trino-User"
+	trinoRoleHeader    = "X-Trino-Role"
 	trinoClientTagsKey = "X-Trino-Client-Tags"
 	bearerPrefix       = "Bearer "
 )
@@ -41,6 +42,10 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req
 		ctx = context.WithValue(ctx, trinoUserHeader, user)
 	}
 
+	if settings.Role != "" {
+		ctx = context.WithValue(ctx, trinoRoleHeader, settings.Role)
+	}
+
 	if settings.ClientTags != "" {
 		ctx = context.WithValue(ctx, trinoClientTagsKey, settings.ClientTags)
 	}

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)
+	role := ctx.Value(trinoRoleHeader)
 	clientTags := ctx.Value(trinoClientTagsKey)
 
 	if user != nil {
@@ -93,6 +94,10 @@ func (s *TrinoDatasource) SetQueryArgs(ctx context.Context, headers http.Header)
 		args = append(args, sql.Named(accessTokenKey, accessToken.(string)))
 	}
 
+	if role != nil {
+		args = append(args, sql.Named(trinoRoleHeader, role.(string)))
+	}
+
 	if clientTags != nil {
 		args = append(args, sql.Named(trinoClientTagsKey, clientTags.(string)))
 	}

pkg/trino/models/settings.go

@@ -20,6 +20,7 @@ type TrinoDatasourceSettings struct {
 	ClientId            string             `json:"clientId"`
 	ClientSecret        string             `json:"clientSecret"`
 	ImpersonationUser   string             `json:"impersonationUser"`
+	Role                string             `json:"role"`
 	ClientTags          string             `json:"clientTags"`
 }
 

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 onRoleChange = (event: ChangeEvent<HTMLInputElement>) => {
+      onOptionsChange({...options, jsonData: {...options.jsonData, role: event.target.value}})
+    };
     const onClientTagsChange = (event: ChangeEvent<HTMLInputElement>) => {
       onOptionsChange({...options, jsonData: {...options.jsonData, clientTags: event.target.value}})
     };
@@ -75,6 +78,19 @@ export class ConfigEditor extends PureComponent<Props, State> {
                 />
             </InlineField>
           </div>
+          <div className="gf-form-inline">
+            <InlineField
+              label="Role"
+              tooltip="For X-Trino-Role"
+              labelWidth={26}
+            >
+              <Input
+                value={options.jsonData?.role ?? ''}
+                onChange={onRoleChange}
+                width={40}
+              />
+            </InlineField>
+          </div>
           <div className="gf-form-inline">
             <InlineField
               label="Client Tags"

src/e2e.test.ts

@@ -43,6 +43,12 @@ async function setupDataSourceWithClientTags(page: Page, clientTags: string) {
     await page.getByTestId('data-testid Data source settings page Save and Test button').click();
 }
 
+async function setupDataSourceWithRole(page: Page, role: string) {
+    await page.getByTestId('data-testid Datasource HTTP settings url').fill('http://trino:8080');
+    await page.locator('div').filter({hasText: /^Role$/}).locator('input').fill(role);
+    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();
@@ -91,3 +97,34 @@ test('test with client tags', async ({ page }) => {
     await setupDataSourceWithClientTags(page, 'tag1,tag2,tag3');
     await runQueryAndCheckResults(page);
 });
+
+test('test with role', async ({ page }) => {
+    await login(page);
+    await goToTrinoSettings(page);
+    await setupDataSourceWithRole(page, 'hive=ROLE{admin}');
+    await runRoleQuery(page);
+    await expect(page.getByTestId('data-testid table body')).toContainText(/.*admin.*/);
+
+});
+
+test('test without role', async ({ page }) => {
+    await login(page);
+    await goToTrinoSettings(page);
+    await setupDataSourceWithRole(page, '');
+    await runRoleQuery(page);
+    await expect(page.getByText(/Access Denied: Cannot show roles/)).toBeVisible();
+});
+
+async function runRoleQuery(page: Page) {
+    await page.getByLabel(EXPORT_DATA).click();
+    await page.locator('div').filter({hasText: /^Format asChoose$/}).locator('svg').click();
+    await page.getByRole('option', {name: 'Table'}).click();
+    await setQuery(page, 'SHOW ROLES FROM hive')
+    await page.getByTestId('data-testid Code editor container').click();
+    await page.getByTestId('data-testid RefreshPicker run button').click();
+}
+
+async function setQuery(page: Page, query: string) {
+    await page.getByTestId('data-testid Code editor container').click({ clickCount: 4 });
+    await page.keyboard.type(query);
+}

src/types.ts

@@ -56,6 +56,7 @@ export interface TrinoDataSourceOptions extends DataSourceJsonData {
   tokenUrl?: string;
   clientId?: string;
   impersonationUser?: string;
+  role?: string;
   clientTags?: string;
 }
 /**

test-data/trino/catalog/hive.properties

@@ -0,0 +1,5 @@
+connector.name=hive
+hive.metastore=file
+hive.metastore.catalog.dir=/tmp/metastore
+hive.security=sql-standard
+fs.hadoop.enabled=true