Skip to content

Commit

Permalink
Bugfix: dataset/view management
Browse files Browse the repository at this point in the history
  • Loading branch information
@darken99
darken99 committed May 12, 2022
1 parent 957646f commit aa8c3c7
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 16 deletions.
39 changes: 24 additions & 15 deletions cid/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -381,7 +381,7 @@ def delete(self, dashboard_id, **kwargs):
datasets = self.qs.dashboards.get(dashboard_id).datasets # save for later
else:
dashboard_definition = self.get_definition("dashboard", id=dashboard_id)
datasets = (dashboard_definition or {}).get('dependsOn', {}).get('datasets', [])
datasets = {d: None for d in (dashboard_definition or {}).get('dependsOn', {}).get('datasets', [])}

try:
# Execute query
Expand All @@ -398,23 +398,23 @@ def delete(self, dashboard_id, **kwargs):
return dashboard_id

print('Processing dependencies')
for dataset in datasets:
self.delete_dataset(dataset)
for dataset_name, dataset_id in datasets.items():
self.delete_dataset(name=dataset_name, id=dataset_id)

return dashboard_id

def delete_dataset(self, dataset_name: str):
if dataset_name not in self.resources['datasets']:
logger.info(f'Dataset {dataset_name} is not managed by CID. Skipping.')
print(f'Dataset {dataset_name} is not managed by CID. Skipping.')
def delete_dataset(self, name: str, id: str=None):
if name not in self.resources['datasets']:
logger.info(f'Dataset {name} is not managed by CID. Skipping.')
print(f'Dataset {name} is not managed by CID. Skipping.')
return False
for dataset in list(self.qs._datasets.values()):
if dataset.name == dataset_name:
if dataset.id == id or dataset.name == name:
# Check if dataset is used in some other dashboard
for dashboard in (self.qs.dashboards or {}).values():
if dataset.id in dashboard.datasets.values():
logger.info(f'Dataset {dataset.name} is still used by dashboard "{dashboard.id}". Skipping.')
print (f'Dataset {dataset.name} is still used by dashboard "{dashboard.id}". Skipping.')
logger.info(f'Dataset {dataset.name} ({dataset.id}) is still used by dashboard "{dashboard.id}". Skipping.')
print (f'Dataset {dataset.name} ({dataset.id}) is still used by dashboard "{dashboard.id}". Skipping.')
return False
else: #not used

Expand All @@ -428,25 +428,30 @@ def delete_dataset(self, dataset_name: str):
message=f'Delete QuickSight Dataset {dataset.name}?',
choices=['yes', 'no'],
default='no') == 'yes':
print(f'Deleting dataset {dataset.name}')
print(f'Deleting dataset {dataset.name} ({dataset.id})')
self.qs.delete_dataset(dataset.id)
else:
logger.info(f'Skipping dataset {dataset.name}')
print (f'Skipping dataset {dataset.name}')
return False
datasources = list(set(self.qs.get_datasets(id=dataset.id)[0].datasources))
athena_datasource = self.qs.datasources.get(datasources[0])
self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
break
else:
print(f'Not found dataset for deletion: {dataset_name}')
for view_name in list(set(self.resources['datasets'][dataset_name].get('dependsOn', {}).get('views', []))):
logger.info(f'Dataset not found for deletion: {name} ({id})')
for view_name in list(set(self.resources['datasets'][name].get('dependsOn', {}).get('views', []))):
self.delete_view(view_name)
return True

def delete_view(self, view_name):
if view_name not in self.resources['views']:
logger.info(f'View {view_name} is not managed by CID. Skipping.')
return False
logger.info(f'Deleting view "{view_name}"')
definition = self.get_definition("view", name=view_name)
if not definition:
print(f"Definition not found for view: {view_name}")
logger.info(f'Definition not found for view: "{view_name}"')
return False

for dashboard in (self.qs.dashboards or {}).values():
Expand Down Expand Up @@ -1003,6 +1008,7 @@ def create_or_update_view(self, view_name: str, recursive: bool=True, update: bo
# For account mappings create a view using a special helper
if view_name in self._visited_views: # avoid checking a views multiple times in one cid session
return
logger.info(f'Processing view: {view_name}')
self._visited_views.append(view_name)

if view_name in ['account_map', 'aws_accounts']:
Expand All @@ -1015,8 +1021,11 @@ def create_or_update_view(self, view_name: str, recursive: bool=True, update: bo
# Create a view
logger.info(f'Getting view definition')
view_definition = self.get_definition("view", name=view_name)
if not view_definition and view_name in self.athena._metadata.keys():
logger.info(f"Definition is unavailable but view exists: {view_name}, skipping")
return
logger.debug(f'View definition: {view_definition}')
logger.info(f'Processing view: {view_name}')

if recursive:
dependency_views = view_definition.get('dependsOn', dict()).get('views', list())
if 'cur' in dependency_views: dependency_views.remove('cur')
Expand Down
6 changes: 5 additions & 1 deletion cid/helpers/athena.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ def get_table_metadata(self, TableName: str) -> dict:
return table_metadata


def execute_query(self, sql_query, sleep_duration=1, database: str=None, catalog: str=None) -> str:
def execute_query(self, sql_query, sleep_duration=1, database: str=None, catalog: str=None, fail: bool=True) -> str:
""" Executes an AWS Athena Query """

# Set execution context
Expand Down Expand Up @@ -229,6 +229,8 @@ def execute_query(self, sql_query, sleep_duration=1, database: str=None, catalog
# Return result, either positive or negative
if (current_status == "SUCCEEDED"):
return query_id
elif not fail:
return False
else:
failure_reason = response['QueryExecution']['Status']['StateChangeReason']
logger.error('Athena query failed: {}'.format(failure_reason))
Expand Down Expand Up @@ -376,6 +378,7 @@ def delete_table(self, name: str, catalog: str=None, database: str=None):
f'DROP TABLE IF EXISTS {name};',
catalog=catalog,
database=database,
fail=False
)
except Exception as exc:
logger.debug(exc, stack_info=True)
Expand All @@ -399,6 +402,7 @@ def delete_view(self, name: str, catalog: str=None, database: str=None):
f'DROP VIEW IF EXISTS {name};',
catalog=catalog,
database=database,
fail=False
)
except Exception as exc:
logger.debug(exc, stack_info=True)
Expand Down

0 comments on commit aa8c3c7

Please sign in to comment.