diff --git a/README.md b/README.md index 52bcfcb..675c9e0 100644 --- a/README.md +++ b/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 for resource group identification ## Macros support diff --git a/pkg/trino/datasource-context.go b/pkg/trino/datasource-context.go index a44ea91..72c723d 100644 --- a/pkg/trino/datasource-context.go +++ b/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) } diff --git a/pkg/trino/datasource.go b/pkg/trino/datasource.go index ae64e2d..9c8d8f0 100644 --- a/pkg/trino/datasource.go +++ b/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 } diff --git a/pkg/trino/models/settings.go b/pkg/trino/models/settings.go index 1e89377..2867ab4 100644 --- a/pkg/trino/models/settings.go +++ b/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 { diff --git a/src/ConfigEditor.tsx b/src/ConfigEditor.tsx index 1d1583e..4d48b87 100644 --- a/src/ConfigEditor.tsx +++ b/src/ConfigEditor.tsx @@ -34,6 +34,9 @@ export class ConfigEditor extends PureComponent { const onImpersonationUserChange = (event: ChangeEvent) => { onOptionsChange({...options, jsonData: {...options.jsonData, impersonationUser: event.target.value}}) }; + const onClientTagsChange = (event: ChangeEvent) => { + onOptionsChange({...options, jsonData: {...options.jsonData, clientTags: event.target.value}}) + }; return (
{ />
+
+ + + +

OAuth Trino Authentication

diff --git a/src/e2e.test.ts b/src/e2e.test.ts index 14caeb6..254da68 100644 --- a/src/e2e.test.ts +++ b/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,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); +}); diff --git a/src/types.ts b/src/types.ts index 834cac9..528a5da 100644 --- a/src/types.ts +++ b/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