diff --git a/admin/unmute_all_datasets.py b/admin/unmute_all_datasets.py new file mode 100644 index 0000000..8fc70cf --- /dev/null +++ b/admin/unmute_all_datasets.py @@ -0,0 +1,144 @@ +##### +# About: +# This script is indended to be used to UNMUTE ALL currently muted datasets within a specified warehouse +# Instructions: +# 1. Run this script, input your API Key ID, Token (generated in Settings -> API within MC UI) +# 2. If applicable, copy/paste the UUID of the warehouse you would like to target to unmute datasets +# Note: the script must be run for one warehouse at a time, run multiple times for multiple warehouses +# 3. Review the list of datasets to be unmuted in the .csv file provided by the prompt +# RECOMMENDATION: Keep this CSV file as a means to audit which datasets were unmuted by this script +# 4. Proceed to unmute the list of datasets +##### + +from pycarlo.core import Client, Query, Session +import csv +import json +from typing import Optional +from datetime import datetime + +def getWarehouses(mcdId,mcdToken): + client=Client(session=Session(mcd_id=mcdId,mcd_token=mcdToken)) + warehousesQuery = """ + query getUser { + getUser { + account { + warehouses { + name + connectionType + uuid + } + } + } + } + """ + + warehouses=client(warehousesQuery).get_user.account.warehouses + + if len(warehouses) == 1: + print(f"Found one warehouse - Name: {warehouses[0].name} - UUID: {warehouses[0].uuid}") + return warehouses[0].uuid + elif len(warehouses) > 1: + print("Found multiple warehouses... ") + for val in warehouses: + print("Name: " + val.name + ", Connection Type: " + val.connection_type + ", UUID: " + val.uuid) + dwId = input("Please copy/paste the full UUID of the warehouse you would like to target: ") + return dwId + +def get_dataset_query(dwId,first: Optional[int] = 1000, after: Optional[str] = None) -> Query: + query = Query() + get_datasets = query.get_datasets(first=first, dw_id=dwId, **(dict(after=after) if after else {})) + get_datasets.edges.node.__fields__("dataset","uuid","is_muted") + get_datasets.page_info.__fields__(end_cursor=True) + get_datasets.page_info.__fields__("has_next_page") + return query + +def getDatasetUuidDict(mcdId,mcdToken,dwId): + client=Client(session=Session(mcd_id=mcdId,mcd_token=mcdToken)) + dataset_uuid_dict={} + next_token=None + while True: + response = client(get_dataset_query(dwId=dwId,after=next_token)).get_datasets + for dataset in response.edges: + if dataset.node.is_muted: + dataset_uuid_dict[dataset.node.dataset.lower()] = dataset.node.uuid + if response.page_info.has_next_page: + next_token = response.page_info.end_cursor + else: + break + return dataset_uuid_dict + +def get_date(): + return datetime.today().strftime('%Y-%m-%d_%H_%M_%S') + +def userReview(uuid_dict, dw_id): + if not uuid_dict: + print(f"No muted datasets found in selected warehouse id {dw_id}. Exiting") + quit() + + fname = f"datasets_to_mute_{get_date()}.csv" + header = ['dataset', 'uuid'] + with open(fname, 'w') as csvfile: + writer = csv.writer(csvfile) + writer.writerow(header) + for dataset, uuid in uuid_dict.items(): + writer.writerow([dataset, uuid]) + userReview = input(f'Datasets to unmute written to file {fname} for your review. OK to proceed? (y/n) ').lower() + + if userReview == 'y': + return + else: + print("Acknowledged do not proceed. Exiting.") + quit() + +def generateVarsInput(uuid_list): + vars_input = { + "input": { + "datasets": uuid_list, + "mute": False + } + } + return vars_input + +def unmute_datasets(mcdId,dwId,mcdToken,uuidDict): + client=Client(session=Session(mcd_id=mcdId,mcd_token=mcdToken)) + uuid_list=[] + unmute_datasets_query = """ + mutation toggleMuteDatasets($input: ToggleMuteDatasetsInput!) { + toggleMuteDatasets(input: $input) { + muted { + uuid + isMuted + } + } + } + """ + + unmuted_dataset_counter = 0 + incremental_datasets = 0 + for uuid in uuidDict.values(): + temp_obj=dict(dsId=uuid, dwId=dwId) + print(temp_obj) + uuid_list.append(temp_obj) + unmuted_dataset_counter += 1 + incremental_datasets += 1 + if incremental_datasets == 99: + vars_input = generateVarsInput(uuid_list) + print(client(unmute_datasets_query, variables=vars_input)) + uuid_list.clear() + incremental_datasets = 0 + if incremental_datasets > 0: + vars_input = generateVarsInput(uuid_list) + print(client(unmute_datasets_query, variables=vars_input)) + print("Successfully Unmuted " + str(unmuted_dataset_counter) + " Datasets") + +if __name__ == '__main__': + #-------------------INPUT VARIABLES--------------------- + mcd_id = input("MCD ID: ") + mcd_token = input("MCD Token: ") + #------------------------------------------------------- + print("Getting warehouses...") + dw_id = getWarehouses(mcd_id, mcd_token) + print("Getting datasets...") + uuid_dict = getDatasetUuidDict(mcd_id,mcd_token,dw_id) + userReview(uuid_dict, dw_id) + unmute_datasets(mcd_id,dw_id,mcd_token,uuid_dict) \ No newline at end of file