Build a GraphQL API with DataLoader and Caching

// src/graphql/schema.ts — GraphQL schema with DataLoader-powered resolvers
import { createSchema, createYoga } from "graphql-yoga";
import DataLoader from "dataloader";
import { pool } from "../db";
import { Redis } from "ioredis";

const redis = new Redis(process.env.REDIS_URL!);

// DataLoaders — batch individual lookups into single queries
function createLoaders(tenantId: string) {
  return {
    user: new DataLoader<string, any>(async (ids) => {
      const { rows } = await pool.query(
        `SELECT * FROM users WHERE id = ANY($1) AND tenant_id = $2`,
        [ids as string[], tenantId]
      );
      const map = new Map(rows.map((r) => [r.id, r]));
      return ids.map((id) => map.get(id) || null);
    }),

    project: new DataLoader<string, any>(async (ids) => {
      const { rows } = await pool.query(
        `SELECT * FROM projects WHERE id = ANY($1) AND tenant_id = $2`,
        [ids as string[], tenantId]
      );
      const map = new Map(rows.map((r) => [r.id, r]));
      return ids.map((id) => map.get(id) || null);
    }),

    // Batch load tasks by project IDs (1-to-many relationship)
    tasksByProject: new DataLoader<string, any[]>(async (projectIds) => {
      const { rows } = await pool.query(
        `SELECT * FROM tasks WHERE project_id = ANY($1) AND tenant_id = $2 ORDER BY created_at DESC`,
        [projectIds as string[], tenantId]
      );
      const grouped = new Map<string, any[]>();
      for (const row of rows) {
        if (!grouped.has(row.project_id)) grouped.set(row.project_id, []);
        grouped.get(row.project_id)!.push(row);
      }
      return projectIds.map((id) => grouped.get(id) || []);
    }),

    // Batch load member counts
    memberCount: new DataLoader<string, number>(async (projectIds) => {
      const { rows } = await pool.query(
        `SELECT project_id, COUNT(*) as count FROM project_members 
         WHERE project_id = ANY($1) GROUP BY project_id`,
        [projectIds as string[]]
      );
      const map = new Map(rows.map((r) => [r.project_id, parseInt(r.count)]));
      return projectIds.map((id) => map.get(id) || 0);
    }),
  };
}

const typeDefs = `
  type Query {
    projects(first: Int, after: String, status: ProjectStatus): ProjectConnection!
    project(id: ID!): Project
    me: User!
  }

  type Mutation {
    createProject(input: CreateProjectInput!): Project!
    createTask(input: CreateTaskInput!): Task!
    updateTask(id: ID!, input: UpdateTaskInput!): Task!
  }

  type ProjectConnection {
    edges: [ProjectEdge!]!
    pageInfo: PageInfo!
    totalCount: Int!
  }

  type ProjectEdge {
    node: Project!
    cursor: String!
  }

  type PageInfo {
    hasNextPage: Boolean!
    endCursor: String
  }

  type Project {
    id: ID!
    name: String!
    description: String
    status: ProjectStatus!
    color: String!
    tasks(first: Int, status: TaskStatus): [Task!]!
    taskCount: Int!
    members: [User!]!
    memberCount: Int!
    owner: User!
    createdAt: DateTime!
    updatedAt: DateTime!
  }

  type Task {
    id: ID!
    title: String!
    description: String
    status: TaskStatus!
    priority: Priority!
    assignee: User
    project: Project!
    dueDate: DateTime
    createdAt: DateTime!
  }

  type User {
    id: ID!
    name: String!
    email: String!
    avatar: String
    role: String!
  }

  enum ProjectStatus { ACTIVE ARCHIVED }
  enum TaskStatus { TODO IN_PROGRESS DONE }
  enum Priority { LOW MEDIUM HIGH URGENT }

  input CreateProjectInput {
    name: String!
    description: String
    color: String
  }

  input CreateTaskInput {
    projectId: ID!
    title: String!
    description: String
    priority: Priority
    assigneeId: ID
    dueDate: DateTime
  }

  input UpdateTaskInput {
    title: String
    status: TaskStatus
    priority: Priority
    assigneeId: ID
  }

  scalar DateTime
`;

