chore(plugins): Refactor the badge backend to store data in database instead of memory. Fix conversation as well

Signed-off-by: Rbillon59 <r.billon@celonis.com>
2023-04-07 16:06:09 +02:00
parent b34749a91e
commit 1aabf9e03f
17 changed files with 844 additions and 270 deletions
@@ -0,0 +1,182 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * 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.
+ */
+
+import {
+  PluginDatabaseManager,
+  resolvePackagePath,
+} from '@backstage/backend-common';
+import { Knex } from 'knex';
+import crypto from 'crypto';
+import { GetEntitiesResponse } from '@backstage/catalog-client';
+
+/**
+ * internal
+ * @public
+ */
+export interface BadgesStore {
+  createAllBadges(
+    entities: GetEntitiesResponse,
+    salt: string | undefined,
+  ): Promise<void>;
+
+  getBadgeFromHash(
+    hash: string,
+  ): Promise<{ name: string; namespace: string; kind: string } | undefined>;
+
+  getHashFromEntityMetadata(
+    name: string,
+    namespace: string,
+    kind: string,
+  ): Promise<{ hash: string } | undefined>;
+
+  deleteObsoleteHashes(
+    entities: GetEntitiesResponse,
+    salt: string | undefined,
+  ): Promise<void>;
+
+  countAllBadges(): Promise<number>;
+
+  getAllBadges(): Promise<
+    { name: string; namespace: string; kind: string; hash: string }[]
+  >;
+}
+
+const migrationsDir = resolvePackagePath(
+  '@backstage/plugin-badges-backend', // Package name
+  'migrations', // Migrations directory
+);
+
+/**
+ * DatabaseBadgesStore
+ * @internal
+ */
+export class DatabaseBadgesStore implements BadgesStore {
+  private constructor(private readonly db: Knex) {}
+
+  static async create({
+    database,
+    skipMigrations,
+  }: {
+    database: PluginDatabaseManager;
+    skipMigrations?: boolean;
+  }): Promise<DatabaseBadgesStore> {
+    const client = await database.getClient();
+
+    if (!database.migrations?.skip && !skipMigrations) {
+      await client.migrate.latest({
+        directory: migrationsDir,
+      });
+    }
+
+    return new DatabaseBadgesStore(client);
+  }
+
+  async createAllBadges(
+    entities: GetEntitiesResponse,
+    salt: string,
+  ): Promise<void> {
+    for (const entity of entities.items) {
+      const name = entity.metadata.name.toLocaleLowerCase();
+      const namespace =
+        entity.metadata.namespace?.toLocaleLowerCase() ?? 'default';
+      const kind = entity.kind.toLocaleLowerCase();
+      const entityHash = crypto
+        .createHash('sha256')
+        .update(`${kind}:${namespace}:${name}:${salt}`)
+        .digest('hex');
+
+      await this.db('badges')
+        .insert({
+          hash: entityHash,
+          namespace: namespace,
+          name: name,
+          kind: kind,
+        })
+        .onConflict()
+        .ignore();
+    }
+  }
+
+  async getAllBadges(): Promise<
+    { name: string; namespace: string; kind: string; hash: string }[]
+  > {
+    const result = await this.db('badges').select('*').orderBy('name', 'asc');
+    return result;
+  }
+
+  async getBadgeFromHash(
+    hash: string,
+  ): Promise<{ name: string; namespace: string; kind: string } | undefined> {
+    const result = await this.db('badges')
+      .select('namespace', 'name', 'kind')
+      .where({ hash: hash })
+      .first();
+
+    return result;
+  }
+
+  async getHashFromEntityMetadata(
+    name: string,
+    namespace: string,
+    kind: string,
+  ): Promise<{ hash: string } | undefined> {
+    const result = await this.db('badges')
+      .select('hash')
+      .where({ name: name, namespace: namespace, kind: kind })
+      .first();
+
+    return result;
+  }
+
+  async deleteObsoleteHashes(
+    entities: GetEntitiesResponse,
+    salt: string,
+  ): Promise<void> {
+    const entityInCatalog: string[] = [];
+    const entityInBadgeDatabase: string[] = [];
+    let entityToDelete: string[] = [];
+
+    for (const entity of entities.items) {
+      const name = entity.metadata.name.toLowerCase();
+      const namespace = entity.metadata.namespace?.toLowerCase() ?? 'default';
+      const kind = entity.kind.toLowerCase();
+      const entityHash = crypto
+        .createHash('sha256')
+        .update(`${kind}:${namespace}:${name}:${salt}`)
+        .digest('hex');
+
+      entityInCatalog.push(entityHash);
+    }
+
+    for (const entity of await this.getAllBadges()) {
+      entityInBadgeDatabase.push(entity.hash);
+    }
+
+    entityToDelete = entityInBadgeDatabase.filter(
+      entity => !entityInCatalog.includes(entity),
+    );
+
+    for (const entity of entityToDelete) {
+      await this.db('badges').where({ hash: entity }).del();
+    }
+  }
+
+  async countAllBadges(): Promise<number> {
+    const result = await this.db('badges').countDistinct('hash as count');
+    const count: number = +result[0].count;
+    return count;
+  }
+}
@@ -24,3 +24,4 @@ export * from './badges';
 export * from './lib';
 export * from './service/router';
 export * from './types';
