自定义指令与模式扩展

# 类型系统指令：定义在 SDL 上
directive @auth(requires: Role = USER) on FIELD_DEFINITION | OBJECT

type Query {
  me: User! @auth
  adminStats: Stats! @auth(requires: ADMIN)
}

# 可执行指令：出现在查询里
query GetData($withOrders: Boolean!) {
  me {
    nickname
    orders @include(if: $withOrders) { edges { node { id } } }
  }
}

directive @constraint(
  minLength: Int
  maxLength: Int
  pattern: String
  format: String
) on INPUT_FIELD_DEFINITION | ARGUMENT_DEFINITION

directive @audit(level: AuditLevel = INFO) on FIELD_DEFINITION
directive @tag(name: String!) repeatable on FIELD_DEFINITION | OBJECT

import { mapSchema, getDirective, MapperKind } from '@graphql-tools/utils';
import { defaultFieldResolver } from 'graphql';

function authDirectiveTransformer(schema, directiveName = 'auth') {
  return mapSchema(schema, {
    [MapperKind.OBJECT_FIELD]: (fieldConfig) => {
      const directive = getDirective(schema, fieldConfig, directiveName)?.[0];
      if (!directive) return fieldConfig;

      const { requires = 'USER' } = directive;
      const { resolve = defaultFieldResolver } = fieldConfig;

      fieldConfig.resolve = async function (source, args, context, info) {
        const role = context.user?.role;
        if (!role || !roleSatisfies(role, requires)) {
          throw new GraphQLError('Forbidden', {
            extensions: { code: 'FORBIDDEN', requiredRole: requires },
          });
        }
        return resolve.call(this, source, args, context, info);
      };
      return fieldConfig;
    },
  });
}

let schema = buildSchema(typeDefs);
schema = authDirectiveTransformer(schema, 'auth');
schema = upperDirectiveTransformer(schema, 'upper');
schema = constraintDirectiveTransformer(schema, 'constraint');

function assertValidAuthArgs(schema, directiveName = 'auth') {
  const validRoles = new Set(['GUEST', 'USER', 'ADMIN', 'OWNER']);
  mapSchema(schema, {
    [MapperKind.OBJECT_FIELD]: (fieldConfig) => {
      const d = getDirective(schema, fieldConfig, directiveName)?.[0];
      if (d && !validRoles.has(d.requires)) {
        throw new Error(`Invalid @auth requires: ${d.requires}`);
      }
      return fieldConfig;
    },
  });
}

const ROLE_LEVEL = { GUEST: 0, USER: 1, ADMIN: 2, OWNER: 3 };

function roleSatisfies(actual, required) {
  return (ROLE_LEVEL[actual] ?? -1) >= (ROLE_LEVEL[required] ?? Infinity);
}

type Query {
  me: User!
  user(id: ID!): User @auth(requires: ADMIN)     # 字段级：仅管理员能查任意用户
  publicFeed: [Post!]! @auth(requires: GUEST)    # 类型级：允许匿名
}

type User @auth(requires: USER) {                # 类型级：整个类型需登录
  id: ID!
  email: String! @auth(requires: OWNER)          # 字段级：仅本人可见
}

orders: async (_, args, ctx) => {
  const rows = await db.orders.findMany({ where: { userId: ctx.user.id } });
  return rows; // 强制按 userId 过滤，而非依赖客户端传参
}

let schema = buildSubgraphSchema([{ typeDefs, resolvers }]);
schema = authDirectiveTransformer(schema, 'auth'); // 在联邦包装之后

directive @upper on FIELD_DEFINITION

type User {
  nickname: String! @upper
}

function upperDirectiveTransformer(schema, directiveName = 'upper') {
  return mapSchema(schema, {
    [MapperKind.OBJECT_FIELD]: (fieldConfig) => {
      const hasUpper = getDirective(schema, fieldConfig, directiveName)?.[0];
      if (!hasUpper) return fieldConfig;
      const { resolve = defaultFieldResolver } = fieldConfig;
      fieldConfig.resolve = async function (source, args, context, info) {
        const value = await resolve.call(this, source, args, context, info);
        return typeof value === 'string' ? value.toUpperCase() : value;
      };
      return fieldConfig;
    },
  });
}

directive @constraint(
  minLength: Int
  maxLength: Int
  pattern: String
  min: Int
  max: Int
) on ARGUMENT_DEFINITION | INPUT_FIELD_DEFINITION

type Mutation {
  createUser(input: CreateUserInput!): User!
}

input CreateUserInput {
  email: String! @constraint(pattern: "^[^@]+@[^@]+$")
  nickname: String! @constraint(minLength: 2, maxLength: 20)
  age: Int @constraint(min: 0, max: 150)
}

function constraintDirectiveTransformer(schema) {
  return mapSchema(schema, {
    [MapperKind.INPUT_OBJECT_FIELD]: (fieldConfig) => {
      const rules = getDirective(schema, fieldConfig, 'constraint')?.[0];
      if (rules) constraintRegistry.set(fieldConfig.astNode.name.value, rules);
      return fieldConfig;
    },
  });
}

query Profile($withEmail: Boolean!, $isMobile: Boolean!) {
  me {
    nickname
    email @include(if: $withEmail)
    avatar @skip(if: $isMobile) { url width }
    ... @defer { heavyStats { visits clicks } }
  }
}

directive @log(level: LogLevel = INFO) on FIELD

function readFieldDirectives(info) {
  const nodes = info.fieldNodes ?? [];
  return nodes.flatMap((n) => n.directives ?? []);
}

import { executeOperation } from '@apollo/server/helpers';

test('@auth 拒绝未授权访问', async () => {
  const res = await executeOperation(server, {
    query: `query { adminStats { totalUsers } }`,
    // 不传 context.user，模拟匿名
  });
  expect(res.body.singleResult.errors?.[0].extensions.code).toBe('FORBIDDEN');
});

test('@upper 转换返回值', async () => {
  const res = await executeOperation(server, { query: `query { me { nickname } }` }, {
    contextValue: { user: { id: '1', role: 'USER' } },
  });
  expect(res.body.singleResult.data.me.nickname).toBe('LEETING');
});