gencow 0.1.153 → 0.1.155

package/core/index.js

@@ -3229,6 +3229,10 @@ async function forEachSetConfig(rls, setOne) {
   }
 }
 var rlsExecClient = new AsyncLocalStorage();
+function isBunSqlSession(session) {
+  const name = session?.constructor?.name;
+  return typeof name === "string" && name.includes("BunSQLSession");
+}
 function isDrizzleTransactionDb(db) {
   const d = db;
   if (typeof d?.rollback === "function") {
@@ -3296,40 +3300,80 @@ async function withRlsConnection(session, rls, reuseOuterConnection, fn) {
   }
   return runInner(c);
 }
-function wrapPreparedQuery(pq, session, rls, reuseOuterConnection) {
+function createClientBoundSession(session, client) {
+  return new Proxy(session, {
+    get(target, prop, receiver) {
+      if (prop === "client") {
+        return client;
+      }
+      const value = Reflect.get(target, prop, receiver);
+      return typeof value === "function" ? value.bind(receiver) : value;
+    }
+  });
+}
+async function withPreparedQueryClient(pq, client, fn) {
+  if (!("client" in pq)) {
+    return await fn();
+  }
+  const prevPq = pq.client;
+  pq.client = client;
+  try {
+    return await fn();
+  } finally {
+    pq.client = prevPq;
+  }
+}
+function createGenericPreparedQueryClientExecutors(pq) {
   const origExecute = pq.execute.bind(pq);
   const origAll = typeof pq.all === "function" ? pq.all.bind(pq) : null;
-  const withPreparedQueryClient = async (client, fn) => {
-    if (!("client" in pq)) {
-      return fn();
-    }
-    const prevPq = pq.client;
-    pq.client = client;
-    try {
-      return await fn();
-    } finally {
-      pq.client = prevPq;
-    }
+  return {
+    execute(client, placeholderValues) {
+      return withPreparedQueryClient(pq, client, () => origExecute(placeholderValues));
+    },
+    all: origAll ? (client, placeholderValues) => withPreparedQueryClient(pq, client, () => origAll(placeholderValues)) : void 0
   };
-  pq.execute = async (placeholderValues) => {
-    const active = rlsExecClient.getStore();
-    if (active) {
-      return withPreparedQueryClient(active, () => origExecute(placeholderValues));
-    }
-    return withRlsConnection(session, rls, reuseOuterConnection, async (client) => {
-      return withPreparedQueryClient(client, () => origExecute(placeholderValues));
-    });
+}
+function createBunSqlPreparedQueryClientExecutors(baseSession, prepareQuery, prepareArgs) {
+  const reprepareOnClient = (client) => {
+    const clientBoundSession = createClientBoundSession(baseSession, client);
+    return Reflect.apply(prepareQuery, clientBoundSession, prepareArgs);
   };
-  if (origAll) {
-    pq.all = async (placeholderValues) => {
-      const active = rlsExecClient.getStore();
-      if (active) {
-        return withPreparedQueryClient(active, () => origAll(placeholderValues));
+  return {
+    execute(client, placeholderValues) {
+      return reprepareOnClient(client).execute(placeholderValues);
+    },
+    all(client, placeholderValues) {
+      const txPreparedQuery = reprepareOnClient(client);
+      if (typeof txPreparedQuery.all === "function") {
+        return txPreparedQuery.all(placeholderValues);
       }
-      return withRlsConnection(session, rls, reuseOuterConnection, async (client) => {
-        return withPreparedQueryClient(client, () => origAll(placeholderValues));
-      });
-    };
+      return txPreparedQuery.execute(placeholderValues);
+    }
+  };
+}
+function createPreparedQueryClientExecutors(params) {
+  if (isBunSqlSession(params.baseSession)) {
+    return createBunSqlPreparedQueryClientExecutors(params.baseSession, params.prepareQuery, params.prepareArgs);
+  }
+  return createGenericPreparedQueryClientExecutors(params.pq);
+}
+async function runPreparedQueryWithRls(session, rls, reuseOuterConnection, runOnClient, placeholderValues) {
+  const active = rlsExecClient.getStore();
+  if (active) {
+    return runOnClient(active, placeholderValues);
+  }
+  return withRlsConnection(session, rls, reuseOuterConnection, (client) => runOnClient(client, placeholderValues));
+}
+function wrapPreparedQuery(pq, session, baseSession, prepareQuery, prepareArgs, rls, reuseOuterConnection) {
+  const clientExecutors = createPreparedQueryClientExecutors({
+    pq,
+    baseSession,
+    prepareQuery,
+    prepareArgs
+  });
+  pq.execute = (placeholderValues) => runPreparedQueryWithRls(session, rls, reuseOuterConnection, clientExecutors.execute, placeholderValues);
+  if (clientExecutors.all) {
+    pq.all = (placeholderValues) => runPreparedQueryWithRls(session, rls, reuseOuterConnection, clientExecutors.all, placeholderValues);
   }
 }
 function wrapSession(session, rls, reuseOuterConnection) {
@@ -3342,7 +3386,7 @@ function wrapSession(session, rls, reuseOuterConnection) {
         return (...args) => {
           const prepareQuery = Reflect.get(sTarget, sProp, sRecv);
           const pq = Reflect.apply(prepareQuery, sRecv, args);
-          wrapPreparedQuery(pq, sRecv, rls, reuseOuterConnection);
+          wrapPreparedQuery(pq, sRecv, sTarget, prepareQuery, args, rls, reuseOuterConnection);
           return pq;
         };
       }