feat: cache warmer implementation

cdn-server/cdn/src/index.ts

@@ -218,6 +218,36 @@ const draftRouterConfig = (storage: BlobStorage) => {
   };
 };
 
+const cacheOperations = (storage: BlobStorage) => {
+  return async (c: Context) => {
+    const organizationId = c.get('authenticatedOrganizationId');
+    const federatedGraphId = c.get('authenticatedFederatedGraphId');
+
+    if (organizationId !== c.req.param('organization_id') || federatedGraphId !== c.req.param('federated_graph_id')) {
+      return c.text('Bad Request', 400);
+    }
+
+    const key = `${organizationId}/${federatedGraphId}/cache_warmup/operations.json`;
+    let blobObject: BlobObject;
+
+    try {
+      blobObject = await storage.getObject({ context: c, key, cacheControl: 'no-cache' });
+    } catch (e: any) {
+      if (e instanceof BlobNotFoundError) {
+        return c.notFound();
+      }
+      throw e;
+    }
+
+    c.header('Content-Type', 'application/json; charset=UTF-8');
+    c.header('Cache-Control', 'no-cache, no-store, must-revalidate');
+
+    return stream(c, async (stream) => {
+      await stream.pipe(blobObject.stream);
+    });
+  };
+};
+
 // eslint-disable-next-line @typescript-eslint/ban-types
 export const cdn = <E extends Env, S extends Schema = {}, BasePath extends string = '/'>(
   hono: Hono<E, S, BasePath>,
@@ -235,4 +265,9 @@ export const cdn = <E extends Env, S extends Schema = {}, BasePath extends strin
   hono
     .use(draftRouterConfigs, jwtMiddleware(opts.authAdmissionJwtSecret))
     .get(draftRouterConfigs, draftRouterConfig(opts.blobStorage));
+
+  const cacheOperationsPath = '/:organization_id/:federated_graph_id/cache_warmup/operations.json';
+  hono
+    .use(cacheOperationsPath, jwtMiddleware(opts.authJwtSecret))
+    .get(cacheOperationsPath, cacheOperations(opts.blobStorage));
 };

cdn-server/cdn/test/cdn.test.ts

@@ -391,4 +391,101 @@ describe('CDN handlers', () => {
       expect(res.status).toBe(404);
     });
   });
+
+  describe('Test cache warmer opeartions handler', async () => {
+    const federatedGraphId = 'federatedGraphId';
+    const organizationId = 'organizationId';
+    const token = await generateToken(organizationId, federatedGraphId, secretKey);
+    const blobStorage = new InMemoryBlobStorage();
+    const requestPath = `${organizationId}/${federatedGraphId}/cache_warmup/operations.json`;
+
+    const app = new Hono();
+
+    cdn(app, {
+      authJwtSecret: secretKey,
+      authAdmissionJwtSecret: secretAdmissionKey,
+      blobStorage,
+    });
+
+    test('it returns a 401 if no Authorization header is provided', async () => {
+      const res = await app.request(requestPath, {
+        method: 'GET',
+      });
+      expect(res.status).toBe(401);
+    });
+
+    test('it returns a 401 if an invalid Authorization header is provided', async () => {
+      const res = await app.request(requestPath, {
+        method: 'GET',
+        headers: {
+          Authorization: `Bearer ${token.slice(0, -1)}}`,
+        },
+      });
+      expect(res.status).toBe(401);
+    });
+
+    test('it returns a 401 if the graph or organization ids does not match with the JWT payload', async () => {
+      const res = await app.request(`/foo/bar/operations/cache_warmup/operations.json`, {
+        method: 'GET',
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
+      });
+      expect(res.status).toBe(400);
+    });
+
+    test('it returns a 401 if the token has expired', async () => {
+      const token = await new SignJWT({
+        organization_id: organizationId,
+        federated_graph_id: federatedGraphId,
+        exp: Math.floor(Date.now() / 1000) - 60,
+      })
+        .setProtectedHeader({ alg: 'HS256' })
+        .sign(new TextEncoder().encode(secretKey));
+      const res = await app.request(`/foo/bar/cache_warmup/operations.json`, {
+        method: 'GET',
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
+      });
+      expect(res.status).toBe(401);
+    });
+
+    test('it returns the cache warmer operations', async () => {
+      const operationContents = JSON.stringify({
+        operations: [
+          {
+            request: {
+              operationName: 'AB',
+              query: 'query AB($a: Int!){employeeAsList(id: $a){tag id derivedMood products}}',
+            },
+            client: { name: 'unknown', version: 'missing' },
+          },
+        ],
+      });
+
+      blobStorage.objects.set(requestPath, {
+        buffer: Buffer.from(operationContents),
+      });
+
+      const res = await app.request(requestPath, {
+        method: 'GET',
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
+      });
+      expect(res.status).toBe(200);
+      expect(await res.text()).toBe(operationContents);
+    });
+
+    test('it returns a 404 if the persisted operation does not exist', async () => {
+      const res = await app.request(`${organizationId}/${federatedGraphId}/cache_warmup/does_not_exist.json`, {
+        method: 'GET',
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
+      });
+      expect(res.status).toBe(404);
+    });
+  });
 });

