refactor: move getConnectedFieldForReferences to separate function

aws-amplify · Apr 25, 2024 · 20f83b1 · 20f83b1
1 parent 14ee319
commit 20f83b1
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 55 deletions.
diff --git a/packages/appsync-modelgen-plugin/src/utils/process-belongs-to.ts b/packages/appsync-modelgen-plugin/src/utils/process-belongs-to.ts
@@ -6,7 +6,7 @@ import {
   flattenFieldDirectives,
   makeConnectionAttributeName,
 } from './process-connections';
-import { getConnectedFieldV2, fieldsAndReferencesErrorMessage } from './process-connections-v2';
+import { getConnectedFieldV2 } from './process-connections-v2';
 
 
 export function processBelongsToConnection(
@@ -28,20 +28,21 @@ export function processBelongsToConnection(
       `A 'belongsTo' field should match to a corresponding 'hasMany' or 'hasOne' field`
     );
   }
+  const references = connectionDirective.arguments.references || [];
+  const isUsingReferences = references.length > 0;
+  if (isUsingReferences) {
+    // ensure there is a matching hasOne/hasMany field with references
+    getConnectedFieldV2(field, model, otherSide, connectionDirective.name)
+  }
+
   const otherSideField = isCustomPKEnabled ? otherSideConnectedFields[0] : getConnectedFieldV2(field, model, otherSide, connectionDirective.name);
   const connectionFields = connectionDirective.arguments.fields || [];
 
-  const references = connectionDirective.arguments.references || [];
-
-  if (connectionFields.length > 0 && references.length > 0) {
-    throw new Error(fieldsAndReferencesErrorMessage);
-  }
   // if a type is connected using name, then amplify-graphql-relational-transformer adds a field to
   //  track the connection and that field is not part of the selection set
   // but if the field are connected using fields argument in connection directive
   // we are reusing the field and it should be preserved in selection set
   const otherSideHasMany = otherSideField.isList;
