adding d1 session support

app/src/adapter/cloudflare/cloudflare-workers.adapter.ts

@@ -9,7 +9,13 @@ import { getDurable } from "./modes/durable";
 import type { App } from "bknd";
 import { $console } from "core";
 
-export type CloudflareEnv = object;
+declare global {
+   namespace Cloudflare {
+      interface Env {}
+   }
+}
+
+export type CloudflareEnv = Cloudflare.Env;
 export type CloudflareBkndConfig<Env = CloudflareEnv> = RuntimeBkndConfig<Env> & {
    mode?: "warm" | "fresh" | "cache" | "durable";
    bindings?: (args: Env) => {
@@ -17,6 +23,11 @@ export type CloudflareBkndConfig<Env = CloudflareEnv> = RuntimeBkndConfig<Env> &
       dobj?: DurableObjectNamespace;
       db?: D1Database;
    };
+   d1?: {
+      session?: boolean;
+      transport?: "header" | "cookie";
+      first?: D1SessionConstraint;
+   };
    static?: "kv" | "assets";
    key?: string;
    keepAliveSeconds?: number;

app/src/adapter/cloudflare/config.ts

@@ -1,11 +1,14 @@
+/// <reference types="@cloudflare/workers-types" />
+
 import { registerMedia } from "./storage/StorageR2Adapter";
 import { getBinding } from "./bindings";
-import { D1Connection } from "./D1Connection";
+import { D1Connection } from "./connection/D1Connection";
 import type { CloudflareBkndConfig, CloudflareEnv } from ".";
 import { App } from "bknd";
 import { makeConfig as makeAdapterConfig } from "bknd/adapter";
-import type { ExecutionContext } from "hono";
+import type { Context, ExecutionContext } from "hono";
 import { $console } from "core";
+import { setCookie } from "hono/cookie";
 
 export const constants = {
    exec_async_event_id: "cf_register_waituntil",
@@ -13,25 +16,96 @@ export const constants = {
    do_endpoint: "/__bknd/do",
 };
 
+export type CfMakeConfigArgs<Env extends CloudflareEnv = CloudflareEnv> = {
+   env: Env;
+   ctx?: ExecutionContext;
+   request?: Request;
+};
+
+function getCookieValue(cookies: string | null, name: string) {
+   if (!cookies) return null;
+
+   for (const cookie of cookies.split("; ")) {
+      const [key, value] = cookie.split("=");
+      if (key === name && value) {
+         return decodeURIComponent(value);
+      }
+   }
+   return null;
+}
+
+export function d1SessionHelper(config: CloudflareBkndConfig<any>) {
+   const headerKey = "x-cf-d1-session";
+   const cookieKey = "cf_d1_session";
+   const transport = config.d1?.transport;
+
+   return {
+      get: (request?: Request): D1SessionBookmark | undefined => {
+         if (!request || !config.d1?.session) return undefined;
+
+         if (!transport || transport === "cookie") {
+            const cookies = request.headers.get("Cookie");
+            if (cookies) {
+               const cookie = getCookieValue(cookies, cookieKey);
+               if (cookie) {
+                  return cookie;
+               }
+            }
+         }
+
+         if (!transport || transport === "header") {
+            if (request.headers.has(headerKey)) {
+               return request.headers.get(headerKey) as any;
+            }
+         }
+
+         return undefined;
+      },
+      set: (c: Context, d1?: D1DatabaseSession) => {
+         if (!d1 || !config.d1?.session) return;
+
+         const session = d1.getBookmark();
+         if (session) {
+            if (!transport || transport === "header") {
+               c.header(headerKey, session);
+            }
+            if (!transport || transport === "cookie") {
+               setCookie(c, cookieKey, session, {
+                  httpOnly: true,
+                  secure: true,
+                  sameSite: "Lax",
+                  maxAge: 60 * 5, // 5 minutes
+               });
+            }
+         }
+      },
+   };
+}
+
 let media_registered: boolean = false;
 export function makeConfig<Env extends CloudflareEnv = CloudflareEnv>(
    config: CloudflareBkndConfig<Env>,
-   args: Env = {} as Env,
+   args: CfMakeConfigArgs<Env>,
 ) {
    if (!media_registered) {
       registerMedia(args as any);
       media_registered = true;
    }
 
-   const appConfig = makeAdapterConfig(config, args);
-   const bindings = config.bindings?.(args);
+   const appConfig = makeAdapterConfig(config, args.env);
+   const bindings = config.bindings?.(args.env);
+
+   const sessionHelper = d1SessionHelper(config);
+   const sessionId = sessionHelper.get(args.request);
+   let session: D1DatabaseSession | undefined;
+
    if (!appConfig.connection) {
       let db: D1Database | undefined;
       if (bindings?.db) {
          $console.log("Using database from bindings");
          db = bindings.db;
       } else if (Object.keys(args).length > 0) {
-         const binding = getBinding(args, "D1Database");
+         const binding = getBinding(args.env, "D1Database");
          if (binding) {
             $console.log(`Using database from env "${binding.key}"`);
             db = binding.value;
@@ -39,12 +113,32 @@ export function makeConfig<Env extends CloudflareEnv = CloudflareEnv>(
       }
 
       if (db) {
-         appConfig.connection = new D1Connection({ binding: db });
+         if (config.d1?.session) {
+            session = db.withSession(sessionId ?? config.d1?.first);
+            appConfig.connection = new D1Connection({ binding: session });
+         } else {
+            appConfig.connection = new D1Connection({ binding: db });
+         }
       } else {
          throw new Error("No database connection given");
       }
    }
 
+   if (config.d1?.session) {
+      appConfig.options = {
+         ...appConfig.options,
+         manager: {
+            ...appConfig.options?.manager,
+            onServerInit: (server) => {
+               server.use(async (c, next) => {
+                  sessionHelper.set(c, session);
+                  await next();
+               });
+            },
+         },
+      };
+   }
+
    return appConfig;
 }
 

app/src/adapter/cloudflare/connection/D1Connection.ts

@@ -5,8 +5,8 @@ import type { QB } from "data/connection/Connection";
 import { type DatabaseIntrospector, Kysely, ParseJSONResultsPlugin } from "kysely";
 import { D1Dialect } from "kysely-d1";
 
-export type D1ConnectionConfig = {
-   binding: D1Database;
+export type D1ConnectionConfig<DB extends D1Database | D1DatabaseSession = D1Database> = {
+   binding: DB;
 };
 
 class CustomD1Dialect extends D1Dialect {
@@ -17,22 +17,24 @@ class CustomD1Dialect extends D1Dialect {
    }
 }
 
-export class D1Connection extends SqliteConnection {
+export class D1Connection<
+   DB extends D1Database | D1DatabaseSession = D1Database,
+> extends SqliteConnection {
    protected override readonly supported = {
       batching: true,
    };
 
-   constructor(private config: D1ConnectionConfig) {
+   constructor(private config: D1ConnectionConfig<DB>) {
       const plugins = [new ParseJSONResultsPlugin()];
 
       const kysely = new Kysely({
-         dialect: new CustomD1Dialect({ database: config.binding }),
+         dialect: new CustomD1Dialect({ database: config.binding as D1Database }),
          plugins,
       });
       super(kysely, {}, plugins);
    }
 
-   get client(): D1Database {
+   get client(): DB {
       return this.config.binding;
    }
 

app/src/adapter/cloudflare/index.ts

@@ -1,4 +1,4 @@
-import { D1Connection, type D1ConnectionConfig } from "./D1Connection";
+import { D1Connection, type D1ConnectionConfig } from "./connection/D1Connection";
 
 export * from "./cloudflare-workers.adapter";
 export { makeApp, getFresh } from "./modes/fresh";

app/src/adapter/cloudflare/modes/cached.ts

@@ -5,8 +5,9 @@ import { makeConfig, registerAsyncsExecutionContext, constants } from "../config
 
 export async function getCached<Env extends CloudflareEnv = CloudflareEnv>(
    config: CloudflareBkndConfig<Env>,
-   { env, ctx, ...args }: Context<Env>,
+   args: Context<Env>,
 ) {
+   const { env, ctx } = args;
    const { kv } = config.bindings?.(env)!;
    if (!kv) throw new Error("kv namespace is not defined in cloudflare.bindings");
    const key = config.key ?? "app";
@@ -20,7 +21,7 @@ export async function getCached<Env extends CloudflareEnv = CloudflareEnv>(
 
    const app = await createRuntimeApp(
       {
-         ...makeConfig(config, env),
+         ...makeConfig(config, args),
          initialConfig,
          onBuilt: async (app) => {
             registerAsyncsExecutionContext(app, ctx);
@@ -41,7 +42,7 @@ export async function getCached<Env extends CloudflareEnv = CloudflareEnv>(
             await config.beforeBuild?.(app);
          },
       },
-      { env, ctx, ...args },
+      args,
    );
 
    if (!cachedConfig) {

app/src/adapter/cloudflare/modes/fresh.ts

@@ -1,13 +1,13 @@
 import { createRuntimeApp, type RuntimeOptions } from "bknd/adapter";
 import type { CloudflareBkndConfig, Context, CloudflareEnv } from "../index";
-import { makeConfig, registerAsyncsExecutionContext } from "../config";
+import { makeConfig, registerAsyncsExecutionContext, type CfMakeConfigArgs } from "../config";
 
 export async function makeApp<Env extends CloudflareEnv = CloudflareEnv>(
    config: CloudflareBkndConfig<Env>,
-   args: Env = {} as Env,
+   args: CfMakeConfigArgs<Env>,
    opts?: RuntimeOptions,
 ) {
-   return await createRuntimeApp<Env>(makeConfig(config, args), args, opts);
+   return await createRuntimeApp<Env>(makeConfig(config, args), args.env, opts);
 }
 
 export async function getFresh<Env extends CloudflareEnv = CloudflareEnv>(
@@ -23,7 +23,7 @@ export async function getFresh<Env extends CloudflareEnv = CloudflareEnv>(
             await config.onBuilt?.(app);
          },
       },
-      ctx.env,
+      ctx,
       opts,
    );
 }