cli/src/commands/router/commands/cache/commands/push.ts

@@ -0,0 +1,84 @@
+import { existsSync } from 'node:fs';
+import { readFile } from 'node:fs/promises';
+import { Command, program } from 'commander';
+import pc from 'picocolors';
+import { resolve } from 'pathe';
+import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb';
+import { BaseCommandOptions } from '../../../../../core/types/types.js';
+import { getBaseHeaders } from '../../../../../core/config.js';
+
+export default (opts: BaseCommandOptions) => {
+  const command = new Command('push');
+  command.description('Pushes new cache warmer operations to the registry');
+  command.argument(
+    '<graph_name>',
+    'The name of the federated graph or monograph on which the check operations are stored.',
+  );
+  command.option('-n, --namespace [string]', 'The namespace of the federated graph or monograph.');
+  command.option(
+    '-o, --operation-name <operation-name>',
+    'The name of the operation. Only needed when working with multi-operation document.',
+  );
+  command.option(
+    '-p, --persisted-operation-id <persisted-operation-id>',
+    'The id of the persisted operation to be pushed.',
+  );
+  command.option(
+    '-f, --file <file>',
+    'The file with the operation to push - supports .graphql, .gql. If both the file and the persisted operation id are provided, the persisted operation id will be used.',
+  );
+
+  command.action(async (name, options) => {
+    if (!options.file && !options.persistedOperationId) {
+      command.error(pc.red('No operation file or the id of persisted operation provided'));
+    }
+
+    if (options.persistedOperationId && options.file) {
+      command.error(pc.red('Only the persisted operation id or the file can be specified'));
+    }
+
+    let operation: string | undefined;
+    if (options.file) {
+      const operationFile = resolve(options.file);
+      if (!existsSync(operationFile)) {
+        program.error(
+          pc.red(
+            pc.bold(
+              `The operation file '${pc.bold(operationFile)}' does not exist. Please check the path and try again.`,
+            ),
+          ),
+        );
+      }
+
+      const schemaBuffer = await readFile(operationFile);
+      operation = new TextDecoder().decode(schemaBuffer);
+      if (operation.trim().length === 0) {
+        program.error(
+          pc.red(pc.bold(`The schema file '${pc.bold(operationFile)}' is empty. Please provide a valid operation.`)),
+        );
+      }
+    }
+
+    const result = await opts.client.platform.pushCacheWarmerOperation(
+      {
+        federatedGraphName: name,
+        operationContent: operation,
+        operationName: options.operationName,
+        operationPersistedId: options.persistedOperationId,
+        namespace: options.namespace,
+      },
+      { headers: getBaseHeaders() },
+    );
+
+    if (result.response?.code === EnumStatusCode.OK) {
+      console.log(pc.green(`The cache warmer operation was pushed successfully.`));
+    } else {
+      console.log(pc.red(`Failed to push the cache warmer operation. Please try again.`));
+      if (result.response?.details) {
+        console.error(pc.red(pc.bold(result.response?.details)));
+      }
+      process.exit(1);
+    }
+  });
+  return command;
+};

cli/src/commands/router/commands/cache/index.ts

@@ -0,0 +1,14 @@
+import { Command } from 'commander';
+import { BaseCommandOptions } from '../../../../core/types/types.js';
+import PushCacheOperation from './commands/push.js';
+
+export default (opts: BaseCommandOptions) => {
+  const schema = new Command('cache');
+  schema.description(
+    'Provides commands for pushing and maintaining router cache warmer operations of a federated graph',
+  );
+
+  schema.addCommand(PushCacheOperation(opts));
+
+  return schema;
+};

cli/src/commands/router/index.ts

@@ -5,6 +5,7 @@ import ComposeRouterConfig from './commands/compose.js';
 import FetchRouterConfig from './commands/fetch.js';
 import RouterTokenCommands from './commands/token/index.js';
 import DownloadRouterBinaryConfig from './commands/download-binary.js';
+import RouterCacheCommands from './commands/cache/index.js';
 
 export default (opts: BaseCommandOptions) => {
   const cmd = new Command('router');
@@ -17,6 +18,11 @@ export default (opts: BaseCommandOptions) => {
     }),
   );
   cmd.addCommand(DownloadRouterBinaryConfig(opts));
+  cmd.addCommand(
+    RouterCacheCommands({
+      client: opts.client,
+    }),
+  );
 
   cmd.hook('preAction', async (thisCmd) => {
     if (['compose', 'download-binary'].includes(thisCmd.args[0])) {