|
10 | 10 | import click |
11 | 11 |
|
12 | 12 | from cid.base import CidBase |
13 | | -from cid.helpers import diff |
| 13 | +from cid.helpers import diff, timezone, randtime |
14 | 14 | from cid.helpers.quicksight.dashboard import Dashboard |
15 | 15 | from cid.helpers.quicksight.dataset import Dataset |
16 | 16 | from cid.helpers.quicksight.datasource import Datasource |
17 | 17 | from cid.helpers.quicksight.template import Template as CidQsTemplate |
18 | | -from cid.utils import get_parameter, get_parameters |
| 18 | +from cid.utils import get_parameter, get_parameters, exec_env, cid_print |
19 | 19 | from cid.exceptions import CidCritical, CidError |
20 | 20 |
|
21 | 21 | logger = logging.getLogger(__name__) |
@@ -1034,6 +1034,123 @@ def update_dataset(self, definition: dict) -> Dataset: |
1034 | 1034 | return self.describe_dataset(dataset_id) |
1035 | 1035 |
|
1036 | 1036 |
|
| 1037 | + def get_dataset_refresh_schedules(self, dataset_id): |
| 1038 | + """Returns refresh schedules for given dataset id""" |
| 1039 | + try: |
| 1040 | + refresh_schedules = self.client.list_refresh_schedules( |
| 1041 | + AwsAccountId=self.account_id, |
| 1042 | + DataSetId=dataset_id |
| 1043 | + ) |
| 1044 | + return refresh_schedules.get("RefreshSchedules") |
| 1045 | + except self.client.exceptions.ResourceNotFoundException as exc: |
| 1046 | + raise CidError(f'DataSource {dataset_id} does not exist') from exc |
| 1047 | + except self.client.exceptions.AccessDeniedException as exc: |
| 1048 | + raise CidError(f'No quicksight:ListDataSets permission') from exc |
| 1049 | + except Exception as exc: |
| 1050 | + raise CidError(f'Unable to list refresh schedules for dataset {dataset_id}: {str(exc)}') from exc |
| 1051 | + |
| 1052 | + |
| 1053 | + def ensure_dataset_refresh_schedule(self, dataset_id, schedules: list): |
| 1054 | + """ Ensures that dataset has scheduled refresh """ |
| 1055 | + # get all existing schedules for the given dataset |
| 1056 | + try: |
| 1057 | + existing_schedules = self.get_dataset_refresh_schedules(dataset_id) |
| 1058 | + except CidError as exc: |
| 1059 | + logger.debug(exc, exc_info=True) |
| 1060 | + logger.warning( |
| 1061 | + f'Cannot read dataset schedules for dataset = {dataset_id}. {str(exc)}. Skipping schedule management.' |
| 1062 | + ' Please make sure scheduled refresh is configured manualy.' |
| 1063 | + ) |
| 1064 | + return |
| 1065 | + |
| 1066 | + if schedules: |
| 1067 | + if exec_env()['terminal'] in ('lambda'): |
| 1068 | + schedule_frequency_timezone = get_parameters().get("timezone", timezone.get_default_timezone()) |
| 1069 | + else: |
| 1070 | + schedule_frequency_timezone = get_parameter("timezone", |
| 1071 | + message='Please select timezone for datasets scheduled refresh time', |
| 1072 | + choices=timezone.get_all_timezones(), |
| 1073 | + default=timezone.get_default_timezone() |
| 1074 | + ) |
| 1075 | + |
| 1076 | + for schedule in schedules: |
| 1077 | + |
| 1078 | + # Get the list of exising schedules with the same id |
| 1079 | + existing_schedule = None |
| 1080 | + for existing in existing_schedules: |
| 1081 | + if schedule["ScheduleId"] == existing["ScheduleId"]: |
| 1082 | + existing_schedule = existing |
| 1083 | + break |
| 1084 | + |
| 1085 | + # Verify that all schedule parameters are set |
| 1086 | + schedule["ScheduleId"] = schedule.get("ScheduleId", "cid") |
| 1087 | + if "ScheduleFrequency" not in schedule: |
| 1088 | + schedule["ScheduleFrequency"] = {} |
| 1089 | + schedule["ScheduleFrequency"]["Timezone"] = schedule_frequency_timezone |
| 1090 | + try: |
| 1091 | + schedule["ScheduleFrequency"]["TimeOfTheDay"] = randtime.get_random_time_from_range( |
| 1092 | + self.account_id + dataset_id, |
| 1093 | + schedule["ScheduleFrequency"].get("TimeOfTheDay", "") |
| 1094 | + ) |
| 1095 | + except Exception as exc: |
| 1096 | + logger.error( |
| 1097 | + f'Invalid timerange for schedule with id "{schedule["ScheduleId"]}"' |
| 1098 | + f' and dataset {dataset_id}: {str(exc)} ... skipping.' |
| 1099 | + f' Please create dataset refresh schedule manually.' |
| 1100 | + ) |
| 1101 | + continue |
| 1102 | + schedule["ScheduleFrequency"]["Interval"] = schedule["ScheduleFrequency"].get("Interval", "DAILY") |
| 1103 | + schedule["RefreshType"] = schedule.get("RefreshType", "FULL_REFRESH") |
| 1104 | + if "providedBy" in schedule: |
| 1105 | + del schedule["providedBy"] |
| 1106 | + |
| 1107 | + if not existing_schedule: |
| 1108 | + # Avoid adding a new schedule when customer already has put a schedule manually as this can lead to additional charges. |
| 1109 | + schedules_with_different_id = [existing for existing in existing_schedules if schedule["ScheduleId"] != existing["ScheduleId"] ] |
| 1110 | + if schedules_with_different_id: |
| 1111 | + logger.info( |
| 1112 | + f'Found the same schedule {schedule.get("RefreshType")} / {schedule["ScheduleFrequency"].get("Interval")}' |
| 1113 | + f' but with different id. Skipping to avoid duplication. Please delete all manually created schedules for dataset {dataset_id}' |
| 1114 | + ) |
| 1115 | + continue |
| 1116 | + logger.debug(f'Creating refresh schedule with id {schedule["ScheduleId"]} for dataset {dataset_id}.') |
| 1117 | + try: |
| 1118 | + self.client.create_refresh_schedule( |
| 1119 | + DataSetId=dataset_id, |
| 1120 | + AwsAccountId=self.account_id, |
| 1121 | + Schedule=schedule |
| 1122 | + ) |
| 1123 | + logger.debug(f'Refresh schedule with id {schedule["ScheduleId"]} for dataset {dataset_id} is created.') |
| 1124 | + except self.client.exceptions.ResourceNotFoundException: |
| 1125 | + logger.error(f'Unable to create refresh schedule with id {schedule["ScheduleId"]}. Dataset {dataset_id} does not exist.') |
| 1126 | + except self.client.exceptions.AccessDeniedException: |
| 1127 | + logger.error(f'Unable to create refresh schedule with id {schedule["ScheduleId"]}. Please add quicksight:CreateDataSet permission.') |
| 1128 | + except Exception as exc: |
| 1129 | + logger.error(f'Unable to create refresh schedule with id {schedule["ScheduleId"]} for dataset "{dataset_id}": {str(exc)}') |
| 1130 | + else: |
| 1131 | + # schedule exists so we need to update |
| 1132 | + logger.debug(f'Updating refresh schedule with id {schedule["ScheduleId"]} for dataset {dataset_id}.') |
| 1133 | + try: |
| 1134 | + self.client.update_refresh_schedule( |
| 1135 | + DataSetId=dataset_id, |
| 1136 | + AwsAccountId=self.account_id, |
| 1137 | + Schedule=schedule |
| 1138 | + ) |
| 1139 | + logger.debug(f'Refresh schedule with id {schedule["ScheduleId"]} for dataset {dataset_id} is updated.') |
| 1140 | + except self.client.exceptions.ResourceNotFoundException: |
| 1141 | + logger.error(f'Unable to update refresh schedule with id {schedule["ScheduleId"]}. Dataset {dataset_id} does not exist.') |
| 1142 | + except self.client.exceptions.AccessDeniedException: |
| 1143 | + logger.error(f'Unable to update refresh schedule with id {schedule["ScheduleId"]}. Please add quicksight:UpdqteDataSet permission.') |
| 1144 | + except Exception as exc: |
| 1145 | + logger.error(f'Unable to update refresh schedule with id {schedule["ScheduleId"]} for dataset "{dataset_id}": {str(exc)}') |
| 1146 | + |
| 1147 | + # Verify if there is at least one schedule and warn user if not |
| 1148 | + try: |
| 1149 | + if not self.get_dataset_refresh_schedules(dataset_id): |
| 1150 | + logger.warning(f'There is no refresh schedule for dataset "{dataset_id}". Please create a refresh schedule manually.' ) |
| 1151 | + except CidError: |
| 1152 | + pass |
| 1153 | + |
1037 | 1154 | def create_dashboard(self, definition: dict) -> Dashboard: |
1038 | 1155 | """ Creates an AWS QuickSight dashboard """ |
1039 | 1156 |
|
|
0 commit comments