From 17e23813bf47595bba5d639e856c5910f9733821 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/query/CHANGELOG.md                    |  2 +-
 packages/query/src/graphql/graphql.module.ts   |  2 +-
 .../graphql/plugins/PgSubscriptionPlugin.ts    | 18 ++++++++++++------
 3 files changed, 14 insertions(+), 8 deletions(-)

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..d6187b755a 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,17 @@ 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 ({id}, args, context, resolveInfo) => {
+          const [row] = await resolveInfo.graphile.selectGraphQLResultFromTable(
+            sql.identifier(table.namespace.name, table.name),
+            (tableAlias, queryBuilder) => {
+              queryBuilder.context.args = {blockHeight: sql.fragment`'999999999999999999'::bigint`};
+              queryBuilder.where(sql.fragment`${tableAlias}.id = ${sql.value(id)}`);
+              queryBuilder.limit(1);
+            }
+          );
+
+          return row;
         },
       },
     };