bugfix: avoid workgroup scan (#460)

Author: iakov-aws

cid/common.py

@@ -526,7 +526,7 @@ def delete_dataset(self, name: str, id: str=None):
                     continue
                 datasources = dataset.datasources
                 athena_datasource = self.qs.datasources.get(datasources[0])
-                if athena_datasource:
+                if athena_datasource and not get_parameters().get('athena-workgroup'):
                     self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
                     break
                 logger.debug(f'Cannot find QuickSight DataSource {datasources[0]}. So cannot define Athena WorkGroup')
@@ -1109,10 +1109,12 @@ def create_or_update_dataset(self, dataset_definition: dict, dataset_id: str=Non
                         )
                         athena_datasource = self.qs.athena_datasources[datasource_id]
                         logger.info(f'Found {len(datasources)} Athena datasources, not using {athena_datasource.id}')
-        if isinstance(athena_datasource, Datasource) and athena_datasource.AthenaParameters.get('WorkGroup', None):
-            self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
-        else:
-            logger.debug('Athena_datasource is not defined. Will only create views')
+        if not get_parameters().get('athena-workgroup'):
+            # set default workgroup from datasource if not provided via parameters
+            if isinstance(athena_datasource, Datasource) and athena_datasource.AthenaParameters.get('WorkGroup', None):
+                self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
+            else:
+                logger.debug('Athena_datasource is not defined. Will only create views')
 
         # Check for required views
         _views = dataset_definition.get('dependsOn', {}).get('views', [])

cid/helpers/athena.py

@@ -17,7 +17,7 @@ class Athena(CidBase):
     defaults = {
         'CatalogName': 'AwsDataCatalog',
         'DatabaseName': 'customer_cur_data',
-        'WorkGroup': 'primary'
+        'WorkGroup': 'CID'
     }
     _CatalogName = None
     _DatabaseName = None
@@ -115,6 +115,9 @@ def DatabaseName(self, database):
     def WorkGroup(self) -> str:
         """ Select AWS Athena workgroup """
         if not self._WorkGroup:
+            if get_parameters().get('athena-workgroup'):
+                self.WorkGroup = get_parameters().get('athena-workgroup')
+                return self._WorkGroup
             logger.info('Selecting Athena workgroup...')
             workgroups = self.list_work_groups()
             logger.info(f'Found {len(workgroups)} workgroups: {", ".join([wg.get("Name") for wg in workgroups])}')