feat: playing around with possible shapes for custom query/mutations/subscriptions

packages/amplify-api-next/__tests__/ModelSchema.test-d.ts

@@ -1,10 +1,12 @@
 import { type ModelType, type InternalModel, model } from '../src/ModelType';
 import { type ModelField, type InternalField, fields } from '../src/ModelField';
+import { query, mutation, subscription } from '../src/CustomType';
 import {
   type ModelSchema,
   type InternalSchema,
   schema,
 } from '../src/ModelSchema';
+import { allow } from '../src/Authorization';
 
 type GetModelTypeArg<T> = T extends ModelType<infer R, any> ? R : never;
 
@@ -24,4 +26,27 @@ describe('basic functionality', () => {
     // TODO: fix
     // const is = s as InternalSchema;
   });
+
+  test('basic accepts custom stuff', () => {
+    const s = schema({
+      Post: model({
+        id: id(),
+        title: string(),
+      }),
+    }, {
+      echo: query()
+        .arguments({
+          message: string(),
+        })
+        .response(string()),
+      put: mutation()
+        .arguments({
+          id: id(),
+          content: string(),
+        })
+        .response(string()),
+      onChange: subscription()
+        .authorization([allow.public().to(['read'])]),
+    });
+  });
 });

packages/amplify-api-next/src/CustomType.ts

@@ -0,0 +1,104 @@
+import { ModelField, InternalField } from './ModelField';
+import { Authorization } from './Authorization';
+import type { SetTypeSubArg } from '@aws-amplify/amplify-api-next-types-alpha';
+
+type CustomArguments = Record<string, ModelField<any, any>>;
+
+type CustomResponse = ModelField<any, any>;
+
+type InternalCustomFields = Record<string, InternalField>;
+type InternalCustomArguments = Record<string, InternalField>;
+type InternalCustomResponse = InternalField;
+
+type CustomTypeName = 'Query' | 'Mutation' | 'Subscription';
+
+type CustomData = {
+  arguments: CustomArguments;
+  response: CustomResponse | null;
+  authorization: Authorization<any, any>[];
+  typeName: CustomTypeName;
+};
+
+type InternalCustomData = CustomData & {
+  fields: InternalCustomFields;
+  arguments: InternalCustomArguments;
+  response: InternalCustomResponse;
+  authorization: Authorization<any, any>[];
+};
+
+export type CustomTypeParamShape = {
+    arguments: CustomArguments;
+    response: CustomResponse | null;
+    authorization: Authorization<any, any>[];
+    typeName: CustomTypeName;
+};
+
+export type CustomType<
+  T extends CustomTypeParamShape,
+  K extends keyof CustomType<T> = never,
+> = Omit<
+  {
+    arguments<Arguments extends CustomArguments>(
+        args: Arguments,
+    ): CustomType<SetTypeSubArg<T, 'arguments', Arguments>, K | 'arguments'>;
+    response<Response extends CustomResponse>(
+        response: Response,
+    ): CustomType<SetTypeSubArg<T, 'response', Response>, K | 'response'>;
+    authorization<AuthRuleType extends Authorization<any, any>>(
+      rules: AuthRuleType[],
+    ): CustomType<
+      SetTypeSubArg<T, 'authorization', AuthRuleType[]>,
+      K | 'authorization'
+    >;
+  },
+  K
+>;
+
+/**
+ * Internal representation of Custom Type that exposes the `data` property.
+ * Used at buildtime.
+ */
+export type InternalCustom = CustomType<any> & {
+  data: InternalCustomData;
+};
+
+function _Custom<T extends CustomTypeParamShape>(typeName: CustomTypeName) {
+  const data: CustomData = {
+    arguments: {},
+    response: null,
+    authorization: [],
+    typeName,
+  };
+
+  const builder: CustomType<T> = {
+    arguments(args) {
+        data.arguments = args;
+
+        return this;
+    },
+    response(response) {
+        data.response = response;
+
+        return this;
+    },
+    authorization(rules) {
+      data.authorization = rules;
+
+      return this;
+    },
+  };
+
+  return { ...builder, data } as InternalCustom as CustomType<T>;
+}
+
+export function query(): CustomType<{ arguments: CustomArguments, response: null, authorization: [], typeName: 'Query' }> {
+  return _Custom('Query');
+}
+
+export function mutation(): CustomType<{ arguments: CustomArguments, response: null, authorization: [], typeName: 'Mutation' }> {
+  return _Custom('Mutation');
+}
+
+export function subscription(): CustomType<{ arguments: CustomArguments, response: null, authorization: [], typeName: 'Subscription' }> {
+  return _Custom('Subscription');
+}

packages/amplify-api-next/src/ModelSchema.ts

@@ -6,27 +6,32 @@ import type {
 } from './ModelType';
 export { __auth } from './ModelField';
 import { defineData } from './SchemaProcessor';
+import { CustomType, CustomTypeParamShape, InternalCustom } from './CustomType';
 
 /*
  * Notes:
  *
  * TSC output diagnostics to benchmark
  */
-
 type ModelSchemaModels = Record<string, ModelType<ModelTypeParamShape, any>>;
+type ModelSchemaCustom = Record<string, CustomType<CustomTypeParamShape, any>>;
 type InternalSchemaModels = Record<string, InternalModel>;
+type InternalSchemaCustom = Record<string, InternalCustom>;
 
 export type ModelSchemaParamShape = {
   models: ModelSchemaModels;
+  custom?: ModelSchemaCustom;
 };
 
 type ModelSchemaData = {
   models: ModelSchemaModels;
+  custom?: ModelSchemaCustom;
 };
 
 export type InternalSchema = {
   data: {
     models: InternalSchemaModels;
+    custom?: InternalSchemaCustom;
   };
 };
 
@@ -51,8 +56,8 @@ export const isModelSchema = (
   return typeof schema === 'object' && schema.data !== undefined;
 };
 
-function _schema<T extends ModelSchemaParamShape>(models: T['models']) {
-  const data: ModelSchemaData = { models };
+function _schema<T extends ModelSchemaParamShape>(models: T['models'], custom?: T['custom']) {
+  const data: ModelSchemaData = { models, custom };
 
   const transform = (): DerivedApiDefinition => {
     const internalSchema: InternalSchema = { data } as InternalSchema;
@@ -63,8 +68,12 @@ function _schema<T extends ModelSchemaParamShape>(models: T['models']) {
   return { data, transform } as ModelSchema<T>;
 }
 
-export function schema<Models extends ModelSchemaModels>(
+export function schema<
+  Models extends ModelSchemaModels,
+  Custom extends ModelSchemaCustom,
+>(
   models: Models,
-): ModelSchema<{ models: Models }> {
-  return _schema(models);
+  custom?: Custom,
+): ModelSchema<{ models: Models, custom: Custom }> {
+  return _schema(models, custom);
 }