-  const isUsingReferences = references.length > 0;
   // New metada type introduced by custom PK v2 support
   let targetNames: string[] = [ ...connectionFields, ...references ];
   if (targetNames.length === 0) {

diff --git a/packages/appsync-modelgen-plugin/src/utils/process-connections-v2.ts b/packages/appsync-modelgen-plugin/src/utils/process-connections-v2.ts
@@ -19,23 +19,8 @@ export function getConnectedFieldV2(
   if (!connectionInfo) {
     throw new Error(`The ${field.name} on model ${model.name} is not connected`);
   }
-
-  const references = connectionInfo.arguments.references;
   if (connectionInfo.name === 'belongsTo') {
     let connectedFieldsBelongsTo = getBelongsToConnectedFields(model, connectedModel);
-    if (references) {
-      const connectedField = connectedFieldsBelongsTo.find((field) => {
-        return field.directives.some((dir) => {
-          return (dir.name === 'hasOne' || dir.name === 'hasMany')
-            && dir.arguments.references
-            && JSON.stringify(dir.arguments.references) === JSON.stringify(connectionInfo.arguments.references);
-        });
-      });
-      if (!connectedField) {
-       throw new Error(`Error processing @belongsTo directive on ${model.name}.${field.name}. @hasOne or @hasMany directive with references ${JSON.stringify(connectionInfo.arguments?.references)} was not found in connected model ${connectedModel.name}`);
-      }
-      return connectedField;
-    }
 
     if (connectedFieldsBelongsTo.length === 1) {
       return connectedFieldsBelongsTo[0];
@@ -44,25 +29,10 @@ export function getConnectedFieldV2(
 
   const indexName = connectionInfo.arguments.indexName;
   const connectionFields = connectionInfo.arguments.fields;
-  if (connectionFields && references) {
-    throw new Error(fieldsAndReferencesErrorMessage);
-  }
-  if (references || connectionFields || directiveName === 'hasOne') {
+  if (connectionFields || directiveName === 'hasOne') {
     let connectionDirective;
-    if (references) {
-      if (connectionInfo) {
-        connectionDirective = flattenFieldDirectives(connectedModel).find((dir) => {
-          return dir.arguments.references
-            && JSON.stringify(dir.arguments.references) === JSON.stringify(connectionInfo.arguments.references);
-        });
-        if (!connectionDirective) {
-         throw new Error(`Error processing @${connectionInfo.name} directive on ${model.name}.${field.name}. @belongsTo directive with references ${JSON.stringify(connectionInfo.arguments?.references)} was not found in connected model ${connectedModel.name}`);
-        }
-      }
-    }
-
     // Find gsi on other side if index is defined
-    else if (indexName) {
+    if (indexName) {
       connectionDirective = flattenFieldDirectives(connectedModel).find(dir => {
         return dir.name === 'index' && dir.arguments.name === indexName;
       });
@@ -154,6 +124,56 @@ export function getConnectedFieldV2(
       };
 }
 
+export function getConnectedFieldForReferences(
+  field: CodeGenField,
+  model: CodeGenModel,
+  connectedModel: CodeGenModel,
+  directiveName: string,
+): CodeGenField {
+  const connectionInfo = getDirective(field)(directiveName);
+  if (!connectionInfo) {
+    throw new Error(`The ${field.name} on model ${model.name} is not connected`);
+  }
+  const references = connectionInfo.arguments.references;
+  if (!references) {
+    throw new Error(`The ${field.name} on model ${model.name} does not have references.`);
+  }
+  const connectionFields = connectionInfo.arguments.fields;
+  if (connectionFields && references) {
+    throw new Error(`'fields' and 'references' cannot be used together.`);
+  }
+
+  if (connectionInfo.name === 'belongsTo') {
+    let connectedFieldsBelongsTo = getBelongsToConnectedFields(model, connectedModel);
+    const connectedField = connectedFieldsBelongsTo.find((field) => {
+      return field.directives.some((dir) => {
+        return (dir.name === 'hasOne' || dir.name === 'hasMany')
+          && dir.arguments.references
+          && JSON.stringify(dir.arguments.references) === JSON.stringify(connectionInfo.arguments.references);
+      });
+    });
+    if (!connectedField) {
+     throw new Error(`Error processing @belongsTo directive on ${model.name}.${field.name}. @hasOne or @hasMany directive with references ${JSON.stringify(connectionInfo.arguments?.references)} was not found in connected model ${connectedModel.name}`);
+    }
+    return connectedField;
+  }
+
+  // hasOne and hasMany
+  const connectionDirective = flattenFieldDirectives(connectedModel).find((dir) => {
+    return dir.arguments.references
+      && JSON.stringify(dir.arguments.references) === JSON.stringify(connectionInfo.arguments.references);
+  });
+  if (!connectionDirective) {
+   throw new Error(`Error processing @${connectionInfo.name} directive on ${model.name}.${field.name}. @belongsTo directive with references ${JSON.stringify(connectionInfo.arguments?.references)} was not found in connected model ${connectedModel.name}`);
+  }
+  const connectedFieldName = ((fieldDir: CodeGenFieldDirective) => {
+    return fieldDir.fieldName;
+  })(connectionDirective as CodeGenFieldDirective)
+
+  const connectedField = connectedModel.fields.find(f => f.name === connectedFieldName);
+  return connectedField!;
+}
+
 export function processConnectionsV2(
   field: CodeGenField,
   model: CodeGenModel,
@@ -177,5 +197,3 @@ export function processConnectionsV2(
     }
   }
 }
-
-export const fieldsAndReferencesErrorMessage = `'fields' and 'references' cannot be used together.`;
diff --git a/packages/appsync-modelgen-plugin/src/utils/process-has-many.ts b/packages/appsync-modelgen-plugin/src/utils/process-has-many.ts
@@ -9,7 +9,7 @@ import {
   flattenFieldDirectives,
   CodeGenFieldConnectionHasMany,
 } from './process-connections';
-import { getConnectedFieldV2, fieldsAndReferencesErrorMessage } from './process-connections-v2';
+import { getConnectedFieldV2, getConnectedFieldForReferences } from './process-connections-v2';
 
 
 export function processHasManyConnection(
@@ -27,15 +27,11 @@ export function processHasManyConnection(
   const connectionFields = connectionDirective.arguments.fields || [];
   const references = connectionDirective.arguments.references || [];
 
-  if (connectionFields.length > 0 && references.length > 0) {
-    throw new Error(fieldsAndReferencesErrorMessage);
-  }
-
   if (references.length > 0) {
     // native uses the connected field instead of associatedWithFields
     // when using references associatedWithFields and associatedWithNative are not the same
-    // getConnectedFieldV2 also ensures there is a matching belongsTo field with references
-    const associatedWithNativeReferences = getConnectedFieldV2(field, model, otherSide, connectionDirective.name, shouldUseModelNameFieldInHasManyAndBelongsTo)
+    // getConnectedFieldForRerences also ensures there is a matching belongsTo field with references
+    const associatedWithNativeReferences = getConnectedFieldForReferences(field, model, otherSide, connectionDirective.name)
     const associatedWithFields = references.map((reference: string) => otherSide.fields.find((field) => reference === field.name))
     return {
       kind: CodeGenConnectionType.HAS_MANY,

diff --git a/packages/appsync-modelgen-plugin/src/utils/process-has-one.ts b/packages/appsync-modelgen-plugin/src/utils/process-has-one.ts
@@ -4,7 +4,7 @@ import {
   CodeGenFieldConnection,
   makeConnectionAttributeName,
 } from './process-connections';
-import { getConnectedFieldV2, fieldsAndReferencesErrorMessage } from './process-connections-v2';
+import { getConnectedFieldV2, getConnectedFieldForReferences } from './process-connections-v2';
 import { getModelPrimaryKeyComponentFields } from './fieldUtils';
 import { getOtherSideBelongsToField } from './fieldUtils';
 
@@ -26,16 +26,12 @@ export function processHasOneConnection(
   const connectionFields = connectionDirective.arguments.fields || [];
   const references = connectionDirective.arguments.references || [];
 
-  if (connectionFields.length > 0 && references.length > 0) {
-    throw new Error(fieldsAndReferencesErrorMessage);
-  }
-
   let associatedWithFields;
   if (references.length > 0) {
     // native uses the connected field instead of associatedWithFields
     // when using references associatedWithFields and associatedWithNative are not the same
-    // getConnectedFieldV2 also ensures there is a matching belongsTo field with references
-    const associatedWithNativeReferences = getConnectedFieldV2(field, model, otherSide, connectionDirective.name);
+    // getConnectedFieldForReferences also ensures there is a matching belongsTo field with references
+    const associatedWithNativeReferences = getConnectedFieldForReferences(field, model, otherSide, connectionDirective.name)
     associatedWithFields = references.map((reference: string) => otherSide.fields.find((field) => reference === field.name))
     return {
       kind: CodeGenConnectionType.HAS_ONE,