From cdb30b5ce4f821392c9a532bda3b55fc362ab3f0 Mon Sep 17 00:00:00 2001
From: Samuel Hurel <samuel.hurel@gmail.com>
Date: Thu, 12 Dec 2024 21:22:20 +0100
Subject: [PATCH] fix: load subscriptions relations

---
 packages/node-core/CHANGELOG.md               |  1 +
 packages/node-core/src/db/sync-helper.ts      |  6 +++---
 packages/query/CHANGELOG.md                   |  2 +-
 packages/query/src/graphql/graphql.module.ts  |  2 +-
 .../graphql/plugins/PgSubscriptionPlugin.ts   | 19 +++++++++++++------
 5 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/packages/node-core/CHANGELOG.md b/packages/node-core/CHANGELOG.md
index a52618403a..7fe35eb0f4 100644
--- a/packages/node-core/CHANGELOG.md
+++ b/packages/node-core/CHANGELOG.md
@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+- Updated send_notification PG function to include `_block height` and entity `_id` (#2626)
 
 ## [16.1.0] - 2024-12-11
 ### Changed
diff --git a/packages/node-core/src/db/sync-helper.ts b/packages/node-core/src/db/sync-helper.ts
index c0d4e73acd..7d2a3d7b2d 100644
--- a/packages/node-core/src/db/sync-helper.ts
+++ b/packages/node-core/src/db/sync-helper.ts
@@ -194,10 +194,10 @@ BEGIN
             'mutation_type', TG_OP,
             '_entity', row);
     IF payload -> '_entity' ? '_block_range' then
-      payload = payload #- '{"_entity","_id"}';
-      payload = payload #- '{"_entity","_block_range"}';
+        payload = payload #- '{"_entity","_block_range"}';
+        payload = payload || jsonb_build_object('_block_height', lower(row._block_range));
         IF NOT upper_inf(row._block_range) then
-         -- Check if a newer version of the entity exists to determine operation
+          -- Check if a newer version of the entity exists to determine operation
           EXECUTE FORMAT(
             'SELECT EXISTS (SELECT 1 FROM "${schema}".%I WHERE id = $1 AND lower(_block_range) = upper($2))',
             TG_TABLE_NAME
diff --git a/packages/query/CHANGELOG.md b/packages/query/CHANGELOG.md
index e444509c87..a4ff354f18 100644
--- a/packages/query/CHANGELOG.md
+++ b/packages/query/CHANGELOG.md
@@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
-- Subscriptions `_entity` field now includes camel cased properties (#2626)
+- Subscriptions `_entity` field now returns all properties (#2626)
 
 ## [2.19.0] - 2024-12-11
 ### Added
diff --git a/packages/query/src/graphql/graphql.module.ts b/packages/query/src/graphql/graphql.module.ts
index 6b47e021de..7c39378ba2 100644
--- a/packages/query/src/graphql/graphql.module.ts
+++ b/packages/query/src/graphql/graphql.module.ts
@@ -220,7 +220,7 @@ export class GraphqlModule implements OnModuleInit, OnModuleDestroy {
         path: WS_ROUTE,
       });
 
-      this.wsCleanup = useServer({schema}, wsServer);
+      this.wsCleanup = useServer({schema, context: {pgClient: this.pgPool}}, wsServer);
     }
 
     app.use(PinoLogger(PinoConfig));
diff --git a/packages/query/src/graphql/plugins/PgSubscriptionPlugin.ts b/packages/query/src/graphql/plugins/PgSubscriptionPlugin.ts
index 19550720b0..5e61d52341 100644
--- a/packages/query/src/graphql/plugins/PgSubscriptionPlugin.ts
+++ b/packages/query/src/graphql/plugins/PgSubscriptionPlugin.ts
@@ -30,7 +30,7 @@ function makePayload(entityType: string): {type: DocumentNode; name: string} {
 }
 
 export const PgSubscriptionPlugin = makeExtendSchemaPlugin((build) => {
-  const {inflection, pgIntrospectionResultsByKind} = build;
+  const {inflection, pgIntrospectionResultsByKind, pgSql: sql} = build;
 
   const typeDefs = [
     gql`
@@ -68,11 +68,18 @@ export const PgSubscriptionPlugin = makeExtendSchemaPlugin((build) => {
 
     resolvers[payloadName] = {
       _entity: {
-        resolve: ({_entity}) => {
-          return Object.entries(_entity).reduce((acc, [key, value]) => {
-            const attr = table.attributes.find((attr) => attr.name === key);
-            return Object.assign(acc, {[inflection.column(attr)]: value});
-          }, _entity);
+        resolve: async ({_block_height, _entity}, args, context, resolveInfo) => {
+          const [row] = await resolveInfo.graphile.selectGraphQLResultFromTable(
+            sql.identifier(table.namespace.name, table.name),
+            (tableAlias, queryBuilder) => {
+              queryBuilder.context.args ??= {};
+              queryBuilder.context.args.blockHeight = sql.fragment`${sql.value(_block_height.toString())}::bigint`;
+              queryBuilder.where(sql.fragment`${tableAlias}._id = ${sql.value(_entity._id)}`);
+              queryBuilder.limit(1);
+            }
+          );
+
+          return row;
         },
       },
     };