From d0ed9ee2e0a74d4ea20e01ef0fd7dd6bfd122f6c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=A9my=20M?= <jeremy.magrin@gmail.com>
Date: Fri, 12 Jan 2024 10:41:38 +0100
Subject: [PATCH] feat: pagination with total count (#3384)

* feat: add totalCount

* feat: add command for production to fix existing tables
---
 .../commands/database-command.module.ts       |  2 +
 .../workspace-add-total-count.command.ts      | 49 +++++++++++++++++++
 .../workspace-migration-runner.service.ts     |  5 ++
 .../connection-type-definition.factory.ts     |  7 ++-
 4 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 packages/twenty-server/src/database/commands/workspace-add-total-count.command.ts

diff --git a/packages/twenty-server/src/database/commands/database-command.module.ts b/packages/twenty-server/src/database/commands/database-command.module.ts
index daa3e2986..20b8163f2 100644
--- a/packages/twenty-server/src/database/commands/database-command.module.ts
+++ b/packages/twenty-server/src/database/commands/database-command.module.ts
@@ -10,6 +10,7 @@ import { DataSeedDemoWorkspaceCommand } from 'src/database/commands/data-seed-de
 import { WorkspaceDataSourceModule } from 'src/workspace/workspace-datasource/workspace-datasource.module';
 import { WorkspaceSyncMetadataModule } from 'src/workspace/workspace-sync-metadata/workspace-sync-metadata.module';
 import { ObjectMetadataModule } from 'src/metadata/object-metadata/object-metadata.module';
+import { WorkspaceAddTotalCountCommand } from 'src/database/commands/workspace-add-total-count.command';
 
 @Module({
   imports: [
@@ -24,6 +25,7 @@ import { ObjectMetadataModule } from 'src/metadata/object-metadata/object-metada
   providers: [
     DataSeedWorkspaceCommand,
     DataSeedDemoWorkspaceCommand,
+    WorkspaceAddTotalCountCommand,
     ConfirmationQuestion,
   ],
 })
diff --git a/packages/twenty-server/src/database/commands/workspace-add-total-count.command.ts b/packages/twenty-server/src/database/commands/workspace-add-total-count.command.ts
new file mode 100644
index 000000000..cdf8aa293
--- /dev/null
+++ b/packages/twenty-server/src/database/commands/workspace-add-total-count.command.ts
@@ -0,0 +1,49 @@
+import { Command, CommandRunner } from 'nest-commander';
+import chalk from 'chalk';
+
+import { TypeORMService } from 'src/database/typeorm/typeorm.service';
+
+@Command({
+  name: 'workspace:add-total-count',
+  description: 'Add pg_graphql total count directive to all workspace tables',
+})
+export class WorkspaceAddTotalCountCommand extends CommandRunner {
+  constructor(private readonly typeORMService: TypeORMService) {
+    super();
+  }
+
+  async run(): Promise<void> {
+    const mainDataSource = this.typeORMService.getMainDataSource();
+
+    try {
+      await mainDataSource.query(`
+        DO $$
+        DECLARE
+            schema_cursor CURSOR FOR SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE 'workspace_%';
+            schema_name text;
+            table_rec record;
+        BEGIN
+            OPEN schema_cursor;
+            LOOP
+                FETCH schema_cursor INTO schema_name;
+                EXIT WHEN NOT FOUND;
+                
+                FOR table_rec IN SELECT t.table_name FROM information_schema.tables t WHERE t.table_schema = schema_name
+                LOOP
+                    EXECUTE 'COMMENT ON TABLE ' || quote_ident(schema_name) || '.' || quote_ident(table_rec.table_name) || ' IS e''@graphql({"totalCount": {"enabled": true}})'';';
+                END LOOP;
+            END LOOP;
+            CLOSE schema_cursor;
+        END $$;    
+      `);
+
+      console.log(
+        chalk.green('Total count directive added to all workspace tables'),
+      );
+    } catch (error) {
+      console.log(
+        chalk.red('Error adding total count directive to all workspace tables'),
+      );
+    }
+  }
+}
diff --git a/packages/twenty-server/src/workspace/workspace-migration-runner/workspace-migration-runner.service.ts b/packages/twenty-server/src/workspace/workspace-migration-runner/workspace-migration-runner.service.ts
index f07c363f3..1ea26d2d3 100644
--- a/packages/twenty-server/src/workspace/workspace-migration-runner/workspace-migration-runner.service.ts
+++ b/packages/twenty-server/src/workspace/workspace-migration-runner/workspace-migration-runner.service.ts
@@ -140,6 +140,11 @@ export class WorkspaceMigrationRunnerService {
       }),
       true,
     );
+
+    // Enable totalCount for the table
+    await queryRunner.query(`
+      COMMENT ON TABLE "${schemaName}"."${tableName}" IS '@graphql({"totalCount": {"enabled": true}})';
+    `);
   }
 
   /**
diff --git a/packages/twenty-server/src/workspace/workspace-schema-builder/factories/connection-type-definition.factory.ts b/packages/twenty-server/src/workspace/workspace-schema-builder/factories/connection-type-definition.factory.ts
index 50e3ccd75..5d40ca1db 100644
--- a/packages/twenty-server/src/workspace/workspace-schema-builder/factories/connection-type-definition.factory.ts
+++ b/packages/twenty-server/src/workspace/workspace-schema-builder/factories/connection-type-definition.factory.ts
@@ -1,6 +1,6 @@
 import { Injectable, Logger } from '@nestjs/common';
 
-import { GraphQLFieldConfigMap, GraphQLObjectType } from 'graphql';
+import { GraphQLFieldConfigMap, GraphQLInt, GraphQLObjectType } from 'graphql';
 
 import { WorkspaceBuildSchemaOptions } from 'src/workspace/workspace-schema-builder/interfaces/workspace-build-schema-optionts.interface';
 import { ObjectMetadataInterface } from 'src/metadata/field-metadata/interfaces/object-metadata.interface';
@@ -71,6 +71,11 @@ export class ConnectionTypeDefinitionFactory {
       ),
     };
 
+    fields.totalCount = {
+      type: GraphQLInt,
+      description: 'Total number of records in the connection',
+    };
+
     return fields;
   }
 }