feat: direct downloads

Author: LouisHaftmann

lib/schemas.ts

@@ -48,6 +48,7 @@ export const envBaseSchema = type({
   'CACHE_CLEANUP_OLDER_THAN_DAYS': 'number = 30',
   'DISABLE_CLEANUP_JOBS?': 'boolean',
   'DEBUG?': 'boolean',
+  'ENABLE_DIRECT_DOWNLOADS': 'boolean = false',
 })
 
 export const envSchema = envBaseSchema.and(envStorageDriverSchema).and(envDbDriverSchema)

lib/storage.ts

@@ -19,6 +19,7 @@ import {
   S3Client,
 } from '@aws-sdk/client-s3'
 import { Upload as S3Upload } from '@aws-sdk/lib-storage'
+import { getSignedUrl } from '@aws-sdk/s3-request-presigner'
 import { Storage as GcsClient } from '@google-cloud/storage'
 import { chunk } from 'remeda'
 import { match } from 'ts-pattern'
@@ -262,7 +263,7 @@ class Storage {
     return { id: uploadId }
   }
 
-  async getCacheEntry({
+  private async getCacheEntryByKeys({
     keys: [primaryKey, ...restoreKeys],
     version,
   }: {
@@ -310,6 +311,35 @@ class Storage {
       if (prefixedMatch) return prefixedMatch
     }
   }
+
+  async getCacheEntryWithDownloadUrl(args: Parameters<typeof this.getCacheEntryByKeys>[0]) {
+    const cacheEntry = await this.getCacheEntryByKeys(args)
+    if (!cacheEntry) return
+
+    const defaultUrl = `${env.API_BASE_URL}/download/${cacheEntry.id}`
+
+    if (!env.ENABLE_DIRECT_DOWNLOADS || !this.adapter.createDownloadUrl)
+      return {
+        downloadUrl: defaultUrl,
+        cacheEntry,
+      }
+
+    const location = await this.db
+      .selectFrom('storage_locations')
+      .where('id', '=', cacheEntry.locationId)
+      .select(['folderName', 'mergedAt'])
+      .executeTakeFirst()
+    if (!location) throw new Error('Storage location not found')
+
+    const downloadUrl = location.mergedAt
+      ? await this.adapter.createDownloadUrl(`${location.folderName}/merged`)
+      : defaultUrl
+
+    return {
+      downloadUrl,
+      cacheEntry,
+    }
+  }
 }
 
 export const getStorage = createSingletonPromise(async () => Storage.fromEnv())
@@ -319,6 +349,7 @@ interface StorageAdapter {
   uploadStream(objectName: string, stream: ReadableStream): Promise<void>
   deleteFolder(folderName: string): Promise<void>
   countFilesInFolder(folderName: string): Promise<number>
+  createDownloadUrl?(objectName: string): Promise<string>
 }
 
 class S3Adapter implements StorageAdapter {
@@ -417,6 +448,19 @@ class S3Adapter implements StorageAdapter {
 
     return listResponse.KeyCount ?? 0
   }
+
+  async createDownloadUrl(objectName: string) {
+    return getSignedUrl(
+      this.s3,
+      new GetObjectCommand({
+        Bucket: this.bucket,
+        Key: `${this.keyPrefix}/${objectName}`,
+      }),
+      {
+        expiresIn: 10 * 60 * 1000, // 10min
+      },
+    )
+  }
 }
 
 class FileSystemAdapter implements StorageAdapter {
@@ -513,4 +557,14 @@ class GcsAdapter implements StorageAdapter {
       })
       .then((res) => res[0].length)
   }
+
+  async createDownloadUrl(objectName: string) {
+    return this.bucket
+      .file(`${this.keyPrefix}/${objectName}`)
+      .getSignedUrl({
+        action: 'read',
+        expires: Date.now() + 10 * 60 * 1000, // 10min
+      })
+      .then((res) => res[0])
+  }
 }

routes/twirp/github.actions.results.api.v1.CacheService/GetCacheEntryDownloadURL.post.ts

@@ -1,5 +1,4 @@
 import { z } from 'zod'
-import { env } from '~/lib/env'
 import { getStorage } from '~/lib/storage'
 
 const bodySchema = z.object({
@@ -19,18 +18,18 @@ export default defineEventHandler(async (event) => {
   const { key, restore_keys, version } = parsedBody.data
 
   const storage = await getStorage()
-  const cacheEntry = await storage.getCacheEntry({
+  const match = await storage.getCacheEntryWithDownloadUrl({
     keys: [key, ...(restore_keys ?? [])],
     version,
   })
-  if (!cacheEntry)
+  if (!match)
     return {
       ok: false,
     }
 
   return {
     ok: true,
-    signed_download_url: `${env.API_BASE_URL}/download/${cacheEntry.id}`,
-    matched_key: cacheEntry.key,
+    signed_download_url: match.downloadUrl,
+    matched_key: match.cacheEntry.key,
   }
 })

tests/setup.ts

@@ -29,7 +29,10 @@ const TESTING_ENV_BASE = {
   RUNNER_TEMP: path.join(TEST_TEMP_DIR, 'runner-temp'),
   ACTIONS_RESULTS_URL: 'http://localhost:3000/',
   ACTIONS_CACHE_URL: 'http://localhost:3000/',
-} satisfies Omit<typeof envBaseSchema.infer, 'CACHE_CLEANUP_OLDER_THAN_DAYS'> &
+} satisfies Omit<
+  typeof envBaseSchema.infer,
+  'CACHE_CLEANUP_OLDER_THAN_DAYS' | 'ENABLE_DIRECT_DOWNLOADS'
+> &
   Record<string, string>
 
 const TESTING_ENV_BY_DB_DRIVER = {