Finish D1 instrumentation

packages/client-library-otel/src/cf-bindings.ts

@@ -4,8 +4,6 @@ import { errorToJson, safelySerializeJSON } from "./utils";
 // TODO - Can we use a Symbol here instead?
 const IS_PROXIED_KEY = "__fpx_proxied";
 
-
-
 /**
  * Proxy a Cloudflare binding to add instrumentation.
  * For now, just wraps all functions on the binding to use a measured version of the function.
@@ -38,14 +36,21 @@ export function proxyCloudflareBinding(o: unknown, bindingName: string) {
     return o;
   }
 
+  // HACK - Clean this up. Special logic for D1.
+  if (isCloudflareD1Binding(o)) {
+    return proxyD1Binding(o, bindingName);
+  }
+
   const proxiedBinding = new Proxy(o, {
     get(target, prop, receiver) {
       const value = Reflect.get(target, prop, receiver);
 
       if (typeof value === "function") {
         const methodName = String(prop);
+
         // OPTIMIZE - Do we want to do these lookups / this wrapping every time the property is accessed?
         const bindingType = getConstructorName(target);
+
         // Use the user's binding name, not the Cloudflare constructor name
         const name = `${bindingName}.${methodName}`;
         const measuredBinding = measure(
@@ -123,7 +128,7 @@ function isCloudflareD1Binding(o: unknown) {
 
 function isCloudflareKVBinding(o: unknown) {
   const constructorName = getConstructorName(o);
-  if (constructorName !== "KVNamespace") {
+  if (constructorName !== "KvNamespace") {
     return false;
   }
 
@@ -184,4 +189,71 @@ function markAsProxied(o: object) {
   });
 }
 
+/**
+ * Proxy a D1 binding to add instrumentation.
+ * 
+ * What this actually does is create a proxy of the `database` prop... I hope this works
+ *
+ * @param o - The D1Database binding to proxy
+ *
+ * @returns A proxied binding, whose `.database._send` and `.database._sendOrThrow` methods are instrumented
+ */
+function proxyD1Binding(o: unknown, bindingName: string) {
+  if (!o || typeof o !== "object") {
+    return o;
+  }
+
+  if (!isCloudflareD1Binding(o)) {
+    return o;
+  }
+
+  if (isAlreadyProxied(o)) {
+    return o;
+  }
+
+  const d1Proxy = new Proxy(o, {
+    get(d1Target, d1Prop) {
+      const d1Method = String(d1Prop);
+      const d1Value = Reflect.get(d1Target, d1Prop);
+      // HACK - These are technically public methods on the database object,
+      // but they have an underscore prefix which usually means "private" by convention...
+      // BEWARE!!!
+      const isSendingMethod = d1Method === "_send" || d1Method === "_sendOrThrow";
+      if (typeof d1Value === "function" && isSendingMethod) {
+        // ...
+        return measure({
+          name: "D1 Call",
+          attributes: {
+            "cf.binding.method": d1Method,
+            "cf.binding.name": bindingName,
+            "cf.binding.type": "D1Database",
+          },
+          onStart: (span, args) => {
+            span.setAttributes({
+              args: safelySerializeJSON(args),
+            });
+          },
+          // TODO - Use this callback to add additional attributes to the span regarding the response...
+          //        But the thing is, the result could be so wildly different depending on the method!
+          //        Might be good to proxy each binding individually, eventually?
+          //
+          // onSuccess: (span, result) => {},
+          onError: (span, error) => {
+            const serializableError =
+              error instanceof Error ? errorToJson(error) : error;
+            const errorAttributes = {
+              "cf.binding.error": safelySerializeJSON(serializableError),
+            };
+            span.setAttributes(errorAttributes);
+          },
+        }, d1Value.bind(d1Target));
+      }
+
+      return d1Value;
+    },
+  });
+
+  markAsProxied(d1Proxy);
 
+  return d1Proxy;
+}

sample-apps/goosify/drizzle/migrations/0001_omniscient_silver_fox.sql

@@ -0,0 +1,6 @@
+CREATE TABLE `goose_images` (
+	`id` integer PRIMARY KEY NOT NULL,
+	`filename` text NOT NULL,
+	`prompt` text NOT NULL,
+	`created_at` text DEFAULT (CURRENT_TIMESTAMP) NOT NULL
+);

sample-apps/goosify/drizzle/migrations/meta/0001_snapshot.json

