feat: Implement multiple AWS authentication methods

apps/dbagent/package.json

@@ -29,6 +29,7 @@
     "@anthropic-ai/sdk": "^0.51.0",
     "@aws-sdk/client-cloudwatch": "^3.812.0",
     "@aws-sdk/client-rds": "^3.812.0",
+    "@aws-sdk/client-sts": "^3.812.0",
     "@fluentui/react-icons": "^2.0.300",
     "@google-cloud/logging": "^11.2.0",
     "@google-cloud/monitoring": "^5.1.0",

apps/dbagent/public/xata-agent-iam-role.yaml

@@ -0,0 +1,72 @@
+AWSTemplateFormatVersion: '2010-09-09'
+Description: >-
+  Creates an IAM Role for the Xata Agent to access AWS RDS cluster information,
+  instance details, logs, and CloudWatch metrics. You need to provide the
+  AWS Account ID where your Xata Agent or the users/roles that will assume this
+  role reside.
+
+Parameters:
+  TrustedAwsAccountId:
+    Type: AWS::AccountId
+    Description: The AWS Account ID that will be trusted to assume this role (e.g., the account where your application or Xata Agent runs).
+  RoleName:
+    Type: String
+    Default: XataAgentRDSAccessRole
+    Description: Name for the IAM role. Keep the default unless you have specific naming conventions.
+
+Resources:
+  XataAgentRole:
+    Type: AWS::IAM::Role
+    Properties:
+      RoleName: !Ref RoleName
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              AWS: !Sub "arn:aws:iam::${TrustedAwsAccountId}:root" # Trusts the root of the specified AWS Account ID
+            Action:
+              - sts:AssumeRole
+      Path: "/" # Default path
+      Policies:
+        - PolicyName: XataAgentRDSAccessPolicy
+          PolicyDocument:
+            Version: '2012-10-17'
+            Statement:
+              - Effect: Allow
+                Action:
+                  - rds:DescribeDBClusters
+                  - rds:DescribeDBInstances
+                  - rds:DescribeDBLogFiles
+                  - rds:DownloadDBLogFilePortion
+                  - cloudwatch:GetMetricStatistics
+                  # The following permission is for future use, e.g., detecting if the agent is on an EC2 instance.
+                  # It can be kept commented if not immediately required by your setup.
+                  # - ec2:DescribeInstances
+                Resource: "*" # For simplicity, this policy allows access to all RDS and CloudWatch resources.
+                               # For a more secure setup, you can restrict this to specific resources if known.
+
+Outputs:
+  XataAgentRoleArn:
+    Description: ARN of the created IAM role for Xata Agent. Use this ARN in the Xata Agent UI.
+    Value: !GetAtt XataAgentRole.Arn
+    Export:
+      Name: XataAgentRoleArn # Makes it easier to reference this role ARN in other CloudFormation stacks.
+
+Metadata:
+  AWS::CloudFormation::Interface:
+    ParameterGroups:
+      - Label:
+          default: "Role Configuration"
+        Parameters:
+          - RoleName
+      - Label:
+          default: "Trusted Principal Configuration"
+        Parameters:
+          - TrustedAwsAccountId
+    ParameterLabels:
+      TrustedAwsAccountId:
+        default: "Trusted AWS Account ID"
+      RoleName:
+        default: "IAM Role Name"
+```

apps/dbagent/src/components/aws-integration/actions.ts

@@ -4,6 +4,7 @@ import {
   getRDSClusterInfo,
   getRDSInstanceInfo,
   initializeRDSClient,
+  isEc2InstanceWithRole,
   listRDSClusters,
   listRDSInstances,
   RDSClusterDetailedInfo,
@@ -16,16 +17,13 @@ import { Connection } from '~/lib/db/schema';
 
 export async function fetchRDSClusters(
   projectId: string,
-  accessKeyId: string,
-  secretAccessKey: string,
-  region: string
+  integration: AwsIntegration
 ): Promise<{ success: boolean; message: string; data: RDSClusterInfo[] }> {
-  const client = initializeRDSClient({ accessKeyId, secretAccessKey, region });
-
   try {
+    const client = await initializeRDSClient(integration);
     const clusters = await listRDSClusters(client);
     const instances = await listRDSInstances(client);
-    // Add standalone instances as "clusters" with single instance
+
     const standaloneInstances = instances.filter((instance) => !instance.dbClusterIdentifier);
     const standaloneAsClusters: RDSClusterInfo[] = standaloneInstances.map((instance) => ({
       identifier: instance.identifier,
@@ -42,45 +40,45 @@ export async function fetchRDSClusters(
     clusters.push(...standaloneAsClusters);
 
     const dbAccess = await getUserSessionDBAccess();
-    await saveIntegration(dbAccess, projectId, 'aws', { accessKeyId, secretAccessKey, region });
-    return { success: true, message: 'RDS instances fetched successfully', data: clusters };
+
+    await saveIntegration(dbAccess, projectId, 'aws', integration);
+    return { success: true, message: 'RDS clusters/instances fetched successfully', data: clusters };
   } catch (error) {
-    console.error('Error fetching RDS instances:', error);
-    return { success: false, message: 'Error fetching RDS instances', data: [] };
+    console.error('Error fetching RDS clusters/instances:', error);
+    const errorMessage = error instanceof Error ? error.message : 'An unknown error occurred';
+    return { success: false, message: `Error fetching RDS clusters/instances: ${errorMessage}`, data: [] };
   }
 }
 
 export async function fetchRDSClusterDetails(
   projectId: string,
-  clusterInfo: RDSClusterInfo
+  clusterInfo: RDSClusterInfo,
+  integration: AwsIntegration
 ): Promise<{ success: boolean; message: string; data: RDSClusterDetailedInfo | null }> {
-  const dbAccess = await getUserSessionDBAccess();
-  const aws = await getIntegration(dbAccess, projectId, 'aws');
-  if (!aws) {
-    return { success: false, message: 'AWS integration not found', data: null };
-  }
-  const client = initializeRDSClient({
-    accessKeyId: aws.accessKeyId,
-    secretAccessKey: aws.secretAccessKey,
-    region: aws.region
-  });
+  try {
+    const client = await initializeRDSClient(integration);
 
-  if (clusterInfo.isStandaloneInstance) {
-    const instance = await getRDSInstanceInfo(clusterInfo.identifier, client);
-    if (!instance) {
-      return { success: false, message: 'RDS instance not found', data: null };
-    }
-    const cluster = {
-      ...clusterInfo,
-      instances: [instance]
-    };
-    return { success: true, message: 'RDS instance details fetched successfully', data: cluster };
-  } else {
-    const cluster = await getRDSClusterInfo(clusterInfo.identifier, client);
-    if (!cluster) {
-      return { success: false, message: 'RDS cluster not found', data: null };
+    if (clusterInfo.isStandaloneInstance) {
+      const instance = await getRDSInstanceInfo(clusterInfo.identifier, client);
+      if (!instance) {
+        return { success: false, message: 'RDS instance not found', data: null };
+      }
+      const cluster = {
+        ...clusterInfo,
+        instances: [instance]
+      };
+      return { success: true, message: 'RDS instance details fetched successfully', data: cluster };
+    } else {
+      const cluster = await getRDSClusterInfo(clusterInfo.identifier, client);
+      if (!cluster) {
+        return { success: false, message: 'RDS cluster not found', data: null };
+      }
+      return { success: true, message: 'RDS cluster details fetched successfully', data: cluster };
     }
-    return { success: true, message: 'RDS cluster details fetched successfully', data: cluster };
+  } catch (error) {
+    console.error('Error fetching RDS cluster/instance details:', error);
+    const errorMessage = error instanceof Error ? error.message : 'An unknown error occurred';
+    return { success: false, message: `Error fetching details: ${errorMessage}`, data: null };
   }
 }
 
@@ -93,70 +91,109 @@ export async function getAWSIntegration(
     if (!aws) {
       return { success: false, message: 'AWS integration not found', data: null };
     }
-    return { success: true, message: 'AWS integration found', data: aws };
+    // Ensure the returned data conforms to AwsIntegration, especially if old data might exist
+    // For example, if 'authMethod' was missing in older records, default it to 'credentials'
+    const validatedAws: AwsIntegration = {
+      authMethod: aws.authMethod || 'credentials',
+      region: aws.region,
+      ...(aws.authMethod === 'credentials' && {
+        accessKeyId: aws.accessKeyId,
+        secretAccessKey: aws.secretAccessKey
+      }),
+      ...(aws.authMethod === 'cloudformation' && {
+        cloudformationStackArn: aws.cloudformationStackArn
+      })
+    } as AwsIntegration;
+
+    return { success: true, message: 'AWS integration found', data: validatedAws };
   } catch (error) {
     console.error('Error fetching AWS integration:', error);
-    return { success: false, message: 'Error fetching AWS integration', data: null };
+    const errorMessage = error instanceof Error ? error.message : 'An unknown error occurred';
+    return { success: false, message: `Error fetching AWS integration: ${errorMessage}`, data: null };
+  }
+}
+
+export async function checkEc2InstanceRoleStatus(): Promise<{
+  success: boolean;
+  message: string;
+  data: { hasIAMRole: boolean } | null;
+}> {
+  try {
+    const hasIAMRole = await isEc2InstanceWithRole();
+    return {
+      success: true,
+      message: 'Successfully checked EC2 instance IAM role status.',
+      data: { hasIAMRole }
+    };
+  } catch (error) {
+    console.error('Error checking EC2 instance IAM role status:', error);
+    const errorMessage = error instanceof Error ? error.message : 'An unknown error occurred';
+    return { success: false, message: `Error checking EC2 status: ${errorMessage}`, data: null };
   }
 }
 
 export async function saveClusterDetails(
   clusterIdentifier: string,
-  region: string,
   connection: Connection
 ): Promise<{ success: boolean; message: string }> {
   const dbAccess = await getUserSessionDBAccess();
-  const aws = await getIntegration(dbAccess, connection.projectId, 'aws');
-  if (!aws) {
-    return { success: false, message: 'AWS integration not found' };
-  }
-  const client = initializeRDSClient({
-    accessKeyId: aws.accessKeyId,
-    secretAccessKey: aws.secretAccessKey,
-    region: region
-  });
-  const cluster = await getRDSClusterInfo(clusterIdentifier, client);
-  if (cluster) {
-    const instanceId = await saveCluster(dbAccess, {
-      projectId: connection.projectId,
-      clusterIdentifier,
-      region,
-      data: cluster
-    });
-    await associateClusterConnection(dbAccess, {
-      projectId: connection.projectId,
-      clusterId: instanceId,
-      connectionId: connection.id
-    });
-    return { success: true, message: 'Cluster details saved successfully' };
-  } else {
-    const instance = await getRDSInstanceInfo(clusterIdentifier, client);
-    if (!instance) {
-      return { success: false, message: 'RDS instance not found' };
+  try {
+    const integration = await getAWSIntegration(connection.projectId);
+    if (!integration.success || !integration.data) {
+      return { success: false, message: 'AWS integration not found' };
     }
-    const instanceId = await saveCluster(dbAccess, {
-      projectId: connection.projectId,
-      clusterIdentifier,
-      region,
-      data: {
-        instances: [instance],
-        identifier: instance.identifier,
-        engine: instance.engine,
-        engineVersion: instance.engineVersion,
-        status: instance.status,
-        endpoint: instance.endpoint?.address,
-        port: instance.endpoint?.port,
-        multiAZ: instance.multiAZ,
-        instanceCount: 1,
-        allocatedStorage: instance.allocatedStorage,
-        isStandaloneInstance: true
+
+    const client = await initializeRDSClient(integration.data);
+
+    const cluster = await getRDSClusterInfo(clusterIdentifier, client);
+    if (cluster) {
+      const clusterId = await saveCluster(dbAccess, {
+        projectId: connection.projectId,
+        clusterIdentifier,
+        region: integration.data.region,
+        data: cluster
+      });
+      await associateClusterConnection(dbAccess, {
+        projectId: connection.projectId,
+        clusterId,
+        connectionId: connection.id
+      });
+      return { success: true, message: 'Cluster details saved successfully' };
+    } else {
+      const instanceInfo = await getRDSInstanceInfo(clusterIdentifier, client);
+      if (instanceInfo) {
+        const standaloneAsClusterData: RDSClusterDetailedInfo = {
+          identifier: instanceInfo.identifier,
+          engine: instanceInfo.engine,
+          engineVersion: instanceInfo.engineVersion,
+          status: instanceInfo.status,
+          endpoint: instanceInfo.endpoint?.address,
+          port: instanceInfo.endpoint?.port,
+          multiAZ: instanceInfo.multiAZ,
+          instanceCount: 1,
+          allocatedStorage: instanceInfo.allocatedStorage,
+          isStandaloneInstance: true,
+          instances: [instanceInfo]
+        };
+        const instanceDbId = await saveCluster(dbAccess, {
+          projectId: connection.projectId,
+          clusterIdentifier,
+          region: integration.data.region,
+          data: standaloneAsClusterData
+        });
+        await associateClusterConnection(dbAccess, {
+          projectId: connection.projectId,
+          clusterId: instanceDbId,
+          connectionId: connection.id
+        });
+        return { success: true, message: 'Instance details saved successfully' };
+      } else {
+        return { success: false, message: 'RDS cluster or instance not found' };
       }
-    });
-    await associateClusterConnection(dbAccess, {
-      projectId: connection.projectId,
-      clusterId: instanceId,
-      connectionId: connection.id
-    });
-    return { success: true, message: 'Instance details saved successfully' };
+    }
+  } catch (error) {
+    console.error('Error saving cluster/instance details:', error);
+    const errorMessage = error instanceof Error ? error.message : 'An unknown error occurred';
+    return { success: false, message: `Error saving details: ${errorMessage}` };
   }
 }