@lark-apaas/nestjs-datapaas 1.0.17 → 1.0.19

package/dist/index.cjs

@@ -225,6 +225,7 @@ var DrizzleDatabaseManager = class {
   requestContext;
   db = null;
   sqlClient = null;
+  connectionOpPromise = null;
   currentConnectionInfo = null;
   isConnected = false;
   skipInitDbConnection = false;
@@ -233,7 +234,7 @@ var DrizzleDatabaseManager = class {
   tokenWatchFilePath = null;
   tokenWatchDebounceTimer = null;
   refreshingByTokenChange = false;
-  reconnectingPromise = null;
+  sqlClientShutdownTimeoutSeconds = 5;
   /**
   * 构造函数：注入配置、日志能力与请求上下文。
   */
@@ -269,32 +270,30 @@ var DrizzleDatabaseManager = class {
       };
       return this.db;
     }
-    try {
-      const resolvedConnectionString = await this.buildConnectionStringWithToken();
-      const { connectionString, connectionConfig } = parseConnectionInfo(resolvedConnectionString, this.config);
-      const baseSql = (0, import_postgres.default)(connectionString, connectionConfig);
-      const sqlProxy = genPostgresProxy(baseSql, this.observable, this.requestContext);
-      this.db = (0, import_postgres_js.drizzle)(sqlProxy, {
-        schema: this.config.schema,
-        logger: this.config.logger
-      });
-      this.sqlClient = baseSql;
-      this.currentResolvedConnectionString = connectionString;
-      this.currentConnectionInfo = {
-        connectionString,
-        connectionConfig: {
-          ssl: connectionConfig.ssl,
-          connectionTimeout: connectionConfig.connect_timeout,
-          idleTimeout: connectionConfig.idle_timeout,
-          maxConnections: connectionConfig.max
-        }
-      };
-      this.isConnected = true;
+    if (this.db && this.isConnected) {
       return this.db;
-    } catch (error) {
-      this.isConnected = false;
-      throw new Error(`Failed to initialize database connection: ${error instanceof Error ? error.message : "Unknown error"}`);
     }
+    if (this.connectionOpPromise) {
+      return this.connectionOpPromise;
+    }
+    return this.runWithConnectionLock(async () => {
+      if (this.db && this.isConnected) {
+        return this.db;
+      }
+      try {
+        const token = await this.readConnectionTokenFromFile();
+        const resolvedConnectionString = this.appendTokenToConnectionString(this.config.connectionString, token);
+        const nextState = await this.createConnectedState(resolvedConnectionString, token);
+        this.applyConnectedState(nextState);
+        return nextState.db;
+      } catch (error) {
+        if (!this.db || !this.isConnected) {
+          this.isConnected = false;
+          this.db = null;
+        }
+        throw new Error(`Failed to initialize database connection: ${error instanceof Error ? error.message : "Unknown error"}`);
+      }
+    });
   }
   /**
   * 获取数据库实例
@@ -327,22 +326,10 @@ var DrizzleDatabaseManager = class {
   * 重新连接数据库。
   * 注意：重连时不关闭 token watcher，避免重连后丢失后续文件变更监听。
   * 同时通过 promise 锁避免并发重连导致状态抖动。
-  * @param forceRefresh 是否强制刷新连接信息
   * @returns 新的数据库实例
   */
-  async reconnect(_forceRefresh = false) {
-    if (this.reconnectingPromise) {
-      return this.reconnectingPromise;
-    }
-    this.reconnectingPromise = (async () => {
-      await this.disconnect(false);
-      return this.initialize();
-    })();
-    try {
-      return await this.reconnectingPromise;
-    } finally {
-      this.reconnectingPromise = null;
-    }
+  async reconnect() {
+    return this.reconnectInternal();
   }
   /**
   * 断开数据库连接
@@ -351,21 +338,19 @@ var DrizzleDatabaseManager = class {
     if (shouldStopWatcher) {
       this.stopTokenWatcher();
     }
-    try {
-      if (this.sqlClient && typeof this.sqlClient.end === "function") {
-        await this.sqlClient.end({
-          timeout: 5
-        });
-      }
-    } catch (error) {
-      console.error("Error disconnecting from database:", error);
-    } finally {
-      this.sqlClient = null;
-      this.db = null;
-      this.isConnected = false;
-      this.currentConnectionInfo = null;
-      this.currentResolvedConnectionString = null;
+    const pendingOp = this.connectionOpPromise;
+    if (pendingOp) {
+      await Promise.allSettled([
+        pendingOp
+      ]);
     }
+    await this.shutdownSqlClientSafely(this.sqlClient, "Error disconnecting from database:");
+    this.sqlClient = null;
+    this.db = null;
+    this.isConnected = false;
+    this.currentConnectionInfo = null;
+    this.currentResolvedConnectionString = null;
+    this.currentToken = null;
   }
   /**
   * 获取当前连接信息
@@ -435,13 +420,80 @@ var DrizzleDatabaseManager = class {
     url.searchParams.set("options", nextOptions);
     return url.toString();
   }
-  /**
-  * 读取当前 token 并构造最终连接串，同时更新内存中的 currentToken。
-  */
-  async buildConnectionStringWithToken() {
-    const token = await this.readConnectionTokenFromFile();
-    this.currentToken = token;
-    return this.appendTokenToConnectionString(this.config.connectionString, token);
+  async createConnectedState(resolvedConnectionString, token) {
+    const { connectionString, connectionConfig } = parseConnectionInfo(resolvedConnectionString, this.config);
+    const baseSql = (0, import_postgres.default)(connectionString, connectionConfig);
+    try {
+      const sqlProxy = genPostgresProxy(baseSql, this.observable, this.requestContext);
+      const db = (0, import_postgres_js.drizzle)(sqlProxy, {
+        schema: this.config.schema,
+        logger: this.config.logger
+      });
+      return {
+        db,
+        sqlClient: baseSql,
+        resolvedConnectionString: connectionString,
+        token,
+        connectionInfo: {
+          connectionString,
+          connectionConfig: {
+            ssl: connectionConfig.ssl,
+            connectionTimeout: connectionConfig.connect_timeout,
+            idleTimeout: connectionConfig.idle_timeout,
+            maxConnections: connectionConfig.max
+          }
+        }
+      };
+    } catch (error) {
+      await this.shutdownSqlClientSafely(baseSql, "Error closing failed new database client:");
+      throw error;
+    }
+  }
+  applyConnectedState(nextState) {
+    this.db = nextState.db;
+    this.sqlClient = nextState.sqlClient;
+    this.currentConnectionInfo = nextState.connectionInfo;
+    this.currentResolvedConnectionString = nextState.resolvedConnectionString;
+    this.currentToken = nextState.token;
+    this.isConnected = true;
+  }
+  reconnectInternal(nextResolvedConnectionString, nextToken) {
+    return this.runWithConnectionLock(() => this.reconnectWithResolvedConnectionString(nextResolvedConnectionString, nextToken));
+  }
+  runWithConnectionLock(operation) {
+    if (this.connectionOpPromise) {
+      return this.connectionOpPromise;
+    }
+    this.connectionOpPromise = (async () => {
+      try {
+        return await operation();
+      } finally {
+        this.connectionOpPromise = null;
+      }
+    })();
+    return this.connectionOpPromise;
+  }
+  async reconnectWithResolvedConnectionString(nextResolvedConnectionString, nextToken) {
+    const resolvedToken = nextToken ?? await this.readConnectionTokenFromFile();
+    const resolvedConnectionString = nextResolvedConnectionString ?? this.appendTokenToConnectionString(this.config.connectionString, resolvedToken);
+    const nextState = await this.createConnectedState(resolvedConnectionString, resolvedToken);
+    const previousSqlClient = this.sqlClient;
+    this.applyConnectedState(nextState);
+    if (previousSqlClient && previousSqlClient !== nextState.sqlClient) {
+      await this.shutdownSqlClientSafely(previousSqlClient, "Error shutting down previous database client after reconnect:");
+    }
+    return nextState.db;
+  }
+  async shutdownSqlClientSafely(client, errorPrefix) {
+    try {
+      if (client && typeof client.end === "function") {
+        await client.end({
+          timeout: this.sqlClientShutdownTimeoutSeconds
+        });
+      }
+    } catch (error) {
+      console.error(errorPrefix, error);
+    }
   }
   /**
   * 启动 token 文件监听；若已存在旧监听则先清理后重建。
@@ -513,7 +565,7 @@ var DrizzleDatabaseManager = class {
       if (sameToken && sameConnectionString && this.isConnected) {
         return;
       }
-      await this.reconnect(true);
+      await this.reconnectInternal(nextResolvedConnectionString, nextToken);
     } catch (error) {
       console.error("Failed to refresh database connection from token file:", error);
     } finally {
@@ -580,11 +632,10 @@ var DataPaasDatabaseService = class {
   }
   /**
   * 重新连接数据库
-  * @param forceRefresh 是否强制刷新连接信息
   * @returns 新的数据库实例
   */
-  async reconnect(forceRefresh = false) {
-    return this.databaseManager.reconnect(forceRefresh);
+  async reconnect() {
+    return this.databaseManager.reconnect();
   }
   /**
   * 获取当前连接信息
@@ -632,7 +683,20 @@ DataPaasDatabaseService = _ts_decorate2([
 
 // src/database/sql-exection-context.middleware.ts
 var import_common4 = require("@nestjs/common");
-var session = __toESM(require("drizzle-orm/postgres-js"), 1);
+
+// src/database/drizzle-auth-context.ts
+var import_node_async_hooks = require("async_hooks");
+var asyncLocalStorage = new import_node_async_hooks.AsyncLocalStorage();
+function getAuthContextStore() {
+  return asyncLocalStorage.getStore();
+}
+__name(getAuthContextStore, "getAuthContextStore");
+function runWithAuthContext(store, fn) {
+  return asyncLocalStorage.run(store, fn);
+}
+__name(runWithAuthContext, "runWithAuthContext");
+
+// src/database/sql-exection-context.middleware.ts
 function _ts_decorate3(decorators, target, key, desc) {
   var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
   if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
@@ -675,7 +739,7 @@ var SqlExecutionContextMiddleware = class {
       // 用户类型
       userType ? `SET LOCAL app.user_type = '${userType}'` : `SET LOCAL app.user_type = ''`
     ].filter(Boolean).join(";");
-    session.runWithAuthContext({
+    runWithAuthContext({
       preSql: sqls
     }, () => {
       next();
@@ -694,6 +758,49 @@ SqlExecutionContextMiddleware = _ts_decorate3([
 // src/database/database.module.ts
 var import_nestjs_common3 = require("@lark-apaas/nestjs-common");
 var import_nestjs_observable = require("@lark-apaas/nestjs-observable");
+
+// src/database/drizzle-monkey-patch.ts
+var patched = false;
+function applyDrizzleMonkeyPatch() {
+  if (patched) return;
+  const pgModule = require("drizzle-orm/postgres-js");
+  const PreparedQuery = pgModule.PostgresJsPreparedQuery;
+  if (!PreparedQuery?.prototype?.execute) {
+    throw new Error("drizzle-orm PostgresJsPreparedQuery.prototype.execute \u4E0D\u5B58\u5728\uFF0C\u7248\u672C\u53EF\u80FD\u4E0D\u517C\u5BB9");
+  }
+  const originalExecute = PreparedQuery.prototype.execute;
+  PreparedQuery.prototype.execute = async function(...args) {
+    const authContext = getAuthContextStore();
+    const preSql = authContext?.preSql;
+    const postSql = authContext?.postSql;
+    if (!preSql && !postSql) {
+      return originalExecute.apply(this, args);
+    }
+    const client = this.client;
+    if (typeof client.begin === "function") {
+      return client.begin(async (tx) => {
+        const savedClient = this.client;
+        this.client = tx;
+        try {
+          if (preSql) await tx.unsafe(preSql);
+          const result2 = await originalExecute.apply(this, args);
+          if (postSql) await tx.unsafe(postSql);
+          return result2;
+        } finally {
+          this.client = savedClient;
+        }
+      });
+    }
+    if (preSql) await client.unsafe(preSql);
+    const result = await originalExecute.apply(this, args);
+    if (postSql) await client.unsafe(postSql);
+    return result;
+  };
+  patched = true;
+}
+__name(applyDrizzleMonkeyPatch, "applyDrizzleMonkeyPatch");
+
+// src/database/database.module.ts
 function _ts_decorate4(decorators, target, key, desc) {
   var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
   if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
@@ -701,20 +808,7 @@ function _ts_decorate4(decorators, target, key, desc) {
   return c > 3 && r && Object.defineProperty(target, key, r), r;
 }
 __name(_ts_decorate4, "_ts_decorate");
-function verifyDrizzlePatched() {
-  try {
-    const postgres2 = require("drizzle-orm/postgres-js");
-    const ok = postgres2 && typeof postgres2.getAuthContextStore === "function" && typeof postgres2.runWithAuthContext === "function";
-    if (!ok) {
-      throw new Error("drizzle-orm \u672A\u68C0\u6D4B\u5230\u8865\u4E01\u5BFC\u51FA\uFF1AgetAuthContextStore/runWithAuthContext");
-    }
-  } catch (err) {
-    const reason = err instanceof Error ? err.message : String(err);
-    throw new Error(`drizzle-orm \u8865\u4E01\u672A\u751F\u6548\u6216\u7248\u672C\u4E0D\u5339\u914D\uFF1A${reason}`);
-  }
-}
-__name(verifyDrizzlePatched, "verifyDrizzlePatched");
-verifyDrizzlePatched();
+applyDrizzleMonkeyPatch();
 var DataPaasModule = class _DataPaasModule {
   static {
     __name(this, "DataPaasModule");