Cloudwatch: Add support for adding account id to sql query (#90880)

Author: idastambuk

public/app/plugins/datasource/cloudwatch/components/QueryEditor/MetricsQueryEditor/SQLBuilderEditor/SQLBuilderEditor.tsx

@@ -26,7 +26,7 @@ export const SQLBuilderEditor = ({ query, datasource, onChange }: React.PropsWit
   const onQueryChange = useCallback(
     (query: CloudWatchMetricsQuery) => {
       const sqlGenerator = new SQLGenerator();
-      const sqlString = sqlGenerator.expressionToSqlQuery(query.sql ?? {});
+      const sqlString = sqlGenerator.expressionToSqlQuery(query.sql ?? {}, query.accountId);
       const fullQuery = {
         ...query,
         sqlExpression: sqlString,
@@ -40,7 +40,7 @@ export const SQLBuilderEditor = ({ query, datasource, onChange }: React.PropsWit
   const [sqlPreview, setSQLPreview] = useState<string | undefined>();
   useEffect(() => {
     const sqlGenerator = new SQLGenerator();
-    const sqlString = sqlGenerator.expressionToSqlQuery(query.sql ?? {});
+    const sqlString = sqlGenerator.expressionToSqlQuery(query.sql ?? {}, query.accountId);
     if (sqlPreview !== sqlString) {
       setSQLPreview(sqlString);
     }

public/app/plugins/datasource/cloudwatch/language/cloudwatch-sql/SQLGenerator.test.ts

@@ -110,12 +110,30 @@ describe('SQLGenerator', () => {
     });
   });
 
-  function assertQueryEndsWith(rest: Partial<SQLExpression>, expectedFilter: string) {
-    expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery, ...rest })).toEqual(
-      `SELECT SUM(CPUUtilization) FROM SCHEMA("AWS/EC2") ${expectedFilter}`
-    );
+  function assertQueryEndsWith(args: { sql: Partial<SQLExpression>; accountId?: string }, expectedFilter: string) {
+    expect(
+      new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery, ...args.sql }, args.accountId)
+    ).toEqual(`SELECT SUM(CPUUtilization) FROM SCHEMA("AWS/EC2") ${expectedFilter}`);
   }
+  describe('accountId', () => {
+    it('should add where clause if account ID is defined', () => {
+      expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery }, '12345')).toEqual(
+        'SELECT SUM(CPUUtilization) FROM SCHEMA("AWS/EC2") WHERE AWS.AccountId = \'12345\''
+      );
+    });
+
+    it('should not add where clause if account ID is not defined', () => {
+      expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery })).toEqual(
+        'SELECT SUM(CPUUtilization) FROM SCHEMA("AWS/EC2")'
+      );
+    });
 
+    it('should not add where clause if account ID is all', () => {
+      expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery })).toEqual(
+        'SELECT SUM(CPUUtilization) FROM SCHEMA("AWS/EC2")'
+      );
+    });
+  });
   describe('filter', () => {
     it('should not add WHERE clause in case its empty', () => {
       expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery })).not.toContain('WHERE');
@@ -126,6 +144,21 @@ describe('SQLGenerator', () => {
       expect(new SQLGenerator(mockTemplateSrv).expressionToSqlQuery({ ...baseQuery, where })).not.toContain('WHERE');
     });
 
+    it('should add where clauses with AND if accountID is defined', () => {
+      const where = createArray([createOperator('Instance-Id', '=', 'I-123')]);
+      assertQueryEndsWith(
+        { sql: { where }, accountId: '12345' },
+        `WHERE AWS.AccountId = '12345' AND "Instance-Id" = 'I-123'`
+      );
+    });
+    it('should add where clauses with WHERE if accountID is not defined', () => {
+      const where = createArray([createOperator('Instance-Id', '=', 'I-123')]);
+      assertQueryEndsWith({ sql: { where } }, `WHERE "Instance-Id" = 'I-123'`);
+    });
+    it('should add where clauses with WHERE if accountID is all', () => {
+      const where = createArray([createOperator('Instance-Id', '=', 'I-123')]);
+      assertQueryEndsWith({ sql: { where }, accountId: 'all' }, `WHERE "Instance-Id" = 'I-123'`);
+    });
     // TODO: We should handle this scenario
     it.skip('should not add WHERE clause when the operator is incomplete', () => {
       const where = createArray([createOperator('Instance-Id', '=')]);
@@ -134,12 +167,12 @@ describe('SQLGenerator', () => {
 
     it('should handle one top level filter with AND', () => {
       const where = createArray([createOperator('Instance-Id', '=', 'I-123')]);
-      assertQueryEndsWith({ where }, `WHERE "Instance-Id" = 'I-123'`);
+      assertQueryEndsWith({ sql: { where } }, `WHERE "Instance-Id" = 'I-123'`);
     });
 
     it('should handle one top level filter with OR', () => {
       assertQueryEndsWith(
-        { where: createArray([createOperator('InstanceId', '=', 'I-123')]) },
+        { sql: { where: createArray([createOperator('InstanceId', '=', 'I-123')]) } },
         `WHERE InstanceId = 'I-123'`
       );
     });
@@ -149,15 +182,15 @@ describe('SQLGenerator', () => {
         [createOperator('InstanceId', '=', 'I-123'), createOperator('Instance-Id', '!=', 'I-456')],
         QueryEditorExpressionType.And
       );
-      assertQueryEndsWith({ where: filter }, `WHERE InstanceId = 'I-123' AND "Instance-Id" != 'I-456'`);
+      assertQueryEndsWith({ sql: { where: filter } }, `WHERE InstanceId = 'I-123' AND "Instance-Id" != 'I-456'`);
     });
 
     it('should handle multiple top level filters combined with OR', () => {
       const filter = createArray(
         [createOperator('InstanceId', '=', 'I-123'), createOperator('InstanceId', '!=', 'I-456')],
         QueryEditorExpressionType.Or
       );
-      assertQueryEndsWith({ where: filter }, `WHERE InstanceId = 'I-123' OR InstanceId != 'I-456'`);
+      assertQueryEndsWith({ sql: { where: filter } }, `WHERE InstanceId = 'I-123' OR InstanceId != 'I-456'`);
     });
 
     it('should handle one top level filters with one nested filter', () => {
@@ -168,7 +201,7 @@ describe('SQLGenerator', () => {
         ],
         QueryEditorExpressionType.And
       );
-      assertQueryEndsWith({ where: filter }, `WHERE InstanceId = 'I-123' AND InstanceId != 'I-456'`);
+      assertQueryEndsWith({ sql: { where: filter } }, `WHERE InstanceId = 'I-123' AND InstanceId != 'I-456'`);
     });
 
     it('should handle one top level filter with two nested filters combined with AND', () => {
@@ -184,7 +217,7 @@ describe('SQLGenerator', () => {
       );
       // In this scenario, the parenthesis are redundant. However, they're not doing any harm and it would be really complicated to remove them
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE "Instance.Type" = 'I-123' AND (InstanceId != 'I-456' AND Type != 'some-type')`
       );
     });
@@ -201,7 +234,7 @@ describe('SQLGenerator', () => {
         QueryEditorExpressionType.And
       );
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE InstanceId = 'I-123' AND (InstanceId != 'I-456' OR Type != 'some-type')`
       );
     });
@@ -222,7 +255,7 @@ describe('SQLGenerator', () => {
       );
 
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE (InstanceId = 'I-123' AND Type != 'some-type') AND (InstanceId != 'I-456' OR Type != 'some-type')`
       );
     });
@@ -242,7 +275,7 @@ describe('SQLGenerator', () => {
         QueryEditorExpressionType.Or
       );
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE (InstanceId = 'I-123' OR Type != 'some-type') OR (InstanceId != 'I-456' OR Type != 'some-type')`
       );
     });
@@ -257,7 +290,7 @@ describe('SQLGenerator', () => {
         QueryEditorExpressionType.Or
       );
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE InstanceId = 'I-123' OR Type != 'some-type' OR InstanceId != 'I-456'`
       );
     });
@@ -272,7 +305,7 @@ describe('SQLGenerator', () => {
         QueryEditorExpressionType.And
       );
       assertQueryEndsWith(
-        { where: filter },
+        { sql: { where: filter } },
         `WHERE InstanceId = 'I-123' AND Type != 'some-type' AND InstanceId != 'I-456'`
       );
     });
@@ -284,14 +317,14 @@ describe('SQLGenerator', () => {
     });
     it('should handle single label', () => {
       const groupBy = createArray([createGroupBy('InstanceId')], QueryEditorExpressionType.And);
-      assertQueryEndsWith({ groupBy }, `GROUP BY InstanceId`);
+      assertQueryEndsWith({ sql: { groupBy } }, `GROUP BY InstanceId`);
     });
     it('should handle multiple label', () => {
       const groupBy = createArray(
         [createGroupBy('InstanceId'), createGroupBy('Type'), createGroupBy('Group')],
         QueryEditorExpressionType.And
       );
-      assertQueryEndsWith({ groupBy }, `GROUP BY InstanceId, Type, Group`);
+      assertQueryEndsWith({ sql: { groupBy } }, `GROUP BY InstanceId, Type, Group`);
     });
   });
 
@@ -301,16 +334,16 @@ describe('SQLGenerator', () => {
     });
     it('should handle SUM ASC', () => {
       const orderBy = createFunction('SUM');
-      assertQueryEndsWith({ orderBy, orderByDirection: 'ASC' }, `ORDER BY SUM() ASC`);
+      assertQueryEndsWith({ sql: { orderBy, orderByDirection: 'ASC' } }, `ORDER BY SUM() ASC`);
     });
 
     it('should handle SUM ASC', () => {
       const orderBy = createFunction('SUM');
-      assertQueryEndsWith({ orderBy, orderByDirection: 'ASC' }, `ORDER BY SUM() ASC`);
+      assertQueryEndsWith({ sql: { orderBy, orderByDirection: 'ASC' } }, `ORDER BY SUM() ASC`);
     });
     it('should handle COUNT DESC', () => {
       const orderBy = createFunction('COUNT');
-      assertQueryEndsWith({ orderBy, orderByDirection: 'DESC' }, `ORDER BY COUNT() DESC`);
+      assertQueryEndsWith({ sql: { orderBy, orderByDirection: 'DESC' } }, `ORDER BY COUNT() DESC`);
     });
   });
   describe('limit', () => {
@@ -319,7 +352,7 @@ describe('SQLGenerator', () => {
     });
 
     it('should be added in case its specified', () => {
-      assertQueryEndsWith({ limit: 10 }, `LIMIT 10`);
+      assertQueryEndsWith({ sql: { limit: 10 } }, `LIMIT 10`);
     });
   });
 

public/app/plugins/datasource/cloudwatch/language/cloudwatch-sql/SQLGenerator.ts

@@ -10,26 +10,24 @@ import {
 } from '../../expressions';
 import { SQLExpression } from '../../types';
 
+const isAccountIdDefined = (accountId: string | undefined): boolean => !!(accountId && accountId !== 'all');
+
 export default class SQLGenerator {
   constructor(private templateSrv: TemplateSrv = getTemplateSrv()) {}
 
-  expressionToSqlQuery({
-    select,
-    from,
-    where,
-    groupBy,
-    orderBy,
-    orderByDirection,
-    limit,
-  }: SQLExpression): string | undefined {
+  expressionToSqlQuery(
+    { select, from, where, groupBy, orderBy, orderByDirection, limit }: SQLExpression,
+    accountId?: string
+  ): string | undefined {
     if (!from || !select?.name || !select?.parameters?.length) {
       return undefined;
     }
 
     let parts: string[] = [];
     this.appendSelect(select, parts);
     this.appendFrom(from, parts);
-    this.appendWhere(where, parts, true, where?.expressions?.length ?? 0);
+    this.appendAccountId(parts, accountId);
+    this.appendWhere(where, parts, true, where?.expressions?.length ?? 0, accountId);
     this.appendGroupBy(groupBy, parts);
     this.appendOrderBy(orderBy, orderByDirection, parts);
     this.appendLimit(limit, parts);
@@ -49,19 +47,31 @@ export default class SQLGenerator {
       : parts.push(this.formatValue(from?.property?.name ?? ''));
   }
 
+  private appendAccountId(parts: string[], accountId?: string) {
+    if (!isAccountIdDefined(accountId)) {
+      return;
+    }
+    parts.push(`WHERE AWS.AccountId = '${accountId}'`);
+  }
+
   private appendWhere(
     filter: QueryEditorExpression | undefined,
     parts: string[],
     isTopLevelExpression: boolean,
-    topLevelExpressionsCount: number
+    topLevelExpressionsCount: number,
+    accountId?: string
   ) {
     if (!filter) {
       return;
     }
 
     const hasChildExpressions = 'expressions' in filter && filter.expressions.length > 0;
     if (isTopLevelExpression && hasChildExpressions) {
-      parts.push('WHERE');
+      if (isAccountIdDefined(accountId)) {
+        parts.push('AND');
+      } else {
+        parts.push('WHERE');
+      }
     }
 
     if (filter.type === QueryEditorExpressionType.And) {