+export * from './database/badgesStore';
@@ -26,11 +26,13 @@ import { Config } from '@backstage/config';
 import { NotFoundError } from '@backstage/errors';
 import { BadgeBuilder, DefaultBadgeBuilder } from '../lib/BadgeBuilder';
 import { BadgeContext, BadgeFactories } from '../types';
-import { isEmpty, isNil } from 'lodash';
+import { isNil } from 'lodash';
 import crypto from 'crypto';
 import { Logger } from 'winston';
 import { IdentityApi } from '@backstage/plugin-auth-node';
 import { getBearerTokenFromAuthorizationHeader } from '@backstage/plugin-auth-node';
+import type { BadgesStore } from '../database/badgesStore';
+
 /** @public */
 export interface RouterOptions {
  badgeBuilder?: BadgeBuilder;
@@ -41,6 +43,7 @@ export interface RouterOptions {
  tokenManager: TokenManager;
  logger: Logger;
  identity: IdentityApi;
+  db: BadgesStore;
 }

 /** @public */
@@ -54,38 +57,46 @@ export async function createRouter(
    new DefaultBadgeBuilder(options.badgeFactories || {});
  const router = Router();

-  const tokenManager = options.tokenManager;
-  let lookupTable: Map<
-    string,
-    {
-      name: string;
-      namespace: string;
-      kind: string;
-      creationDate: number;
-    }
-  > = new Map();
+  const { db, config, logger, tokenManager, discovery, identity } = options;
+  const baseUrl = await discovery.getExternalBaseUrl('badges');
+  const salt = config.getOptionalString('custom.badges-backend.salt');
+  const cacheTimeToLive =
+    config.getOptionalNumber('badgeDatabaseRefreshCacheTimeToLive') ?? 3600;
+
+  let lastDatabaseRefresh = 0;

  // Check if the users have enabled the obfuscation of the entity name
-  const obfuscate: boolean =
-    options.config.getOptionalBoolean('app.badges.obfuscate') ?? false;
+  if (config.getOptionalBoolean('app.badges.obfuscate')) {
+    logger.info('Badges obfuscation is enabled');

-  if (obfuscate) {
-    options.logger.info('Badges obfuscation is enabled');
+    if (isNil(salt)) {
+      throw new Error(
+        'Badges obfuscation is enabled but no salt has been provided',
+      );
+    }

    // Use the generated hash instead of the triplet namespace/kind/name
    router.get('/entity/:entityHash/badge-specs', async (req, res) => {
      const { entityHash } = req.params;

-      if (isLookupTableToRefresh(lookupTable)) {
-        lookupTable = await generateHashLookupTable();
+      // Chech if the database needs to be refreshed
+      if (await isBadgeDatabaseRefreshNeeded(lastDatabaseRefresh)) {
+        lastDatabaseRefresh = await refreshBadgeDatabase();
+      }
+
+      // Retrieve the badge info from the database
+      const badgeInfos = await db.getBadgeFromHash(entityHash);
+
+      if (isNil(badgeInfos)) {
+        throw new NotFoundError(
+          `No badge found for entity hash "${entityHash}"`,
+        );
      }
-      // Try to match the queried hash with a key in the lookup table
-      const entityInfo = await getEntityInfoFromLookupTable(entityHash);

      // If a mapping is found, map name, namespace and kind
-      const name = entityInfo.name;
-      const namespace = entityInfo.namespace;
-      const kind = entityInfo.kind;
+      const name = badgeInfos.name;
+      const namespace = badgeInfos.namespace;
+      const kind = badgeInfos.kind;
      const token = await tokenManager.getToken();

      // Query the catalog with the name, namespace, kind to get the entity informations
@@ -107,13 +118,8 @@ export async function createRouter(
      const specs = [];
      for (const badgeInfo of await badgeBuilder.getBadges()) {
        const context: BadgeContext = {
-          badgeUrl: await getBadgeObfuscatedUrl(
-            namespace,
-            kind,
-            name,
-            badgeInfo.id,
-          ),
-          config: options.config,
+          badgeUrl: await getBadgeObfuscatedUrl(entityHash, badgeInfo.id),
+          config: config,
          entity,
        };

@@ -129,19 +135,25 @@ export async function createRouter(

    // Use the generated hash instead of the triplet namespace/kind/name
    router.get('/entity/:entityHash/:badgeId', async (req, res) => {
-      const { entityHash, badgeId } = req.params;
-
-      if (isLookupTableToRefresh(lookupTable)) {
-        lookupTable = await generateHashLookupTable();
+      if (await isBadgeDatabaseRefreshNeeded(lastDatabaseRefresh)) {
+        lastDatabaseRefresh = await refreshBadgeDatabase();
      }

-      // Try to match the queried hash with a key in the lookup table
-      const entityInfo = await getEntityInfoFromLookupTable(entityHash);
+      const { entityHash, badgeId } = req.params;
+
+      // Retrieve the badge info from the database
+      const badgeInfo = await db.getBadgeFromHash(entityHash);
+
+      if (isNil(badgeInfo)) {
+        throw new NotFoundError(
+          `No badge found for entity hash "${entityHash}"`,
+        );
+      }

      // If a mapping is found, map name, namespace and kind
-      const name = entityInfo.name;
-      const namespace = entityInfo.namespace;
-      const kind = entityInfo.kind;
+      const name = badgeInfo.name;
+      const namespace = badgeInfo.namespace;
+      const kind = badgeInfo.kind;
      const token = await tokenManager.getToken();
      const entity = await catalog.getEntityByRef(
        {
@@ -168,8 +180,8 @@ export async function createRouter(
      const badgeOptions = {
        badgeInfo: { id: badgeId },
        context: {
-          badgeUrl: await getBadgeObfuscatedUrl(namespace, kind, name, badgeId),
-          config: options.config,
+          badgeUrl: await getBadgeObfuscatedUrl(entityHash, badgeId),
+          config: config,
          entity,
        },
      };
@@ -203,22 +215,42 @@ export async function createRouter(
        }

        try {
-          req.user = options.identity.getIdentity({ request: req });
+          req.user = identity.getIdentity({ request: req });
          next();
        } catch (error) {
-          options.tokenManager.authenticate(token.toString());
+          tokenManager.authenticate(token.toString());
          next(error);
        }
      },
      async (req, res) => {
+        if (await isBadgeDatabaseRefreshNeeded(lastDatabaseRefresh)) {
+          lastDatabaseRefresh = await refreshBadgeDatabase();
+        }
+
        const { namespace, kind, name } = req.params;
-        const salt = options.config.getString('custom.badges-backend.salt');
-        const hash = crypto
+        const storedEntityHash: { hash: string } | undefined =
+          await db.getHashFromEntityMetadata(name, namespace, kind);
+
+        if (isNil(storedEntityHash)) {
+          throw new NotFoundError(
+            `No hash found for entity "${namespace}/${kind}/${name}"`,
+          );
+        }
+
+        // Compare a live calculated hash with the stored one to avoid returning an obsolete entityHash (in case of a salt change)
+
+        const liveHash = crypto
          .createHash('sha256')
          .update(`${kind}:${namespace}:${name}:${salt}`)
          .digest('hex');

-        return res.status(200).json({ hash });
+        if (storedEntityHash.hash !== liveHash) {
+          throw new NotFoundError(
+            "Stored Hash and live Hash don't match, did you change the salt? Try to refresh the badge database table and try again",
+          );
+        }
+
+        return res.status(200).json(storedEntityHash);
      },
    );

@@ -248,7 +280,7 @@ export async function createRouter(
        for (const badgeInfo of await badgeBuilder.getBadges()) {
          const context: BadgeContext = {
            badgeUrl: await getBadgeUrl(namespace, kind, name, badgeInfo.id),
-            config: options.config,
+            config: config,
            entity,
          };

@@ -288,7 +320,7 @@ export async function createRouter(
          badgeInfo: { id: badgeId },
          context: {
            badgeUrl: await getBadgeUrl(namespace, kind, name, badgeId),
-            config: options.config,
+            config: config,
            entity,
          },
        };
@@ -314,117 +346,11 @@ export async function createRouter(
    return router;
  }

-  // This function generate a table that maps the hash to the namespace/kind/name triplet
-  async function generateHashLookupTable() {
-    const logger = options.logger.child({ service: 'badges-backend' });
-    const generationStartTimestamp = Date.now();
-    logger.info('Start Generating lookup table');
-    const token = await options.tokenManager.getToken();
-
-    // The salt is used to increase the hash entropy
-    const salt = options.config.getString('custom.badges-backend.salt');
-
-    // Get all the entities in the catalog of kind "Component"
-    const entitiesList = await catalog.getEntities(
-      {
-        filter: {
-          kind: 'Component',
-        },
-      },
-      token,
-    );
-    if (isEmpty(entitiesList)) {
-      throw new NotFoundError(`No entities found`);
-    }
-
-    // For each entity, generate the hash with the triplet kind:namespace:name then add the salt. Finally add it to the lookup table
-    entitiesList.items.map(async entity => {
-      const name = entity.metadata.name.toLocaleLowerCase();
-      const creationDate: number = Date.now();
-      const namespace =
-        entity.metadata.namespace?.toLocaleLowerCase() ?? 'default';
-      const kind = entity.kind.toLocaleLowerCase();
-      const hash = crypto
-        .createHash('sha256')
-        .update(`${kind}:${namespace}:${name}:${salt}`)
-        .digest('hex');
-
-      lookupTable.set(hash, { name, namespace, kind, creationDate });
-
-      return lookupTable;
-    });
-
-    // Monitor the lookup generation time and log it
-    logger.info(
-      `Finished generating lookup table in ${
-        Date.now() - generationStartTimestamp
-      } ms`,
-    );
-    // Monitor the lookup table size in byte
-    logger.info(`Lookup table size: ${lookupTable.size} entries`);
-
-    return lookupTable;
-  }
-
-  // This function check if the lookup table is empty or if it is expired (based on the cacheTimeToLive config) and need to be refreshed
-  function isLookupTableToRefresh(
-    table: Map<
-      string,
-      {
-        name: string;
-        namespace: string;
-        kind: string;
-        creationDate: number;
-      }
-    >,
-  ): boolean {
-    const lookupTableLength = table.size;
-
-    if (lookupTableLength === 0) {
-      return true;
-    }
-
-    const now = Date.now();
-    const cacheTimeToLive =
-      options.config.getOptionalNumber(
-        'custom.badges-backend.cacheTimeToLive',
-      ) ?? 10 * 60 * 1000;
-
-    const valuesArray = Array.from(table.values());
-    const lastEntry = valuesArray[valuesArray.length - 1];
-    const lastCreationDate = lastEntry.creationDate;
-
-    return lastCreationDate + cacheTimeToLive < now;
-  }
-
-  // This function return the namespace/kind/name triplet from the lookup table based on the hash
-  async function getEntityInfoFromLookupTable(
-    entityHash: string,
-  ): Promise<{ name: string; namespace: string; kind: string }> {
-    if (lookupTable.get(entityHash) === undefined) {
-      throw new NotFoundError(`No entity with hash "${entityHash}"`);
-    }
-
-    const name = lookupTable.get(entityHash)!.name;
-    const namespace = lookupTable.get(entityHash)!.namespace;
-    const kind = lookupTable.get(entityHash)!.kind;
-
-    return { name, namespace, kind };
-  }
-
  // This function return the obfuscated badge url based on the namespace/kind/name triplet
  async function getBadgeObfuscatedUrl(
-    namespace: string,
-    kind: string,
-    name: string,
+    hash: string,
    badgeId: string,
  ): Promise<string> {
-    const baseUrl = await options.discovery.getExternalBaseUrl('badges');
-    const salt = options.config.getString('custom.badges-backend.salt');
-    const hash = crypto
-      .createHash('sha256')
-      .update(`${kind}:${namespace}:${name}:${salt}`)
-      .digest('hex');
    return `${baseUrl}/entity/${hash}/${badgeId}`;
  }

@@ -435,7 +361,47 @@ export async function createRouter(
    name: string,
    badgeId: string,
  ): Promise<string> {
-    const baseUrl = await options.discovery.getExternalBaseUrl('badges');
    return `${baseUrl}/entity/${namespace}/${kind}/${name}/badge/${badgeId}`;
  }
+
+  async function isBadgeDatabaseRefreshNeeded(
+    lastDatabaseRefreshTimestamp: number,
+  ): Promise<boolean> {
+    if ((await db.countAllBadges()) === 0) {
+      logger.info('Badge database is empty, refreshing it');
+      return true;
+    }
+
+    if (
+      lastDatabaseRefreshTimestamp === 0 ||
+      Date.now() - lastDatabaseRefreshTimestamp > cacheTimeToLive * 1000
+    ) {
+      logger.info(
+        'Badge database refresh cache to live exceeded, refreshing it',
+      );
+      return true;
+    }
+
+    return false;
+  }
+
+  async function refreshBadgeDatabase(): Promise<number> {
+    const token = await tokenManager.getToken();
+    const entities = await catalog.getEntities(
+      {
+        filter: {
+          kind: ['Component'],
+        },
+      },
+      token,
+    );
+    logger.info(
+      `Refreshing badge database with ${entities.items.length} entities`,
+    );
+    await db.createAllBadges(entities, salt);
+    logger.info('Badge database refreshed, deleting obsolete hashes');
+    await db.deleteObsoleteHashes(entities, salt);
+
+    return Date.now();
+  }
 }
@@ -21,9 +21,12 @@ import {
  loadBackendConfig,
  ServerTokenManager,
  SingleHostDiscovery,
+  useHotMemoize,
 } from '@backstage/backend-common';
 import { createRouter } from './router';
 import { IdentityApi } from '@backstage/plugin-auth-node';
+import { DatabaseBadgesStore } from '../database/badgesStore';
+import Knex from 'knex';

 export interface ServerOptions {
  port: number;
@@ -37,6 +40,13 @@ export async function startStandaloneServer(
  const logger = options.logger.child({ service: 'badges-backend' });
  const config = await loadBackendConfig({ logger, argv: process.argv });
  const discovery = SingleHostDiscovery.fromConfig(config);
+  const database = useHotMemoize(module, () => {
+    return Knex({
+      client: 'better-sqlite3',
+      connection: ':memory:',
+      useNullAsDefault: true,
+    });
+  });

  logger.debug('Creating application...');

@@ -60,6 +70,9 @@ export async function startStandaloneServer(
    tokenManager,
    logger,
    identity,
+    db: await DatabaseBadgesStore.create({
+      database: { getClient: async () => database },
+    }),
  });

  let service = createServiceBuilder(module)
@@ -0,0 +1,294 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * 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.
+ */
+
+import { DatabaseBadgesStore } from '../database/badgesStore';
+import { TestDatabaseId, TestDatabases } from '@backstage/backend-test-utils';
+import { Knex } from 'knex';
+import { Entity } from '@backstage/catalog-model';
+import { ConfigReader } from '@backstage/config';
+import { GetEntitiesResponse } from '@backstage/catalog-client';
+import crypto from 'crypto';
+
+describe('DatabaseBadgesStore', () => {
+  const config = new ConfigReader({
+    backend: {
+      baseUrl: 'http://127.0.0.1',
+      listen: {
+        port: 7007,
+      },
+    },
+    app: {
+      badges: {
+        obfuscate: true,
+      },
+    },
+    custom: {
+      'badges-backend': {
+        salt: 'random-string',
+        cacheTimeToLive: '60',
+      },
+    },
+  });
+
+  const entity: Entity = {
+    apiVersion: 'v1',
+    kind: 'Component',
+    metadata: {
+      name: 'test',
+    },
+  };
+
+  const entities: Entity[] = [
+    entity,
+    {
+      apiVersion: 'v1',
+      kind: 'Component',
+      metadata: {
+        name: 'test-2',
+      },
+    },
+    {
+      apiVersion: 'v1',
+      kind: 'Component',
+      metadata: {
+        name: 'test-3',
+      },
+    },
+  ];
+
+  const defaultEntityListResponse: GetEntitiesResponse = {
+    items: entities,
+  };
+
+  const databases = TestDatabases.create();
+
+  async function createDatabaseBadgesStore(databaseId: TestDatabaseId) {
+    const knex = await databases.init(databaseId);
+    return {
+      knex,
+      badgeStore: await DatabaseBadgesStore.create({
+        database: { getClient: async () => knex },
+      }),
+    };
+  }
+
+  describe.each(databases.eachSupportedId())('%p', databaseId => {
+    let knex: Knex;
+    let badgeStore: DatabaseBadgesStore;
+
+    beforeEach(async () => {
+      ({ knex, badgeStore } = await createDatabaseBadgesStore(databaseId));
+    });
+
+    it('createAllBadges', async () => {
+      await badgeStore.createAllBadges(
+        defaultEntityListResponse,
+        config.getString('custom.badges-backend.salt'),
+      );
+
+      const storedBadges = await badgeStore.getAllBadges();
+      const testedEntities: {
+        name: string;
+        namespace: string;
+        kind: string;
+        hash: string;
+      }[] = [];
+
+      entities.forEach(entityInLoop => {
+        const kind = entityInLoop.kind.toLowerCase();
+        const namespace = entityInLoop.metadata.namespace || 'default';
+        const name = entityInLoop.metadata.name.toLocaleLowerCase();
+        const salt = config.getString('custom.badges-backend.salt');
+        const entityHash = crypto
+          .createHash('sha256')
+          .update(`${kind}:${namespace}:${name}:${salt}`)
+          .digest('hex');
+
+        testedEntities.push({
+          hash: entityHash,
+          name: name,
+          namespace: namespace,
+          kind: kind,
+        });
+        expect(storedBadges).toEqual(expect.arrayContaining(testedEntities));
+        expect(storedBadges.length).toEqual(3);
+      });
+    });
+
+    it('getBadgeFromHash', async () => {
+      await knex('badges').truncate();
+      await knex('badges').insert([
+        {
+          hash: 'hash1',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+      ]);
+
+      const storedEntity = await badgeStore.getBadgeFromHash('hash1');
+
+      expect(storedEntity).toEqual({
+        name: 'test',
+        namespace: 'default',
+        kind: 'component',
+      });
+    });
+
+    it('getHashFromEntityMetadata', async () => {
+      await knex('badges').truncate();
+      await knex('badges').insert([
+        {
+          hash: 'hash1',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+      ]);
+
+      const storedHash = await badgeStore.getHashFromEntityMetadata(
+        'test',
+        'default',
+        'component',
+      );
+
+      expect(storedHash).toEqual({
+        hash: 'hash1',
+      });
+    });
+
+    it('deleteObsoleteHashes', async () => {
+      const entityHash = crypto
+        .createHash('sha256')
+        .update(
+          `component:default:test:${config.getString(
+            'custom.badges-backend.salt',
+          )}`,
+        )
+        .digest('hex');
+
+      await knex('badges').insert([
+        {
+          hash: 'obsoleteHash',
+          name: 'obsoleteEntity',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: entityHash,
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+      ]);
+
+      await badgeStore.deleteObsoleteHashes(
+        defaultEntityListResponse,
+        config.getString('custom.badges-backend.salt'),
+      );
+
+      const storedBadges = await badgeStore.getAllBadges();
+
+      expect(storedBadges).toEqual(
+        expect.not.arrayContaining([
+          {
+            hash: 'obsoleteHash',
+            name: 'obsoleteEntity',
+            namespace: 'default',
+            kind: 'component',
+          },
+        ]),
+      );
+      expect(storedBadges.length).toEqual(1);
+      expect(storedBadges).toEqual(
+        expect.arrayContaining([
+          {
+            hash: entityHash,
+            name: 'test',
+            namespace: 'default',
+            kind: 'component',
+          },
+        ]),
+      );
+    });
+
+    it('countAllBadges', async () => {
+      await knex('badges').truncate();
+      await knex('badges').insert([
+        {
+          hash: 'hash1',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash2',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash3',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash4',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+      ]);
+
+      const storedBadgesCount = await badgeStore.countAllBadges();
+      expect(storedBadgesCount).toEqual(4);
+    });
+
+    it('getAllBadges', async () => {
+      await knex('badges').truncate();
+      await knex('badges').insert([
+        {
+          hash: 'hash1',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash2',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash3',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+        {
+          hash: 'hash4',
+          name: 'test',
+          namespace: 'default',
+          kind: 'component',
+        },
+      ]);
+
+      const storedBadgesCount = await badgeStore.countAllBadges();
+      expect(storedBadgesCount).toEqual(4);
+    });
+  });
+});
@@ -31,10 +31,16 @@ import {
  BackstageIdentityResponse,
  IdentityApiGetIdentityRequest,
 } from '@backstage/plugin-auth-node';
+import { BadgesStore } from '../database/badgesStore';
+import crypto from 'crypto';

 describe('createRouter', () => {
  let app: express.Express;
-  let badgeBuilder: jest.Mocked<BadgeBuilder>;
+  const badgeBuilder: jest.Mocked<BadgeBuilder> = {
+    getBadges: jest.fn(),
+    createBadgeJson: jest.fn(),
+    createBadgeSvg: jest.fn(),
+  };

  const catalog = {
    addLocation: jest.fn(),
@@ -49,10 +55,46 @@ describe('createRouter', () => {
    getEntityFacets: jest.fn(),
    validateEntity: jest.fn(),
  };
-  let config: Config;
-  let discovery: PluginEndpointDiscovery;
+  const getIdentity = jest
+    .fn()
+    .mockImplementation(
+      async ({
+        request: _request,
+      }: IdentityApiGetIdentityRequest): Promise<
+        BackstageIdentityResponse | undefined
+      > => {
+        return {
+          identity: {
+            userEntityRef: 'user:default/guest',
+            ownershipEntityRefs: [],
+            type: 'user',
+          },
+          token: 'token',
+        };
+      },
+    );

-  const getIdentity = jest.fn();
+  const config: Config = new ConfigReader({
+    backend: {
+      baseUrl: 'http://127.0.0.1',
+      listen: {
+        port: 7007,
+      },
+    },
+    app: {
+      badges: {
+        obfuscate: true,
+      },
+    },
+    custom: {
+      'badges-backend': {
+        salt: 'random-string',
+        cacheTimeToLive: '60',
+      },
+    },
+  });
+
+  let discovery: PluginEndpointDiscovery;

  const entity: Entity = {
    apiVersion: 'v1',
@@ -83,49 +125,47 @@ describe('createRouter', () => {
    markdown: '[![...](...)]',
  };

+  const badgeEntity = {
+    name: 'test',
+    namespace: 'default',
+    kind: 'component',
+  };
+
+  const badgeEntities = [
+    {
+      hash: 'hash1',
+      name: 'test',
+      namespace: 'default',
+      kind: 'component',
+    },
+    {
+      hash: 'hash2',
+      name: 'test2',
+      namespace: 'default',
+      kind: 'component',
+    },
+  ];
+
+  const badgeStore: jest.Mocked<BadgesStore> = {
+    createAllBadges: jest.fn(),
+    getBadgeFromHash: jest.fn().mockImplementation(async () => {
+      return badgeEntity;
+    }),
+    getHashFromEntityMetadata: jest
+      .fn()
+      .mockImplementation(async () => 'niceHash'),
+    deleteObsoleteHashes: jest.fn(),
+    countAllBadges: jest.fn().mockImplementation(async () => 4),
+    getAllBadges: jest.fn().mockImplementation(async () => badgeEntities),
+  };
+
+  const salt = config.getString('custom.badges-backend.salt');
+  const entityHash = crypto
+    .createHash('sha256')
+    .update(`component:default:test:${salt}`)
+    .digest('hex');
+
  beforeAll(async () => {
-    getIdentity.mockImplementation(
-      async ({
-        request: _request,
-      }: IdentityApiGetIdentityRequest): Promise<
-        BackstageIdentityResponse | undefined
-      > => {
-        return {
-          identity: {
-            userEntityRef: 'user:default/guest',
-            ownershipEntityRefs: [],
-            type: 'user',
-          },
-          token: 'token',
-        };
-      },
-    );
-
-    badgeBuilder = {
-      getBadges: jest.fn(),
-      createBadgeJson: jest.fn(),
-      createBadgeSvg: jest.fn(),
-    };
-    config = new ConfigReader({
-      backend: {
-        baseUrl: 'http://127.0.0.1',
-        listen: {
-          port: 7007,
-        },
-      },
-      app: {
-        badges: {
-          obfuscate: true,
-        },
-      },
-      custom: {
-        'badges-backend': {
-          salt: 'random-string',
-          cacheTimeToLive: '60',
-        },
-      },
-    });
-
    discovery = SingleHostDiscovery.fromConfig(config);
    const tokenManager = ServerTokenManager.noop();
    const router = await createRouter({
@@ -136,12 +176,13 @@ describe('createRouter', () => {
      tokenManager,
      logger: getVoidLogger(),
      identity: { getIdentity },
+      db: badgeStore,
    });
    app = express().use(router);
  });

  beforeEach(() => {
-    jest.resetAllMocks();
+    jest.clearAllMocks();
  });

  it('works', async () => {
@@ -154,6 +195,7 @@ describe('createRouter', () => {
      tokenManager,
      logger: getVoidLogger(),
      identity: { getIdentity },
+      db: badgeStore,
    });
    expect(router).toBeDefined();
  });
@@ -206,9 +248,7 @@ describe('createRouter', () => {
      badgeBuilder.getBadges.mockResolvedValueOnce([{ id: badge.id }]);
      badgeBuilder.createBadgeJson.mockResolvedValueOnce(badge);

-      const response = await request(app).get(
-        '/entity/3a5f91c1e66519be5394c37a8ba69c3087b7c322c600e7497dc9d517353e5bed/badge-specs',
-      );
+      const response = await request(app).get(`/entity/hash1/badge-specs`);
      expect(response.status).toEqual(200);
      expect(response.body).toEqual([badge]);

@@ -228,7 +268,7 @@ describe('createRouter', () => {
        badgeInfo: { id: badge.id },
        context: {
          badgeUrl: expect.stringMatching(
-            /http:\/\/127.0.0.1\/api\/badges\/entity\/3a5f91c1e66519be5394c37a8ba69c3087b7c322c600e7497dc9d517353e5bed\/test-badge/,
+            /http:\/\/127.0.0.1\/api\/badges\/entity\/hash1\/test-badge/,
          ),
          config,
          entity,
@@ -293,6 +333,9 @@ describe('createRouter', () => {
  describe('GET /entity/:namespace/:kind/:name/obfuscated', () => {
    catalog.getEntityByRef.mockResolvedValueOnce(entity);
    catalog.getEntities.mockResolvedValueOnce({ items: entities });
+    badgeStore.getHashFromEntityMetadata.mockResolvedValue({
+      hash: entityHash,
+    });

    it('returns obfuscated entity', async () => {
      const obfuscatedEntity = await request(app)
@@ -301,7 +344,7 @@ describe('createRouter', () => {
      expect(obfuscatedEntity.status).toEqual(200);
      // echo -n  "component:default:test:random-string" | openssl dgst -sha256
      expect(obfuscatedEntity.body).toEqual({
-        hash: '3a5f91c1e66519be5394c37a8ba69c3087b7c322c600e7497dc9d517353e5bed',
+        hash: entityHash,
      });
    });

@@ -315,6 +358,7 @@ describe('createRouter', () => {

  describe('Errors', () => {
    it('returns 404 for unknown entity hash', async () => {
+      badgeStore.getBadgeFromHash.mockResolvedValue(undefined);
      catalog.getEntityByRef.mockResolvedValueOnce(entity);
      catalog.getEntities.mockResolvedValueOnce({ items: entities });
      badgeBuilder.getBadges.mockResolvedValueOnce([{ id: badge.id }]);
@@ -31,6 +31,7 @@ import {
  BackstageIdentityResponse,
  IdentityApiGetIdentityRequest,
 } from '@backstage/plugin-auth-node';
+import { BadgesStore } from '../database/badgesStore';

 describe('createRouter', () => {
  let app: express.Express;
@@ -73,6 +74,15 @@ describe('createRouter', () => {
    markdown: '[![...](...)]',
  };

+  const badgeStore: jest.Mocked<BadgesStore> = {
+    createAllBadges: jest.fn(),
+    getBadgeFromHash: jest.fn(),
+    getHashFromEntityMetadata: jest.fn(),
+    deleteObsoleteHashes: jest.fn(),
+    countAllBadges: jest.fn(),
+    getAllBadges: jest.fn(),
+  };
+
  beforeAll(async () => {
    badgeBuilder = {
      getBadges: jest.fn(),
@@ -113,6 +123,7 @@ describe('createRouter', () => {
      tokenManager,
      logger: getVoidLogger(),
      identity: { getIdentity },
+      db: badgeStore,
    });
    app = express().use(router);
  });
@@ -131,6 +142,7 @@ describe('createRouter', () => {
      tokenManager,
      logger: getVoidLogger(),
      identity: { getIdentity },
+      db: badgeStore,
    });
    expect(router).toBeDefined();
  });