From 224d003d64e69498bec3e50bbbce50f4d4a09e93 Mon Sep 17 00:00:00 2001
From: Alexander Onnikov <aonnikov@hardcoreeng.com>
Date: Wed, 11 Dec 2024 19:42:39 +0700
Subject: [PATCH] UBERF-8884 Use direct connection to database (#7426)

Signed-off-by: Alexander Onnikov <Alexander.Onnikov@xored.com>
---
 workers/datalake/src/blob.ts               | 16 +++++++++++-----
 workers/datalake/src/db.ts                 |  6 ++++--
 workers/datalake/worker-configuration.d.ts |  1 +
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/workers/datalake/src/blob.ts b/workers/datalake/src/blob.ts
index 1c97a1675b..98cbf414e2 100644
--- a/workers/datalake/src/blob.ts
+++ b/workers/datalake/src/blob.ts
@@ -44,9 +44,13 @@ export async function handleBlobGet (
   const { workspace, name } = request
 
   const cache = new LoggedCache(caches.default, metrics)
-  const cached = await cache.match(request)
-  if (cached !== undefined) {
-    return cached
+
+  const cacheControl = request.headers.get('Cache-Control') ?? ''
+  if (!cacheControl.includes('no-cache')) {
+    const cached = await cache.match(request)
+    if (cached !== undefined) {
+      return cached
+    }
   }
 
   const { bucket } = selectStorage(env, workspace)
@@ -75,8 +79,10 @@ export async function handleBlobGet (
   const response = new Response(object?.body, { headers, status })
 
   if (response.status === 200) {
-    const clone = metrics.withSync('response.clone', () => response.clone())
-    ctx.waitUntil(cache.put(request, clone))
+    if (!cacheControl.includes('no-store')) {
+      const clone = metrics.withSync('response.clone', () => response.clone())
+      ctx.waitUntil(cache.put(request, clone))
+    }
   }
 
   return response
diff --git a/workers/datalake/src/db.ts b/workers/datalake/src/db.ts
index 8af838fdb9..91e9bfe282 100644
--- a/workers/datalake/src/db.ts
+++ b/workers/datalake/src/db.ts
@@ -50,12 +50,14 @@ export async function withPostgres<T> (
   fn: (db: BlobDB) => Promise<T>
 ): Promise<T> {
   const sql = metrics.withSync('db.connect', () => {
-    return postgres(env.HYPERDRIVE.connectionString, {
+    return postgres(env.DB_URL, {
       connection: {
         application_name: 'datalake'
-      }
+      },
+      fetch_types: false
     })
   })
+
   const db = new LoggedDB(new PostgresDB(sql), metrics)
 
   try {
diff --git a/workers/datalake/worker-configuration.d.ts b/workers/datalake/worker-configuration.d.ts
index dfc1e51e8e..532c6991ce 100644
--- a/workers/datalake/worker-configuration.d.ts
+++ b/workers/datalake/worker-configuration.d.ts
@@ -12,6 +12,7 @@ interface Env {
   STREAMS_ACCOUNT_ID: string;
   STREAMS_AUTH_KEY: string;
   R2_ACCOUNT_ID: string;
+  DB_URL: string;
   DATALAKE_APAC_ACCESS_KEY: string;
   DATALAKE_APAC_SECRET_KEY: string;
   DATALAKE_APAC_BUCKET_NAME: string;