@@ -0,0 +1,98 @@
+{
+  "version": "6",
+  "dialect": "sqlite",
+  "id": "9beaed9b-6535-4fd9-a201-6872fc1cf72d",
+  "prevId": "7a67a8d6-2ac2-4242-948d-048e03cba163",
+  "tables": {
+    "geese": {
+      "name": "geese",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "integer",
+          "primaryKey": true,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "avatar": {
+          "name": "avatar",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false,
+          "default": "(CURRENT_TIMESTAMP)"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false,
+          "default": "(CURRENT_TIMESTAMP)"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "goose_images": {
+      "name": "goose_images",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "integer",
+          "primaryKey": true,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "filename": {
+          "name": "filename",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "prompt": {
+          "name": "prompt",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false,
+          "default": "(CURRENT_TIMESTAMP)"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    }
+  },
+  "enums": {},
+  "_meta": {
+    "schemas": {},
+    "tables": {},
+    "columns": {}
+  }
+}

sample-apps/goosify/drizzle/migrations/meta/_journal.json

@@ -8,6 +8,13 @@
       "when": 1724758784931,
       "tag": "0000_new_human_robot",
       "breakpoints": true
+    },
+    {
+      "idx": 1,
+      "version": "6",
+      "when": 1724775919239,
+      "tag": "0001_omniscient_silver_fox",
+      "breakpoints": true
     }
   ]
 }

sample-apps/goosify/src/db/schema.ts

@@ -8,3 +8,10 @@ export const geese = sqliteTable("geese", {
   createdAt: text("created_at").notNull().default(sql`(CURRENT_TIMESTAMP)`),
   updatedAt: text("updated_at").notNull().default(sql`(CURRENT_TIMESTAMP)`),
 });
+
+export const gooseImages = sqliteTable("goose_images", {
+  id: integer("id", { mode: "number" }).primaryKey(),
+  filename: text("filename").notNull(),
+  prompt: text("prompt").notNull(),
+  createdAt: text("created_at").notNull().default(sql`(CURRENT_TIMESTAMP)`),
+});

sample-apps/goosify/src/index.ts

@@ -13,6 +13,26 @@ type Bindings = {
 
 const app = new Hono<{ Bindings: Bindings }>();
 
+const LATEST_LOCALE_KEY = "latest_locale";
+
+// Middleware to set the locale in kv-storage
+app.use(async (c, next) => {
+  console.log("Setting locale");
+  const storedLocale = await c.env.GOOSIFY_KV.get(LATEST_LOCALE_KEY);
+
+  let locale = c.req.header("Accept-Language") || "en";
+
+  // Optional: Parse the "Accept-Language" header to get the most preferred language
+  locale = parseAcceptLanguage(locale);
+
+  if (storedLocale !== locale) {
+    console.log(`Setting latest locale to ${locale}`);
+    await c.env.GOOSIFY_KV.put(LATEST_LOCALE_KEY, locale);
+  }
+
+  await next();
+});
+
 app.get("/", (c) => {
   return c.text("Hello Hono!");
 });
@@ -25,4 +45,39 @@ app.get("/api/geese", async (c) => {
   return c.json({ geese });
 });
 
+// TODO
+app.get("/api/cyberpunk-goose", async (c) => {
+  const inputs = {
+    prompt: "cyberpunk goose",
+  };
+  const cyberpunkGooseImage = await c.env.AI.run(
+    "@cf/lykon/dreamshaper-8-lcm",
+    inputs
+  );
+
+  const blob = new Blob([cyberpunkGooseImage], { type: 'image/png' });
+  const filename = `cyberpunk-goose--${crypto.randomUUID()}.png`;
+  await c.env.GOOSIFY_R2.put(filename, blob);
+
+  const db = drizzle(c.env.DB);
+  await db.insert(schema.gooseImages).values({
+    filename,
+    prompt: inputs.prompt,
+  });
+
+  c.header("Content-Type", "image/png");
+  return c.body(cyberpunkGooseImage);
+});
+
 export default instrument(app);
+
+
+function parseAcceptLanguage(acceptLanguage: string) {
+  // Simple parser to get the most preferred language
+  const locales = acceptLanguage.split(',').map(lang => {
+    const parts = lang.split(';');
+    return { locale: parts[0], q: parts[1] ? Number.parseFloat(parts[1].split('=')[1]) : 1 };
+  });
+  locales.sort((a, b) => b.q - a.q);
+  return locales[0].locale; // Return the most preferred language
+}