Autodetect identity region and other fixes (#373)

* paginate datasources and fix identiy region check * add share in tests * autodetect identity region in cfn * manage flag parameters in commands Co-authored-by: Iakov Gan <[email protected]>
aws-samples · Sep 25, 2022 · 845f42b · 845f42b
1 parent 521d61b
commit 845f42b
Show file tree

Hide file tree

Showing 12 changed files with 61 additions and 28 deletions.
diff --git a/cfn-templates/cid-cfn.tests.bats b/cfn-templates/cid-cfn.tests.bats
@@ -76,7 +76,6 @@ setup_file() {
       CURTableName=""\
       CidVersion="$cid_version"\
       QuickSightDataSetRefreshSchedule="cron(0 4 * * ? *)"\
-      QuicksightIdentityRegion="us-east-1"\
       LambdaLayerBucketPrefix="aws-managed-cost-intelligence-dashboards"\
       Suffix=""\
     --stack-name "$stackname"

diff --git a/cfn-templates/cid-cfn.yml b/cfn-templates/cid-cfn.yml
@@ -32,7 +32,6 @@ Metadata:
           - CURTableName
           - CidVersion
           - Suffix
-          - QuicksightIdentityRegion
           - QuickSightDataSetRefreshSchedule
           - LambdaLayerBucketPrefix
   cfn-lint:
@@ -65,11 +64,6 @@ Parameters:
     Description: >
         REQUIRED - User name of QuickSight user from default namespace (as displayed in QuickSight admin panel).
         Dashboard created by this template with be owned by this user. See https://quicksight.aws.amazon.com/sn/admin#users
-  QuicksightIdentityRegion: #CID can detect QS Idenity region, but we need in permissions
-    Type: String
-    MinLength: 8
-    Description: Quicksight Identity region(can be different than the region you deploy the dashboard, typically us-east-1)
-    Default: us-east-1
   QuickSightDataSetRefreshSchedule:
     Type: String
     MinLength: 3
@@ -325,6 +319,8 @@ Resources:
           import botocore
           import urllib3
 
+          from cid.helpers import QuickSight # from layer
+
           BUCKET = os.environ['BUCKET']
           WORKGROUP = os.environ['WORKGROUP']
           CRAWLER = os.environ['CRAWLER']
@@ -333,10 +329,12 @@ Resources:
               print(event)
               type_ = event.get('RequestType', 'Undef')
               res = (True, f"Un error on {type_}. Check logs")
+              identity_region = ''
               try:
                   if type_ == 'Create': res = on_create()
                   elif type_ == 'Delete': res = on_delete()
                   else: res = (True, f"Not supported operation: {type_}")
+                  identity_region = get_identity_region()
               finally:
                   url = event.get('ResponseURL')
                   body = {}
@@ -347,7 +345,7 @@ Resources:
                   body['RequestId'] = event.get('RequestId')
                   body['LogicalResourceId'] = event.get('LogicalResourceId')
                   body['NoEcho'] = False
-                  body['Data'] =  {'Reason': res[1], 'uuid': str(uuid.uuid1()) }
+                  body['Data'] =  {'Reason': res[1], 'uuid': str(uuid.uuid1()), 'IdentityRegion': identity_region}
                   print(body)
                   if not url: return
                   json_body=json.dumps(body)
@@ -358,6 +356,10 @@ Resources:
                   except Exception as exc:
                       print("Failed sending PUT to CFN: " + str(exc))
 
+          def get_identity_region():
+              qs = QuickSight(boto3.session.Session())
+              return qs.identityRegion
+
           def on_create():
               if CRAWLER:
                   # FIXME: this can be replaced with AWS::Glue::Trigger
@@ -411,6 +413,8 @@ Resources:
                       log.append(f'ERROR: WorkGroup {WORKGROUP} Error: {exc}')
               print('\n'.join(log))
               return (True, '\n'.join(log))
+      Layers:
+        - !Ref CidResourceLambdaLayer
       Environment:
         Variables:
           BUCKET: !If [NeedAthenaQueryResultsBucket, !Ref MyAthenaQueryResultsBucket, '']
@@ -452,6 +456,10 @@ Resources:
                 Action:
                   - glue:StartCrawler
                 Resource: '*' # FIXME: use MyGlueCURCrawler
+              - Effect: Allow
+                Action:
+                  - quicksight:DescribeUser
+                Resource: '*' #FIXME: use !Sub 'arn:${AWS::Partition}:quicksight:*:${AWS::AccountId}:user/default/${QuickSightUser}'
       ManagedPolicyArns:
         - !Sub arn:${AWS::Partition}:iam::aws:policy/AWSLambdaExecute
   InitialSetup:
@@ -747,7 +755,7 @@ Resources:
             - 'quicksight:UpdateDataSource'
             - 'quicksight:DeleteDataSource'
             - 'quicksight:UpdateDataSourcePermissions'
-          Principal: !Sub 'arn:${AWS::Partition}:quicksight:${QuicksightIdentityRegion}:${AWS::AccountId}:user/default/${QuickSightUser}'
+          Principal: !Sub 'arn:${AWS::Partition}:quicksight:${InitialSetup.IdentityRegion}:${AWS::AccountId}:user/default/${QuickSightUser}'
 
   CidExecRole:
     Type: AWS::IAM::Role

diff --git a/cid/cli.py b/cid/cli.py
@@ -55,6 +55,8 @@ def main(ctx, **kwargs):
     ctx.obj = Cid(**kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @cid_command
 def map(ctx, **kwargs):
     """Create account mapping
@@ -70,6 +72,10 @@ def map(ctx, **kwargs):
     ctx.obj.map(**kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
+@click.option('--share-with-account', help='Share dashboard with all users in the current account', is_flag=True, default=False)
+@click.option('--quicksight-delete-failed-datasource', help='Delete datasoruce if creation failed', is_flag=True, default=False)
 @cid_command
 def deploy(ctx, **kwargs):
     """Deploy Dashboard
@@ -82,18 +88,18 @@ def deploy(ctx, **kwargs):
      --glue-data-catalog TEXT              Glue data catalog
      --cur-table-name TEXT                 CUR table name
      --quicksight-datasource-id TEXT      QuickSight Datasource ARN (if not found one with provided Athena workgroup)
-     --quicksight-delete-failed-datasource (yes|no) Delete datasoruce if creation failed
      --quicksight-user TEXT                QuickSight user
      --dataset-{dataset_name}-id TEXT      QuickSight dataset id for a specific dataset
      --view-{view_name}-{parameter} TEXT   a custom parameter for a view creation, can use variable: {account_id}
      --account-map-source TEXT             csv, dummy, organization (if autodiscovery impossible)
      --account-map-file TEXT               csv file path relative to current directory (if autodiscovery impossible and csv selected as a source )
      --resources TEXT                      CID resources file (yaml)
-     --share-with-account ['yes/no']       Share dashboard with all users in the current account
     """
     ctx.obj.deploy(**kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @cid_command
 def export(ctx, **kwargs):
     """Deploy Dashboard
@@ -111,14 +117,17 @@ def export(ctx, **kwargs):
 
 
 @click.option('--dashboard-id', help='QuickSight dashboard id', default=None)
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @cid_command
 def status(ctx, dashboard_id, **kwargs):
     """Show Dashboard status"""
     ctx.obj.status(dashboard_id, **kwargs)
 
 
 @click.option('--dashboard-id', help='QuickSight dashboard id', default=None)
-@click.option('-y', '--yes', help='Answer Yes to all confirmation questions', default=False, is_flag=True)
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @cid_command
 def delete(ctx, dashboard_id, **kwargs):
     """Delete Dashboard and all dependencies unused by other CID-managed dasboards
@@ -132,6 +141,8 @@ def delete(ctx, dashboard_id, **kwargs):
     ctx.obj.delete(dashboard_id, **kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @click.option('--dashboard-id', help='QuickSight dashboard id', default=None)
 @click.option('--force/--noforce', help='allow selecting  up to date dashboards (flags must be before options)', default=False)
 @click.option('--recursive/--norecursive', help='Recursive update all Datasets and Views (flags must be before options)', default=False)
@@ -141,19 +152,25 @@ def update(ctx, dashboard_id, force, recursive, **kwargs):
     ctx.obj.update(dashboard_id, force=force, recursive=recursive, **kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @click.option('--dashboard-id', help='QuickSight dashboard id', default=None)
 @cid_command
 def open(ctx, dashboard_id, **kwargs):
     """Open Dashboard in browser"""
     ctx.obj.open(dashboard_id, **kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @cid_command
 def cleanup(ctx, **kwargs):
     """Delete unused resources (QuickSight datasets, Athena views)"""
     ctx.obj.cleanup(**kwargs)
 
 
+@click.option('-v', '--verbose', count=True)
+@click.option('-y', '--yes', help='confirm all', is_flag=True, default=False)
 @click.option('--dashboard-id', help='QuickSight dashboard id', default=None)
 @click.option('--share-method', help='Sharing method', default=None, type=click.Choice(['folder', 'user', 'account']))
 @click.option('--folder-method', help='Folder to use', default=None, type=click.Choice(['new', 'existing']))

diff --git a/cid/common.py b/cid/common.py
@@ -29,11 +29,6 @@
 class Cid():
 
     def __init__(self, **kwargs) -> None:
-        set_cid_logger(
-            verbosity=kwargs.pop('verbose'),
-            log_filename=kwargs.pop('log_filename', 'cid.log')
-        )
-        logger.info(f'Initializing CID {__version__}')
         self.base: CidBase = None
         # Defined resources
         self.resources = dict()
@@ -148,7 +143,14 @@ def command(func):
         '''
         @functools.wraps(func)
         def wrap(self, *args, **kwargs):
-            set_parameters(kwargs)
+            set_cid_logger(
+                verbosity=kwargs.pop('verbose'),
+                log_filename=kwargs.pop('log_filename', 'cid.log')
+            )
+            logger.info(f'Initializing CID {__version__} for {func.__name__}')
+            all_yes = kwargs.get('yes', None)
+            set_parameters(kwargs, all_yes=all_yes)
+            logger.debug(json.dumps(get_parameters()))
             self.aws_login()
             self.load_resources()
             return func(self, *args, **kwargs)

diff --git a/cid/helpers/quicksight/__init__.py b/cid/helpers/quicksight/__init__.py
@@ -14,7 +14,7 @@
 from cid.helpers.quicksight.dashboard import Dashboard
 from cid.helpers.quicksight.dataset import Dataset
 from cid.helpers.quicksight.datasource import Datasource
-from cid.utils import get_parameter
+from cid.utils import get_parameter, get_parameters
 
 logger = logging.getLogger(__name__)
 
@@ -50,8 +50,9 @@ def AthenaWorkGroup(self, value):
     @property
     def user(self) -> dict:
         if not self._user:
+            username = get_parameters().get('quicksight-user', self.username)
             try:
-                self._user =  self.describe_user(self.username)
+                self._user =  self.describe_user(username)
             except Exception as exc:
                 logger.debug(exc, stack_info=True)
                 logger.error(f'Failed to find your QuickSight username ({exc}). Is QuickSight activated?')
@@ -69,9 +70,10 @@ def identityRegion(self) -> str:
         if not self._identityRegion:
             try:
                 logger.info(f'Detecting QuickSight identity region, trying {self.region}')
+                username = get_parameters().get('quicksight-user', self.username)
                 parameters = {
                     'AwsAccountId': self.account_id,
-                    'UserName': '/'.join(self.username),
+                    'UserName': username,
                     'Namespace': 'default'
                 }
                 self.client.describe_user(**parameters)
@@ -410,14 +412,11 @@ def list_data_sources(self) -> list:
         parameters = {
             'AwsAccountId': self.account_id
         }
+        data_sources = []
         try:
-            result = self.client.list_data_sources(**parameters)
-            logger.debug(result)
-            if result.get('Status') != 200:
-                print(f'Error, {result}')
-                exit()
-            else:
-                return result.get('DataSources')
+            for page in self.client.get_paginator('list_data_sources').paginate(**parameters):
+                data_sources += page.get('DataSources',[])
+            return data_sources
         except self.client.exceptions.AccessDeniedException:
             logger.info('Access denied listing data sources')
             raise

diff --git a/cid/logger.py b/cid/logger.py
@@ -21,6 +21,7 @@ def log_method(self, message, *args, **kwargs):
     def log_to_root(message, *args, **kwargs):
         logging.log(num, message, *args, **kwargs)
 
+    if hasattr(logging, name): return # Already set
     logging.addLevelName(num, name)
     setattr(logging, name, num)
     setattr(logging.getLoggerClass(), method, log_method)

diff --git a/cid/test/bats/10-deploy-update-delete/compute-optimizer-dashboard.bats b/cid/test/bats/10-deploy-update-delete/compute-optimizer-dashboard.bats
@@ -5,6 +5,7 @@ account_id=$(aws sts get-caller-identity --query "Account" --output text )
 @test "Install" {
   run cid-cmd -vv deploy  \
     --dashboard-id compute-optimizer-dashboard \
+    --share-with-account \
     --athena-database 'optimization_data' \
     --view-compute-optimizer-lambda-lines-s3FolderPath       's3://costoptimizationdata{account_id}/Compute_Optimizer/Compute_Optimizer_ec2_lambda' \
     --view-compute-optimizer-ebs-volume-lines-s3FolderPath   's3://costoptimizationdata{account_id}/Compute_Optimizer/Compute_Optimizer_ebs_volume' \

diff --git a/cid/test/bats/10-deploy-update-delete/cost_intelligence_dashboard.bats b/cid/test/bats/10-deploy-update-delete/cost_intelligence_dashboard.bats
@@ -9,6 +9,7 @@ database_name="${database_name:-athenacurcfn_cur1}" # If variable not set or nul
     --dashboard-id cost_intelligence_dashboard \
     --athena-database $database_name\
     --account-map-source dummy \
+    --share-with-account \
 
   [ "$status" -eq 0 ]
 }

diff --git a/cid/test/bats/10-deploy-update-delete/cudos.bats b/cid/test/bats/10-deploy-update-delete/cudos.bats
@@ -11,6 +11,7 @@ quicksight_user="${quicksight_user:-cicd-staging}" # If variable not set or null
     --athena-database $database_name\
     --account-map-source dummy \
     --quicksight-user $quicksight_user \
+    --share-with-account \
 
   [ "$status" -eq 0 ]
 }

diff --git a/cid/test/bats/10-deploy-update-delete/kpi_dashboard.bats b/cid/test/bats/10-deploy-update-delete/kpi_dashboard.bats
@@ -9,6 +9,7 @@ database_name="${database_name:-athenacurcfn_cur1}" # If variable not set or nul
     --dashboard-id kpi_dashboard \
     --athena-database $database_name\
     --account-map-source dummy \
+    --share-with-account \
 
   [ "$status" -eq 0 ]
 }

diff --git a/cid/test/bats/10-deploy-update-delete/ta-organizational-view.bats b/cid/test/bats/10-deploy-update-delete/ta-organizational-view.bats
@@ -6,6 +6,8 @@ account_id=$(aws sts get-caller-identity --query "Account" --output text )
   run cid-cmd -vv deploy  \
     --dashboard-id ta-organizational-view \
     --athena-database 'optimization_data' \
+    --share-with-account \
+
     --view-ta-organizational-view-reports-s3FolderPath "s3://costoptimizationdata$account_id/optics-data-collector/ta-data'"
 
   [ "$status" -eq 0 ]

diff --git a/cid/test/bats/10-deploy-update-delete/trends-dashboard.bats b/cid/test/bats/10-deploy-update-delete/trends-dashboard.bats
@@ -9,6 +9,7 @@ database_name="${database_name:-athenacurcfn_cur1}" # If variable not set or nul
     --dashboard-id trends-dashboard \
     --athena-database $database_name\
     --account-map-source dummy \
+    --share-with-account \
 
   [ "$status" -eq 0 ]
 }