From 377813f19a3c26bc8f70b65bac1956c207f41d81 Mon Sep 17 00:00:00 2001
From: Peter Vaiciulis <petesv@amazon.com>
Date: Mon, 16 Sep 2024 13:15:03 -0700
Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20support=20auto=20increme?=
 =?UTF-8?q?nt?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Implements support for auto increment (serial) fields from Postgres
datasources. Such fields are denoted by an empty `@default` applied to
an `Int` field.
---
 .../src/graphql-default-value-transformer.ts  | 26 ++++++++++++++-----
 .../__snapshots__/index.test.ts.snap          |  2 +-
 .../src/directives/default.ts                 |  2 +-
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/packages/amplify-graphql-default-value-transformer/src/graphql-default-value-transformer.ts b/packages/amplify-graphql-default-value-transformer/src/graphql-default-value-transformer.ts
index 15c76e22a5..40c373a4f2 100644
--- a/packages/amplify-graphql-default-value-transformer/src/graphql-default-value-transformer.ts
+++ b/packages/amplify-graphql-default-value-transformer/src/graphql-default-value-transformer.ts
@@ -4,6 +4,7 @@ import {
   InputObjectDefinitionWrapper,
   InvalidDirectiveError,
   isDynamoDbModel,
+  isPostgresModel,
   MappingTemplate,
   TransformerPluginBase,
 } from '@aws-amplify/graphql-transformer-core';
@@ -22,7 +23,6 @@ import {
   Kind,
   ObjectTypeDefinitionNode,
   StringValueNode,
-  TypeNode,
 } from 'graphql';
 import { methodCall, printBlock, qref, raw, ref, str } from 'graphql-mapping-template';
 import { getBaseType, isEnum, isListType, isScalarOrEnum, ModelResourceIDs, toCamelCase } from 'graphql-transformer-common';
@@ -31,16 +31,24 @@ import { TypeValidators } from './validators';
 
 const nonStringTypes = ['Int', 'Float', 'Boolean', 'AWSTimestamp', 'AWSJSON'];
 
-const validateFieldType = (ctx: TransformerSchemaVisitStepContextProvider, type: TypeNode): void => {
+const validateFieldType = (ctx: TransformerSchemaVisitStepContextProvider, config: DefaultValueDirectiveConfiguration): void => {
+  const type = config.field.type;
+  const argc = config.directive.arguments!.length;
   const enums = ctx.output.getTypeDefinitionsOfKind(Kind.ENUM_TYPE_DEFINITION) as EnumTypeDefinitionNode[];
   if (isListType(type) || !isScalarOrEnum(type, enums)) {
     throw new InvalidDirectiveError('The @default directive may only be added to scalar or enum field types.');
   }
+  if (isPostgresModel(ctx, config.object.name.value) && argc === 0 && getBaseType(type) !== 'Int') {
+    throw new InvalidDirectiveError('The empty @default (auto-increment) may only be applied to integer fields.');
+  }
 };
 
-const validateDirectiveArguments = (directive: DirectiveNode): void => {
-  if (directive.arguments!.length === 0) throw new InvalidDirectiveError('The @default directive must have a value property');
-  if (directive.arguments!.length > 1) throw new InvalidDirectiveError('The @default directive only takes a value property');
+const validateDirectiveArguments = (ctx: TransformerSchemaVisitStepContextProvider, config: DefaultValueDirectiveConfiguration): void => {
+  const argc = config.directive.arguments!.length;
+  const isPostgres = isPostgresModel(ctx, config.object.name.value);
+  if (!isPostgres && argc === 0)
+    throw new InvalidDirectiveError('The @default directive requires a value property on non Postgres datasources.');
+  if (argc > 1) throw new InvalidDirectiveError('The @default directive only takes a value property');
 };
 
 const validateModelDirective = (config: DefaultValueDirectiveConfiguration): void => {
@@ -74,8 +82,8 @@ const validateDefaultValueType = (ctx: TransformerSchemaVisitStepContextProvider
 
 const validate = (ctx: TransformerSchemaVisitStepContextProvider, config: DefaultValueDirectiveConfiguration): void => {
   validateModelDirective(config);
-  validateFieldType(ctx, config.field.type);
-  validateDirectiveArguments(config.directive);
+  validateFieldType(ctx, config);
+  validateDirectiveArguments(ctx, config);
 
   // Validate the default values only for the DynamoDB datasource.
   // For SQL, the database determines and sets the default value. We will not validate the value in transformers.
@@ -123,6 +131,10 @@ export class DefaultValueTransformer extends TransformerPluginBase {
         const input = InputObjectDefinitionWrapper.fromObject(name, config.object, ctx.inputDocument);
         const fieldWrapper = input.fields.find((f) => f.name === config.field.name.value);
         fieldWrapper?.makeNullable();
+
+        if (isPostgresModel(ctx, typeName)) {
+          ctx.output.updateInput(input.serialize());
+        }
       }
     }
   };
diff --git a/packages/amplify-graphql-directives/src/__tests__/__snapshots__/index.test.ts.snap b/packages/amplify-graphql-directives/src/__tests__/__snapshots__/index.test.ts.snap
index 8e5b36eecb..c0a20f5bdf 100644
--- a/packages/amplify-graphql-directives/src/__tests__/__snapshots__/index.test.ts.snap
+++ b/packages/amplify-graphql-directives/src/__tests__/__snapshots__/index.test.ts.snap
@@ -255,7 +255,7 @@ exports[`Directive Definitions DefaultDirective 1`] = `
 Object {
   "defaults": Object {},
   "definition": "
-  directive @default(value: String!) on FIELD_DEFINITION
+  directive @default(value: String) on FIELD_DEFINITION
 ",
   "name": "default",
 }
diff --git a/packages/amplify-graphql-directives/src/directives/default.ts b/packages/amplify-graphql-directives/src/directives/default.ts
index 4aa8c33020..c378c08073 100644
--- a/packages/amplify-graphql-directives/src/directives/default.ts
+++ b/packages/amplify-graphql-directives/src/directives/default.ts
@@ -2,7 +2,7 @@ import { Directive } from './directive';
 
 const name = 'default';
 const definition = /* GraphQL */ `
-  directive @${name}(value: String!) on FIELD_DEFINITION
+  directive @${name}(value: String) on FIELD_DEFINITION
 `;
 const defaults = {};