const resolvers = {
  Query: {
    projects: async (_: any, args: any, ctx: any) => {
      const limit = Math.min(args.first || 20, 100);
      const cursor = args.after ? Buffer.from(args.after, "base64").toString() : null;

      // Check cache
      const cacheKey = `gql:projects:${ctx.tenantId}:${args.status || "all"}:${cursor || "0"}:${limit}`;
      const cached = await redis.get(cacheKey);
      if (cached) return JSON.parse(cached);

      let query = "SELECT * FROM projects WHERE tenant_id = $1";
      const params: any[] = [ctx.tenantId];

      if (args.status) {
        params.push(args.status.toLowerCase());
        query += ` AND status = $${params.length}`;
      }
      if (cursor) {
        params.push(cursor);
        query += ` AND id > $${params.length}`;
      }

      params.push(limit + 1);
      query += ` ORDER BY id LIMIT $${params.length}`;

      const { rows } = await pool.query(query, params);
      const hasNextPage = rows.length > limit;
      const edges = rows.slice(0, limit).map((node: any) => ({
        node,
        cursor: Buffer.from(node.id).toString("base64"),
      }));

      // Total count (cached separately)
      const countCacheKey = `gql:projects:count:${ctx.tenantId}`;
      let totalCount = parseInt(await redis.get(countCacheKey) || "0");
      if (!totalCount) {
        const { rows: [{ count }] } = await pool.query(
          "SELECT COUNT(*) as count FROM projects WHERE tenant_id = $1",
          [ctx.tenantId]
        );
        totalCount = parseInt(count);
        await redis.setex(countCacheKey, 60, String(totalCount));
      }

      const result = {
        edges,
        pageInfo: {
          hasNextPage,
          endCursor: edges.length > 0 ? edges[edges.length - 1].cursor : null,
        },
        totalCount,
      };

      await redis.setex(cacheKey, 30, JSON.stringify(result));
      return result;
    },

    project: async (_: any, { id }: any, ctx: any) => {
      return ctx.loaders.project.load(id);
    },

    me: async (_: any, __: any, ctx: any) => {
      return ctx.loaders.user.load(ctx.userId);
    },
  },

  Project: {
    // DataLoader batches these — 20 projects = 1 SQL query, not 20
    tasks: (project: any, args: any, ctx: any) => ctx.loaders.tasksByProject.load(project.id),
    owner: (project: any, _: any, ctx: any) => ctx.loaders.user.load(project.owner_id),
    memberCount: (project: any, _: any, ctx: any) => ctx.loaders.memberCount.load(project.id),
    taskCount: async (project: any) => {
      const { rows: [{ count }] } = await pool.query(
        "SELECT COUNT(*) as count FROM tasks WHERE project_id = $1",
        [project.id]
      );
      return parseInt(count);
    },
  },

  Task: {
    assignee: (task: any, _: any, ctx: any) =>
      task.assignee_id ? ctx.loaders.user.load(task.assignee_id) : null,
    project: (task: any, _: any, ctx: any) => ctx.loaders.project.load(task.project_id),
  },

  Mutation: {
    createProject: async (_: any, { input }: any, ctx: any) => {
      const { rows: [project] } = await pool.query(
        `INSERT INTO projects (name, description, color, status, owner_id, tenant_id, created_at, updated_at)
         VALUES ($1, $2, $3, 'active', $4, $5, NOW(), NOW()) RETURNING *`,
        [input.name, input.description, input.color || "#3b82f6", ctx.userId, ctx.tenantId]
      );

      // Invalidate cache
      const keys = await redis.keys(`gql:projects:${ctx.tenantId}:*`);
      if (keys.length) await redis.del(...keys);

      return project;
    },

    createTask: async (_: any, { input }: any, ctx: any) => {
      const { rows: [task] } = await pool.query(
        `INSERT INTO tasks (title, description, status, priority, project_id, assignee_id, tenant_id, due_date, created_at)
         VALUES ($1, $2, 'todo', $3, $4, $5, $6, $7, NOW()) RETURNING *`,
        [input.title, input.description, input.priority || "medium",
         input.projectId, input.assigneeId, ctx.tenantId, input.dueDate]
      );
      return task;
    },
  },
};

export { typeDefs, resolvers, createLoaders };