From 9de310a858c1457d428d28e877b1c1ae3c328285 Mon Sep 17 00:00:00 2001
From: Denis Carriere <carriere.denis@gmail.com>
Date: Sun, 8 Oct 2023 13:41:13 -0400
Subject: [PATCH] Validate moduleHashByChain fixes:
 https://github.com/pinax-network/substreams-sink-websockets/issues/19

---
 README.md                  | 15 ++++++++-------
 src/GET.ts                 |  1 +
 src/POST.ts                |  5 +++--
 src/banner.ts              |  1 +
 src/sqlite.ts              |  1 +
 src/websocket/subscribe.ts | 12 ++++++++++--
 6 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 0c52ccf..41b60e7 100644
--- a/README.md
+++ b/README.md
@@ -26,13 +26,14 @@ Send Text Message with the following JSON payloads
 
 ## REST API
 
-| Pathname          | Description           |
-|-------------------|-----------------------|
-| GET `/health`     | Health check
-| GET `/metrics`    | Prometheus metrics
-| GET `/chain`      | Returns all available `chain`
-| GET `/traceId`    | Returns all `traceId` by `chain`
-| GET `/moduleHash` | Returns all available `moduleHash` by `chain`
+| Pathname                  | Description           |
+|---------------------------|-----------------------|
+| GET `/health`             | Health check
+| GET `/metrics`            | Prometheus metrics
+| GET `/chain`              | Returns all available `chain`
+| GET `/traceId`            | Returns all `traceId` by `chain`
+| GET `/moduleHash`         | Returns all available `moduleHash`
+| GET `/moduleHashByChain`  | Returns all available `moduleHash` by `chain`
 | POST `/` {timestamp, signature, body}   | Webhook HTTP POST (Ed25519 signature)
 | POST `/` {"message": "PING"}            | Webhook HTTP POST Ping
 
diff --git a/src/GET.ts b/src/GET.ts
index d3700b0..e183893 100644
--- a/src/GET.ts
+++ b/src/GET.ts
@@ -22,6 +22,7 @@ export default async function (req: Request, server: Server) {
     if ( pathname === "/health") return checkHealth();
     if ( pathname === "/metrics") return new Response(await prometheus.registry.metrics());
     if ( pathname === "/moduleHash") return toJSON(sqlite.selectAll(db, "moduleHash"));
+    if ( pathname === "/moduleHashByChain") return toJSON(sqlite.selectAll(db, "moduleHashByChain"));
     if ( pathname === "/traceId") return toJSON(sqlite.selectAll(db, "traceId"));
     if ( pathname === "/chain") return toJSON(sqlite.selectAll(db, "chain"));
     return new Response("Not found", { status: 400 });
diff --git a/src/POST.ts b/src/POST.ts
index 276fbdc..16abb59 100644
--- a/src/POST.ts
+++ b/src/POST.ts
@@ -56,9 +56,10 @@ export default async function (req: Request, server: Server) {
     prometheus.trace_id.labels({traceId}).inc(1);
 
     // Upsert moduleHash into SQLite DB
-    sqlite.replace(db, "moduleHash", moduleHash, timestamp);
     sqlite.replace(db, "chain", chain, timestamp);
-    sqlite.replace(db, "traceId", traceId, timestamp);
+    sqlite.replace(db, "moduleHash", moduleHash, timestamp);
+    sqlite.replace(db, "moduleHashByChain", `${chain}:${moduleHash}`, timestamp);
+    sqlite.replace(db, "traceId", `${chain}:${traceId}`, timestamp);
 
     return new Response("OK");
 }
\ No newline at end of file
diff --git a/src/banner.ts b/src/banner.ts
index 5a19abf..e0885e9 100644
--- a/src/banner.ts
+++ b/src/banner.ts
@@ -20,6 +20,7 @@ export function banner()  {
         /health
         /metrics
         /moduleHash
+        /moduleHashByChain
         /traceId
         /chain
 
diff --git a/src/sqlite.ts b/src/sqlite.ts
index 9b357d0..31e65e4 100644
--- a/src/sqlite.ts
+++ b/src/sqlite.ts
@@ -14,6 +14,7 @@ export function createDb(filename: string) {
 
     // create tables if does not exists
     create(db, "moduleHash");
+    create(db, "moduleHashByChain");
     create(db, "traceId");
     create(db, "chain");
     return db;
diff --git a/src/websocket/subscribe.ts b/src/websocket/subscribe.ts
index 5d216bc..587554f 100644
--- a/src/websocket/subscribe.ts
+++ b/src/websocket/subscribe.ts
@@ -8,8 +8,16 @@ export default function (ws: ServerWebSocket<ServerWebSocketData>, params: {[key
     const { moduleHash, chain } = params;
     const topic = chain ? `${chain}:${moduleHash}` : moduleHash;
     if ( ws.isSubscribed(topic) ) throw new Error(`Already subscribed to [${topic}] topic.`);
-    if ( !sqlite.exists(db, "moduleHash", moduleHash) ) throw new Error(`ModuleHash [${moduleHash}] not found.`);
-    if ( chain && !sqlite.exists(db, "chain", chain) ) throw new Error(`Chain [${chain}] not found.`);
+
+    // Subscribe to ModuleHash by Chain
+    if ( chain ) {
+        if ( !sqlite.exists(db, "chain", chain) ) throw new Error(`Chain [${chain}] not found.`);
+        if ( !sqlite.exists(db, "moduleHashByChain", topic) ) throw new Error(`ModuleHash [${moduleHash}] from Chain [${chain}] not found.`);
+
+    // Subscribe to all ModuleHash
+    } else {
+        if ( !sqlite.exists(db, "moduleHash", moduleHash) ) throw new Error(`ModuleHash [${moduleHash}] not found.`);
+    }
     ws.subscribe(topic);
     logger.info('subscribed', {id, key: ws.data.key, remoteAddress: ws.remoteAddress, topic, params});
 }
\ No newline at end of file