Refactor to use context argument

packages/catalog-client/src/CatalogClient.ts

@@ -24,13 +24,14 @@ import fetch from 'cross-fetch';
 import {
   AddLocationRequest,
   AddLocationResponse,
+  ApiContext,
+  CatalogApi,
   CatalogEntitiesRequest,
   CatalogListResponse,
-  CatalogRequestOptions,
   DiscoveryApi,
 } from './types';
 
-export class CatalogClient {
+export class CatalogClient implements CatalogApi {
   private readonly discoveryApi: DiscoveryApi;
 
   constructor(options: { discoveryApi: DiscoveryApi }) {
@@ -39,14 +40,14 @@ export class CatalogClient {
 
   async getLocationById(
     id: String,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<Location | undefined> {
-    return await this.getOptional(`/locations/${id}`, options);
+    return await this.getOptional(`/locations/${id}`, context);
   }
 
   async getEntities(
     request?: CatalogEntitiesRequest,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<CatalogListResponse<Entity>> {
     const { filter = {}, fields = [] } = request ?? {};
     const params: string[] = [];
@@ -68,31 +69,31 @@ export class CatalogClient {
     const query = params.length ? `?${params.join('&')}` : '';
     const entities: Entity[] = await this.getRequired(
       `/entities${query}`,
-      options,
+      context,
     );
     return { items: entities };
   }
 
   async getEntityByName(
     compoundName: EntityName,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<Entity | undefined> {
     const { kind, namespace = 'default', name } = compoundName;
     return this.getOptional(
       `/entities/by-name/${kind}/${namespace}/${name}`,
-      options,
+      context,
     );
   }
 
   async addLocation(
     { type = 'url', target, dryRun }: AddLocationRequest,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<AddLocationResponse> {
     const headers = {
       'Content-Type': 'application/json',
     } as { [header: string]: string };
-    if (options?.token) {
-      headers.authorization = `Bearer ${options.token}`;
+    if (context?.token) {
+      headers.authorization = `Bearer ${context.token}`;
     }
     const response = await fetch(
       `${await this.discoveryApi.getBaseUrl('catalog')}/locations${
@@ -128,27 +129,24 @@ export class CatalogClient {
 
   async getLocationByEntity(
     entity: Entity,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<Location | undefined> {
     const locationCompound = entity.metadata.annotations?.[LOCATION_ANNOTATION];
     const all: { data: Location }[] = await this.getRequired(
       '/locations',
-      options,
+      context,
     );
     return all
       .map(r => r.data)
       .find(l => locationCompound === `${l.type}:${l.target}`);
   }
 
-  async removeEntityByUid(
-    uid: string,
-    options?: CatalogRequestOptions,
-  ): Promise<void> {
+  async removeEntityByUid(uid: string, context?: ApiContext): Promise<void> {
     const response = await fetch(
       `${await this.discoveryApi.getBaseUrl('catalog')}/entities/by-uid/${uid}`,
       {
-        headers: options?.token
-          ? { authorization: `Bearer ${options.token}` }
+        headers: context?.token
+          ? { authorization: `Bearer ${context.token}` }
           : {},
         method: 'DELETE',
       },
@@ -166,14 +164,11 @@ export class CatalogClient {
   // Private methods
   //
 
-  private async getRequired(
-    path: string,
-    options?: CatalogRequestOptions,
-  ): Promise<any> {
+  private async getRequired(path: string, context?: ApiContext): Promise<any> {
     const url = `${await this.discoveryApi.getBaseUrl('catalog')}${path}`;
     const response = await fetch(url, {
-      headers: options?.token
-        ? { authorization: `Bearer ${options.token}` }
+      headers: context?.token
+        ? { authorization: `Bearer ${context.token}` }
         : {},
     });
 
@@ -188,12 +183,12 @@ export class CatalogClient {
 
   private async getOptional(
     path: string,
-    options?: CatalogRequestOptions,
+    context?: ApiContext,
   ): Promise<any | undefined> {
     const url = `${await this.discoveryApi.getBaseUrl('catalog')}${path}`;
     const response = await fetch(url, {
-      headers: options?.token
-        ? { authorization: `Bearer ${options.token}` }
+      headers: context?.token
+        ? { authorization: `Bearer ${context.token}` }
         : {},
     });
 

packages/catalog-client/src/index.ts

@@ -18,6 +18,7 @@ export { CatalogClient } from './CatalogClient';
 export type {
   AddLocationRequest,
   AddLocationResponse,
+  ApiContext,
   CatalogApi,
   CatalogEntitiesRequest,
   CatalogListResponse,

packages/catalog-client/src/types.ts

@@ -25,19 +25,32 @@ export type CatalogListResponse<T> = {
   items: T[];
 };
 
-export type CatalogRequestOptions = {
-  token: string | undefined;
+export type ApiContext = {
+  token?: string;
 };
 
 export interface CatalogApi {
-  getLocationById(id: String): Promise<Location | undefined>;
-  getEntityByName(name: EntityName): Promise<Entity | undefined>;
+  getLocationById(
+    id: String,
+    context?: ApiContext,
+  ): Promise<Location | undefined>;
+  getEntityByName(
+    name: EntityName,
+    context?: ApiContext,
+  ): Promise<Entity | undefined>;
   getEntities(
     request?: CatalogEntitiesRequest,
+    context?: ApiContext,
   ): Promise<CatalogListResponse<Entity>>;
-  addLocation(location: AddLocationRequest): Promise<AddLocationResponse>;
-  getLocationByEntity(entity: Entity): Promise<Location | undefined>;
-  removeEntityByUid(uid: string): Promise<void>;
+  addLocation(
+    location: AddLocationRequest,
+    context?: ApiContext,
+  ): Promise<AddLocationResponse>;
+  getLocationByEntity(
+    entity: Entity,
+    context?: ApiContext,
+  ): Promise<Location | undefined>;
+  removeEntityByUid(uid: string, context?: ApiContext): Promise<void>;
 }
 
 export type AddLocationRequest = {

plugins/auth-backend/src/lib/catalog/CatalogIdentityClient.test.ts

@@ -14,34 +14,38 @@
  * limitations under the License.
  */
 
-import { CatalogClient } from '@backstage/catalog-client';
+import { CatalogApi } from '@backstage/catalog-client';
+import { UserEntity } from '@backstage/catalog-model';
 import { CatalogIdentityClient } from './CatalogIdentityClient';
 
-jest.mock('@backstage/catalog-client');
-const MockedCatalogClient = CatalogClient as jest.Mock<CatalogClient>;
-
 describe('CatalogIdentityClient', () => {
-  const token = 'fake-id-token';
+  const catalogApi: jest.Mocked<CatalogApi> = {
+    getLocationById: jest.fn(),
+    getEntityByName: jest.fn(),
+    getEntities: jest.fn(),
+    addLocation: jest.fn(),
+    getLocationByEntity: jest.fn(),
+    removeEntityByUid: jest.fn(),
+  };
 
   afterEach(() => jest.resetAllMocks());
 
   it('passes through the correct search params', async () => {
+    catalogApi.getEntities.mockResolvedValueOnce({ items: [{} as UserEntity] });
     const client = new CatalogIdentityClient({
-      catalogClient: new MockedCatalogClient(),
+      catalogApi: catalogApi as CatalogApi,
     });
 
-    client.findUser({ annotations: { key: 'value' } }, { token });
+    client.findUser({ annotations: { key: 'value' } });
 
-    const getEntities = MockedCatalogClient.mock.instances[0].getEntities;
-    expect(getEntities).toHaveBeenCalledWith(
+    expect(catalogApi.getEntities).toBeCalledWith(
       {
         filter: {
           kind: 'user',
           'metadata.annotations.key': 'value',
         },
       },
-      { token },
+      undefined,
     );
-    expect(getEntities).toHaveBeenCalledTimes(1);
   });
 });

plugins/auth-backend/src/lib/catalog/CatalogIdentityClient.ts

@@ -15,9 +15,8 @@
  */
 
 import { ConflictError, NotFoundError } from '@backstage/backend-common';
-import { CatalogClient } from '@backstage/catalog-client';
+import { ApiContext, CatalogApi } from '@backstage/catalog-client';
 import { UserEntity } from '@backstage/catalog-model';
-import { CatalogIdentityRequestOptions } from './types';
 
 type UserQuery = {
   annotations: Record<string, string>;
@@ -27,10 +26,10 @@ type UserQuery = {
  * A catalog client tailored for reading out identity data from the catalog.
  */
 export class CatalogIdentityClient {
-  private readonly catalogClient: CatalogClient;
+  private readonly catalogApi: CatalogApi;
 
-  constructor(options: { catalogClient: CatalogClient }) {
-    this.catalogClient = options.catalogClient;
+  constructor(options: { catalogApi: CatalogApi }) {
+    this.catalogApi = options.catalogApi;
   }
 
   /**
@@ -38,11 +37,7 @@ export class CatalogIdentityClient {
    *
    * Throws a NotFoundError or ConflictError if 0 or multiple users are found.
    */
-  async findUser(
-    query: UserQuery,
-    options?: CatalogIdentityRequestOptions,
-  ): Promise<UserEntity> {
-    const token = options?.token;
+  async findUser(query: UserQuery, context?: ApiContext): Promise<UserEntity> {
     const filter: Record<string, string> = {
       kind: 'user',
     };
@@ -50,10 +45,7 @@ export class CatalogIdentityClient {
       filter[`metadata.annotations.${key}`] = value;
     }
 
-    const { items } = await this.catalogClient.getEntities(
-      { filter },
-      { token },
-    );
+    const { items } = await this.catalogApi.getEntities({ filter }, context);
 
     if (items.length !== 1) {
       if (items.length > 1) {

plugins/auth-backend/src/lib/catalog/types.ts

@@ -1,19 +0,0 @@
-/*
- * Copyright 2020 Spotify AB
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-export type CatalogIdentityRequestOptions = {
-  token: string | undefined;
-};

plugins/auth-backend/src/providers/aws-alb/provider.test.ts

@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 import { getVoidLogger } from '@backstage/backend-common';
-import { CatalogClient } from '@backstage/catalog-client';
 import express from 'express';
 import { JWT } from 'jose';
 
@@ -44,9 +43,6 @@ jest.mock('cross-fetch', () => ({
   },
 }));
 
-jest.mock('@backstage/catalog-client');
-const MockedCatalogClient = CatalogClient as jest.Mock<CatalogClient>;
-
 const identityResolutionCallbackMock = async (): Promise<AuthResponse<any>> => {
   return {
     backstageIdentity: {
@@ -71,7 +67,15 @@ beforeEach(() => {
 });
 
 describe('AwsALBAuthProvider', () => {
-  const catalogClient = new MockedCatalogClient();
+  const catalogApi = {
+    /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
+    addLocation: jest.fn(),
+    getEntities: jest.fn(),
+    getLocationByEntity: jest.fn(),
+    getLocationById: jest.fn(),
+    removeEntityByUid: jest.fn(),
+    getEntityByName: jest.fn(),
+  };
 
   const mockResponseSend = jest.fn();
   const mockRequest = ({
@@ -91,7 +95,7 @@ describe('AwsALBAuthProvider', () => {
 
   describe('should transform to type OAuthResponse', () => {
     it('when JWT is valid and identity is resolved successfully', async () => {
-      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogClient, {
+      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogApi, {
         region: 'us-west-2',
         identityResolutionCallback: identityResolutionCallbackMock,
         issuer: 'foo',
@@ -117,7 +121,7 @@ describe('AwsALBAuthProvider', () => {
   });
   describe('should fail when', () => {
     it('JWT is missing', async () => {
-      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogClient, {
+      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogApi, {
         region: 'us-west-2',
         identityResolutionCallback: identityResolutionCallbackMock,
         issuer: 'foo',
@@ -129,7 +133,7 @@ describe('AwsALBAuthProvider', () => {
     });
 
     it('JWT is invalid', async () => {
-      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogClient, {
+      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogApi, {
         region: 'us-west-2',
         identityResolutionCallback: identityResolutionCallbackMock,
         issuer: 'foo',
@@ -145,7 +149,7 @@ describe('AwsALBAuthProvider', () => {
     });
 
     it('issuer is invalid', async () => {
-      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogClient, {
+      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogApi, {
         region: 'us-west-2',
         identityResolutionCallback: identityResolutionCallbackMock,
         issuer: 'foobar',
@@ -159,7 +163,7 @@ describe('AwsALBAuthProvider', () => {
     });
 
     it('identity resolution callback rejects', async () => {
-      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogClient, {
+      const provider = new AwsAlbAuthProvider(getVoidLogger(), catalogApi, {
         region: 'us-west-2',
         identityResolutionCallback: identityResolutionCallbackRejectedMock,
         issuer: 'foo',

plugins/auth-backend/src/providers/aws-alb/provider.ts

@@ -25,7 +25,7 @@ import { KeyObject } from 'crypto';
 import { Logger } from 'winston';
 import NodeCache from 'node-cache';
 import { JWT } from 'jose';
-import { CatalogClient } from '@backstage/catalog-client';
+import { CatalogApi } from '@backstage/catalog-client';
 
 const ALB_JWT_HEADER = 'x-amzn-oidc-data';
 /**
@@ -44,13 +44,13 @@ export const getJWTHeaders = (input: string) => {
 
 export class AwsAlbAuthProvider implements AuthProviderRouteHandlers {
   private logger: Logger;
-  private readonly catalogClient: CatalogClient;
+  private readonly catalogClient: CatalogApi;
   private options: AwsAlbAuthProviderOptions;
   private readonly keyCache: NodeCache;
 
   constructor(
     logger: Logger,
-    catalogClient: CatalogClient,
+    catalogClient: CatalogApi,
     options: AwsAlbAuthProviderOptions,
   ) {
     this.logger = logger;
@@ -108,14 +108,14 @@ export class AwsAlbAuthProvider implements AuthProviderRouteHandlers {
 
 export const createAwsAlbProvider = ({
   logger,
-  catalogClient,
+  catalogApi,
   config,
   identityResolver,
 }: AuthProviderFactoryOptions) => {
   const region = config.getString('region');
   const issuer = config.getOptionalString('iss');
   if (identityResolver !== undefined) {
-    return new AwsAlbAuthProvider(logger, catalogClient, {
+    return new AwsAlbAuthProvider(logger, catalogApi, {
       region,
       issuer,
       identityResolutionCallback: identityResolver,

plugins/auth-backend/src/providers/google/provider.ts

@@ -190,7 +190,7 @@ export const createGoogleProvider: AuthProviderFactory = ({
   config,
   logger,
   tokenIssuer,
-  catalogClient,
+  catalogApi,
 }) =>
   OAuthEnvironmentHandler.mapConfig(config, envConfig => {
     const clientId = envConfig.getString('clientId');
@@ -203,7 +203,7 @@ export const createGoogleProvider: AuthProviderFactory = ({
       callbackUrl,
       logger,
       tokenIssuer,
-      identityClient: new CatalogIdentityClient({ catalogClient }),
+      identityClient: new CatalogIdentityClient({ catalogApi }),
     });
 
     return OAuthAdapter.fromConfig(globalConfig, provider, {

plugins/auth-backend/src/providers/types.ts

@@ -15,7 +15,7 @@
  */
 
 import { PluginEndpointDiscovery } from '@backstage/backend-common';
-import { CatalogClient } from '@backstage/catalog-client';
+import { CatalogApi } from '@backstage/catalog-client';
 import { Config } from '@backstage/config';
 import express from 'express';
 import { Logger } from 'winston';
@@ -122,7 +122,7 @@ export type ExperimentalIdentityResolver = (
    * An object containing information specific to the auth provider.
    */
   payload: object,
-  catalogClient: CatalogClient,
+  catalogApi: CatalogApi,
 ) => Promise<AuthResponse<any>>;
 
 export type AuthProviderFactoryOptions = {
@@ -132,8 +132,8 @@ export type AuthProviderFactoryOptions = {
   logger: Logger;
   tokenIssuer: TokenIssuer;
   discovery: PluginEndpointDiscovery;
+  catalogApi: CatalogApi;
   identityResolver?: ExperimentalIdentityResolver;
-  catalogClient: CatalogClient;
 };
 
 export type AuthProviderFactory = (

plugins/auth-backend/src/service/router.ts

@@ -66,7 +66,7 @@ export async function createRouter({
     keyDurationSeconds,
     logger: logger.child({ component: 'token-factory' }),
   });
-  const catalogClient = new CatalogClient({ discoveryApi: discovery });
+  const catalogApi = new CatalogClient({ discoveryApi: discovery });
 
   const secret = config.getOptionalString('auth.session.secret');
   if (secret) {
@@ -103,7 +103,7 @@ export async function createRouter({
         logger,
         tokenIssuer,
         discovery,
-        catalogClient,
+        catalogApi,
       });
 
       const r = Router();

plugins/catalog/src/CatalogClientWrapper.ts

@@ -18,6 +18,7 @@ import { Entity, EntityName, Location } from '@backstage/catalog-model';
 import {
   AddLocationRequest,
   AddLocationResponse,
+  ApiContext,
   CatalogApi,
   CatalogEntitiesRequest,
   CatalogListResponse,
@@ -25,6 +26,9 @@ import {
 } from '@backstage/catalog-client';
 import { IdentityApi } from '@backstage/core';
 
+/**
+ * CatalogClient wrapper that injects identity token for all requests
+ */
 export class CatalogClientWrapper implements CatalogApi {
   private readonly identityApi: IdentityApi;
   private readonly client: CatalogClient;
@@ -34,35 +38,69 @@ export class CatalogClientWrapper implements CatalogApi {
     this.identityApi = options.identityApi;
   }
 
-  async getLocationById(id: String): Promise<Location | undefined> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.getLocationById(id, { token });
+  private async injectToken(context?: ApiContext) {
+    const result: ApiContext = context ?? {};
+    // Inject identity token if not provided
+    if (!result.token) {
+      result.token = await this.identityApi.getIdToken();
+    }
+    return result;
+  }
+
+  async getLocationById(
+    id: String,
+    context?: ApiContext,
+  ): Promise<Location | undefined> {
+    return await this.client.getLocationById(
+      id,
+      await this.injectToken(context),
+    );
   }
 
   async getEntities(
     request?: CatalogEntitiesRequest,
+    context?: ApiContext,
   ): Promise<CatalogListResponse<Entity>> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.getEntities(request, { token });
+    return await this.client.getEntities(
+      request,
+      await this.injectToken(context),
+    );
   }
 
-  async getEntityByName(compoundName: EntityName): Promise<Entity | undefined> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.getEntityByName(compoundName, { token });
+  async getEntityByName(
+    compoundName: EntityName,
+    context?: ApiContext,
+  ): Promise<Entity | undefined> {
+    return await this.client.getEntityByName(
+      compoundName,
+      await this.injectToken(context),
+    );
   }
 
-  async addLocation(request: AddLocationRequest): Promise<AddLocationResponse> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.addLocation(request, { token });
+  async addLocation(
+    request: AddLocationRequest,
+    context?: ApiContext,
+  ): Promise<AddLocationResponse> {
+    return await this.client.addLocation(
+      request,
+      await this.injectToken(context),
+    );
   }
 
-  async getLocationByEntity(entity: Entity): Promise<Location | undefined> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.getLocationByEntity(entity, { token });
+  async getLocationByEntity(
+    entity: Entity,
+    context?: ApiContext,
+  ): Promise<Location | undefined> {
+    return await this.client.getLocationByEntity(
+      entity,
+      await this.injectToken(context),
+    );
   }
 
-  async removeEntityByUid(uid: string): Promise<void> {
-    const token = await this.identityApi.getIdToken();
-    return await this.client.removeEntityByUid(uid, { token });
+  async removeEntityByUid(uid: string, context?: ApiContext): Promise<void> {
+    return await this.client.removeEntityByUid(
+      uid,
+      await this.injectToken(context),
+    );
   }
 }

plugins/catalog/src/plugin.ts

@@ -27,7 +27,6 @@ import {
   catalogRouteRef,
   entityRouteRef,
 } from '@backstage/plugin-catalog-react';
-import { catalogRouteRef, entityRouteRef } from './routes';
 import { CatalogClientWrapper } from './CatalogClientWrapper';
 
 export const catalogPlugin = createPlugin({

plugins/scaffolder-backend/src/lib/catalog/CatalogEntityClient.ts

@@ -15,13 +15,12 @@
  */
 
 import { TemplateEntityV1alpha1 } from '@backstage/catalog-model';
-import { CatalogClient } from '@backstage/catalog-client';
+import { ApiContext, CatalogClient } from '@backstage/catalog-client';
 import {
   ConflictError,
   NotFoundError,
   PluginEndpointDiscovery,
 } from '@backstage/backend-common';
-import { CatalogEntityRequestOptions } from './types';
 
 /**
  * A catalog client tailored for reading out entity data from the catalog.
@@ -42,9 +41,8 @@ export class CatalogEntityClient {
    */
   async findTemplate(
     templateName: string,
-    options?: CatalogEntityRequestOptions,
+    context?: ApiContext,
   ): Promise<TemplateEntityV1alpha1> {
-    const token = options?.token;
     const { items: templates } = (await this.catalogClient.getEntities(
       {
         filter: {
@@ -52,7 +50,7 @@ export class CatalogEntityClient {
           'metadata.name': templateName,
         },
       },
-      { token },
+      context,
     )) as { items: TemplateEntityV1alpha1[] };
 
     if (templates.length !== 1) {

plugins/scaffolder-backend/src/lib/catalog/types.ts

@@ -1,32 +0,0 @@
-/*
- * Copyright 2021 Spotify AB
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-/*
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-export type CatalogEntityRequestOptions = {
-  token: string | undefined;
-};