Minor fixes (#610)

aws-samples · Sep 27, 2023 · ad78189 · ad78189
1 parent 2ea36ca
commit ad78189
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/cid/helpers/athena.py b/cid/helpers/athena.py
@@ -236,15 +236,17 @@ def execute_query(self, sql_query, sleep_duration=1, database: str=None, catalog
 
             # Get Query ID
             query_id = response.get('QueryExecutionId')
-
-            # Get Query Status
+            if not query_id:
+                logger.debug(f'Full query: {sql_query}')
+                raise CidCritical(f'Athena cannot start query. Answer is: {response}')
+            # Get Query Status for the first time
             query_status = self.client.get_query_execution(QueryExecutionId=query_id)
-        except self.client.exceptions.InvalidRequestException as e:
+        except self.client.exceptions.InvalidRequestException as exc:
             logger.debug(f'Full query: {sql_query}')
-            raise CidCritical(f'InvalidRequestException: {e}')
-        except Exception as e:
+            raise CidCritical(f'InvalidRequestException: {exc}') from exc
+        except Exception as exc:
             logger.debug(f'Full query: {sql_query}')
-            raise CidCritical(f'Athena query failed: {e}')
+            raise CidCritical(f'Athena query failed: {exc}') from exc
 
         current_status = query_status['QueryExecution']['Status']['State']
 
@@ -256,16 +258,14 @@ def execute_query(self, sql_query, sleep_duration=1, database: str=None, catalog
             # Sleep before polling again
             time.sleep(sleep_duration)
 
-        if (current_status == "SUCCEEDED"):
+        if current_status == "SUCCEEDED":
             return query_id
-        else:
-            failure_reason = response.get('QueryExecution', {}).get('Status', {}).get('StateChangeReason',repr(response))
-            logger.info(f'Athena query failed: {failure_reason}')
-            logger.debug(f'Full query: {sql_query}')
-            if fail:
-                raise CidCritical(f'Athena query failed: {failure_reason}')
-            else:
-                return False
+        failure_reason = response.get('QueryExecution', {}).get('Status', {}).get('StateChangeReason',repr(response))
+        logger.info(f'Athena query failed: {failure_reason}')
+        logger.debug(f'Full query: {sql_query}')
+        if fail:
+            raise CidCritical(f'Athena query failed: {failure_reason}')
+        return False
 
     def get_query_results(self, query_id):
         """ Get Query Results """

diff --git a/cid/helpers/quicksight/__init__.py b/cid/helpers/quicksight/__init__.py
@@ -1072,9 +1072,9 @@ def get_dataset_refresh_schedules(self, dataset_id):
             )
             return refresh_schedules.get("RefreshSchedules")
         except self.client.exceptions.ResourceNotFoundException as exc:
-            raise CidError(f'DataSource {dataset_id} does not exist') from exc
+            raise CidError(f'DataSet {dataset_id} does not exist') from exc
         except self.client.exceptions.AccessDeniedException as exc:
-            raise CidError(f'No quicksight:ListDataSets permission') from exc
+            raise CidError(f'AccessDenied') from exc
         except Exception as exc:
             raise CidError(f'Unable to list refresh schedules for dataset {dataset_id}: {str(exc)}') from exc
 
@@ -1085,7 +1085,7 @@ def ensure_dataset_refresh_schedule(self, dataset_id, schedules: list):
         try:
             existing_schedules = self.get_dataset_refresh_schedules(dataset_id)
         except CidError as exc:
-            logger.debug(exc, exc_info=True)
+            logger.debug(f'List refresh schedule throws: {exc}')
             logger.warning(
                 f'Cannot read dataset schedules for dataset = {dataset_id}. {str(exc)}. Skipping schedule management.'
                 ' Please make sure scheduled refresh is configured manualy.'
@@ -1103,7 +1103,7 @@ def ensure_dataset_refresh_schedule(self, dataset_id, schedules: list):
                     default=default_timezone
                 )
             if not schedule_frequency_timezone:
-                logger.warning(f'Cannot get timezone. Please provide --timezone parameter. Please make sure scheduled refresh is configured manualy.')
+                logger.warning('Cannot get timezone. Please provide --timezone parameter. Please make sure scheduled refresh is configured manualy.')
                 return
 
         for schedule in schedules: