From 84020b3305b2c919ea3bfb0f3b1ef20ca439d9fa Mon Sep 17 00:00:00 2001
From: Andrey Sobolev <haiodo@users.noreply.github.com>
Date: Wed, 25 Dec 2024 16:08:11 +0700
Subject: [PATCH] UBERF-9017: Reduce createTable calls (#7550)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
---
 dev/tool/src/db.ts             |  2 +-
 server/postgres/src/storage.ts | 13 +++++++++++--
 server/postgres/src/utils.ts   | 34 ++++++++++++++++++++++++----------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/dev/tool/src/db.ts b/dev/tool/src/db.ts
index ed4db3d78a..30ecccb7dd 100644
--- a/dev/tool/src/db.ts
+++ b/dev/tool/src/db.ts
@@ -81,7 +81,7 @@ async function moveWorkspace (
       tables = tables.filter((t) => include.has(t))
     }
 
-    await createTables(new MeasureMetricsContext('', {}), pgClient, tables)
+    await createTables(new MeasureMetricsContext('', {}), pgClient, '', tables)
     const token = generateToken(systemAccountEmail, wsId)
     const endpoint = await getTransactorEndpoint(token, 'external')
     const connection = (await connect(endpoint, wsId, undefined, {
diff --git a/server/postgres/src/storage.ts b/server/postgres/src/storage.ts
index 32d82b533d..ba02249c4f 100644
--- a/server/postgres/src/storage.ts
+++ b/server/postgres/src/storage.ts
@@ -37,6 +37,7 @@ import core, {
   type ModelDb,
   type ObjQueryType,
   type Projection,
+  RateLimiter,
   type Ref,
   type ReverseLookups,
   type SessionData,
@@ -1506,13 +1507,18 @@ interface OperationBulk {
   mixins: TxMixin<Doc, Doc>[]
 }
 
+const initRateLimit = new RateLimiter(1)
+
 class PostgresAdapter extends PostgresAdapterBase {
   async init (ctx: MeasureContext, domains?: string[], excludeDomains?: string[]): Promise<void> {
     let resultDomains = domains ?? this.hierarchy.domains()
     if (excludeDomains !== undefined) {
       resultDomains = resultDomains.filter((it) => !excludeDomains.includes(it))
     }
-    await createTables(ctx, this.client, resultDomains)
+    const url = this.refClient.url()
+    await initRateLimit.exec(async () => {
+      await createTables(ctx, this.client, url, resultDomains)
+    })
     this._helper.domains = new Set(resultDomains as Domain[])
   }
 
@@ -1789,7 +1795,10 @@ class PostgresAdapter extends PostgresAdapterBase {
 class PostgresTxAdapter extends PostgresAdapterBase implements TxAdapter {
   async init (ctx: MeasureContext, domains?: string[], excludeDomains?: string[]): Promise<void> {
     const resultDomains = domains ?? [DOMAIN_TX, DOMAIN_MODEL_TX]
-    await createTables(ctx, this.client, resultDomains)
+    await initRateLimit.exec(async () => {
+      const url = this.refClient.url()
+      await createTables(ctx, this.client, url, resultDomains)
+    })
     this._helper.domains = new Set(resultDomains as Domain[])
   }
 
diff --git a/server/postgres/src/utils.ts b/server/postgres/src/utils.ts
index 0d985d3f0a..f05927f68c 100644
--- a/server/postgres/src/utils.ts
+++ b/server/postgres/src/utils.ts
@@ -72,8 +72,13 @@ export const NumericTypes = [
   core.class.Collection
 ]
 
-export async function createTables (ctx: MeasureContext, client: postgres.Sql, domains: string[]): Promise<void> {
-  const filtered = domains.filter((d) => !loadedDomains.has(d))
+export async function createTables (
+  ctx: MeasureContext,
+  client: postgres.Sql,
+  url: string,
+  domains: string[]
+): Promise<void> {
+  const filtered = domains.filter((d) => !loadedDomains.has(url + translateDomain(d)))
   if (filtered.length === 0) {
     return
   }
@@ -90,17 +95,15 @@ export async function createTables (ctx: MeasureContext, client: postgres.Sql, d
   const exists = new Set(tables.map((it) => it.table_name))
 
   await retryTxn(client, async (client) => {
-    await ctx.with('load-schemas', {}, () =>
-      getTableSchema(
-        client,
-        mapped.filter((it) => exists.has(it))
-      )
-    )
+    const domainsToLoad = mapped.filter((it) => exists.has(it))
+    if (domainsToLoad.length > 0) {
+      await ctx.with('load-schemas', {}, () => getTableSchema(client, domainsToLoad))
+    }
     for (const domain of mapped) {
       if (!exists.has(domain)) {
         await ctx.with('create-table', {}, () => createTable(client, domain))
       }
-      loadedDomains.add(domain)
+      loadedDomains.add(url + domain)
     }
   })
 }
@@ -188,6 +191,8 @@ export async function shutdown (): Promise<void> {
 export interface PostgresClientReference {
   getClient: () => Promise<postgres.Sql>
   close: () => void
+
+  url: () => string
 }
 
 class PostgresClientReferenceImpl {
@@ -195,6 +200,7 @@ class PostgresClientReferenceImpl {
   client: postgres.Sql | Promise<postgres.Sql>
 
   constructor (
+    readonly connectionString: string,
     client: postgres.Sql | Promise<postgres.Sql>,
     readonly onclose: () => void
   ) {
@@ -202,6 +208,10 @@ class PostgresClientReferenceImpl {
     this.client = client
   }
 
+  url (): string {
+    return this.connectionString
+  }
+
   async getClient (): Promise<postgres.Sql> {
     if (this.client instanceof Promise) {
       this.client = await this.client
@@ -233,6 +243,10 @@ export class ClientRef implements PostgresClientReference {
     clientRefs.set(this.id, this)
   }
 
+  url (): string {
+    return this.client.url()
+  }
+
   closed = false
   async getClient (): Promise<postgres.Sql> {
     if (!this.closed) {
@@ -274,7 +288,7 @@ export function getDBClient (connectionString: string, database?: string): Postg
       ...extraOptions
     })
 
-    existing = new PostgresClientReferenceImpl(sql, () => {
+    existing = new PostgresClientReferenceImpl(connectionString, sql, () => {
       connections.delete(key)
     })
     connections.set(key, existing)