implement ref and getAtt for related resources (#484)

* related resources

* implement getAtt for related resources

* add parent resource selection for related resources

Author: gemammercado

src/handlers/RelatedResourcesHandler.ts

@@ -3,6 +3,7 @@ import { TopLevelSection } from '../context/CloudFormationEnums';
 import { getEntityMap } from '../context/SectionContextBuilder';
 import { Resource } from '../context/semantic/Entity';
 import {
+    AuthoredResource,
     GetRelatedResourceTypesParams,
     InsertRelatedResourcesParams,
     RelatedResourcesCodeAction,
@@ -19,22 +20,25 @@ import {
 
 export function getAuthoredResourceTypesHandler(
     components: ServerComponents,
-): RequestHandler<TemplateUri, string[], void> {
+): RequestHandler<TemplateUri, AuthoredResource[], void> {
     return (rawParams) => {
         try {
             const templateUri = parseWithPrettyError(parseTemplateUriParams, rawParams);
             const syntaxTree = components.syntaxTreeManager.getSyntaxTree(templateUri);
             if (syntaxTree) {
                 const resourcesMap = getEntityMap(syntaxTree, TopLevelSection.Resources);
                 if (resourcesMap) {
-                    const resourceTypes = [...resourcesMap.values()]
-                        .map((context) => {
-                            const resource = context.entity as Resource;
-                            return resource?.Type;
-                        })
-                        .filter((type): type is string => type !== undefined && type !== null);
-
-                    return [...new Set(resourceTypes)];
+                    const resources: AuthoredResource[] = [];
+                    for (const [logicalId, context] of resourcesMap) {
+                        const resource = context.entity as Resource;
+                        if (resource?.Type) {
+                            resources.push({
+                                logicalId,
+                                type: resource.Type,
+                            });
+                        }
+                    }
+                    return resources;
                 }
             }
 
@@ -52,26 +56,74 @@ export function getRelatedResourceTypesHandler(
         try {
             const { parentResourceType } = parseWithPrettyError(parseGetRelatedResourceTypesParams, rawParams);
             const relatedTypes = components.relationshipSchemaService.getAllRelatedResourceTypes(parentResourceType);
-            return [...relatedTypes];
+
+            const filtered = [...relatedTypes].filter((relatedType) =>
+                hasExactlyOnePopulatableRelationship(relatedType, parentResourceType, components),
+            );
+
+            return filtered;
         } catch (error) {
             handleLspError(error, 'Failed to get related resource types');
         }
     };
 }
 
+/**
+ * Checks if a related resource type has exactly one top-level property
+ * that references the parent resource type, and that property is not an array.
+ */
+function hasExactlyOnePopulatableRelationship(
+    relatedType: string,
+    parentResourceType: string,
+    components: ServerComponents,
+): boolean {
+    const relationships = components.relationshipSchemaService.getRelationshipsForResourceType(relatedType);
+    if (!relationships) {
+        return false;
+    }
+
+    const schema = components.schemaRetriever.getDefault().schemas.get(relatedType);
+
+    const topLevelParentRefs: { property: string; isArray: boolean }[] = [];
+    for (const rel of relationships.relationships) {
+        if (rel.property.includes('/')) {
+            continue;
+        }
+
+        const referencesParent = rel.relatedResourceTypes.some((rt) => rt.typeName === parentResourceType);
+        if (!referencesParent) {
+            continue;
+        }
+
+        const isArray = schema?.properties?.[rel.property]?.type === 'array';
+        topLevelParentRefs.push({ property: rel.property, isArray });
+
+        if (topLevelParentRefs.length > 1) {
+            break;
+        }
+    }
+
+    if (topLevelParentRefs.length !== 1) {
+        return false;
+    }
+
+    return !topLevelParentRefs[0].isArray;
+}
+
 export function insertRelatedResourcesHandler(
     components: ServerComponents,
 ): RequestHandler<InsertRelatedResourcesParams, RelatedResourcesCodeAction, void> {
     return (rawParams) => {
         try {
-            const { templateUri, relatedResourceTypes, parentResourceType } = parseWithPrettyError(
+            const { templateUri, relatedResourceTypes, parentResourceType, parentLogicalId } = parseWithPrettyError(
                 parseInsertRelatedResourcesParams,
                 rawParams,
             );
             return components.relatedResourcesSnippetProvider.insertRelatedResources(
                 templateUri,
                 relatedResourceTypes,
                 parentResourceType,
+                parentLogicalId,
             );
         } catch (error) {
             handleLspError(error, 'Failed to insert related resources');

src/handlers/RelatedResourcesParser.ts

@@ -16,6 +16,7 @@ const InsertRelatedResourcesParamsSchema = z.object({
     templateUri: NonEmptyZodString,
     relatedResourceTypes: z.array(NonEmptyZodString).min(1),
     parentResourceType: NonEmptyZodString,
+    parentLogicalId: NonEmptyZodString.optional(),
 });
 
 export function parseTemplateUriParams(input: unknown): TemplateUri {

src/protocol/LspRelatedResourcesHandlers.ts

@@ -1,5 +1,6 @@
 import { Connection, RequestHandler } from 'vscode-languageserver';
 import {
+    AuthoredResource,
     GetAuthoredResourceTypesRequest,
     GetRelatedResourceTypesParams,
     GetRelatedResourceTypesRequest,
@@ -12,7 +13,7 @@ import {
 export class LspRelatedResourcesHandlers {
     constructor(private readonly connection: Connection) {}
 
-    onGetAuthoredResourceTypes(handler: RequestHandler<TemplateUri, string[], void>) {
+    onGetAuthoredResourceTypes(handler: RequestHandler<TemplateUri, AuthoredResource[], void>) {
         this.connection.onRequest(GetAuthoredResourceTypesRequest.method, handler);
     }
 

src/protocol/RelatedResourcesProtocol.ts

@@ -6,10 +6,16 @@ export type GetRelatedResourceTypesParams = {
     parentResourceType: string;
 };
 
+export type AuthoredResource = {
+    logicalId: string;
+    type: string;
+};
+
 export type InsertRelatedResourcesParams = {
     templateUri: string;
     relatedResourceTypes: string[];
     parentResourceType: string;
+    parentLogicalId?: string;
 };
 
 export interface RelatedResourcesCodeAction extends CodeAction {
@@ -19,7 +25,7 @@ export interface RelatedResourcesCodeAction extends CodeAction {
     };
 }
 
-export const GetAuthoredResourceTypesRequest = new RequestType<TemplateUri, string[], void>(
+export const GetAuthoredResourceTypesRequest = new RequestType<TemplateUri, AuthoredResource[], void>(
     'aws/cfn/template/resources/authored',
 );