qs_migration.py

import configparser
import boto3
import time

config = configparser.ConfigParser()
config.sections()
config.read('credentials.conf')

class QsMigration:

    def __init__(self):
        self.origin_acc = config["origin"]["account_id"]
        self.origin = boto3.client(
            service_name="quicksight",
            region_name=config["origin"]["region"],
            aws_access_key_id=config["origin"]["access_key"],
            aws_secret_access_key=config["origin"]["secret"],
        )

        self.destination_acc = config["destination"]["account_id"]
        self.destination = boto3.client(
            service_name="quicksight",
            region_name=config["destination"]["region"],
            aws_access_key_id=config["destination"]["access_key"],
            aws_secret_access_key=config["destination"]["secret"],
        )

        self.users_in_destination = self.get_all_users(self.destination, self.destination_acc)

    def _list_all(self, client: boto3.client, account_id: str, action: str, payload: dict = {}, next_token: str = ""):
        call_method_by_name = getattr(client, action) # Calls method `action` within class `client`

        if not next_token:
            request = call_method_by_name(AwsAccountId=account_id, **payload)
        else:
            request = call_method_by_name(AwsAccountId=account_id, **payload, NextToken=next_token)

        if request["Status"] != 200:
            return
        
        next_token = request.get("NextToken")
        del(request["ResponseMetadata"])
        del(request["Status"])

        return_key = next(iter(request))
        yield request[return_key]

        if next_token:
            yield from self.list_all_users(client, account_id, action, payload, next_token)

    def get_all_users(self, client: boto3.client, account_id: str):
        list_all_users = self._list_all(client, account_id, "list_users", {"Namespace": "default"})
        all_users = list()

        for page in list_all_users:
            for user in page:
                all_users.append(user["Arn"])

        return all_users                

    def create_data_set(self, data_set_id: str):
        describe_data_set = self.origin.describe_data_set(AwsAccountId=self.origin_acc, DataSetId=data_set_id)

        k_i = next(iter(describe_data_set["DataSet"]["PhysicalTableMap"]))
        k_n = next(iter(describe_data_set["DataSet"]["PhysicalTableMap"][k_i]))
        data_source_id = str(describe_data_set["DataSet"]["PhysicalTableMap"][k_i][k_n]["DataSourceArn"]).split("/")[-1]

        describe_data_source = self.origin.describe_data_source(AwsAccountId=self.origin_acc, DataSourceId=data_source_id)
        if describe_data_source["DataSource"]["Type"] != "REDSHIFT": # Temporary - only migrating Redshift first, for testing purposes. In the future, add feature to map data sources
            return
        
        describe_data_set["DataSet"]["PhysicalTableMap"][k_i][k_n]["DataSourceArn"] = "arn:aws:quicksight:us-west-2:492208021373:datasource/e6f9b64a-435d-4594-80a1-af5ffa9e3dd0" # ARN for Redshift data source in Fuji Dev account
        
        # If Row-level security enabled, replace dataset with the one in origin (needs to be created upfront)
        if "RowLevelPermissionDataSet" in describe_data_set["DataSet"].keys() and describe_data_set["DataSet"]["RowLevelPermissionDataSet"]["Status"] == "ENABLED":
            describe_data_set["DataSet"]["RowLevelPermissionDataSet"]["Arn"] = "arn:aws:quicksight:us-west-2:492208021373:dataset/86e4197a-5f2b-45eb-acb6-6d9563c5dec9"
        
        payload = {
            "AwsAccountId": self.destination_acc,
            "DataSetId": data_set_id,
            "Name": describe_data_set["DataSet"].get("Name"),
            "PhysicalTableMap": describe_data_set["DataSet"].get("PhysicalTableMap"),
            "LogicalTableMap": describe_data_set["DataSet"].get("LogicalTableMap"),
            "ImportMode": describe_data_set["DataSet"].get("ImportMode"),
            "ColumnGroups": describe_data_set["DataSet"].get("ColumnGroups"),
            "Permissions": [
                {
                    "Principal": user,
                    "Actions": ["quicksight:PassDataSet", "quicksight:DescribeIngestion", "quicksight:CreateIngestion", "quicksight:UpdateDataSet", "quicksight:DeleteDataSet", "quicksight:DescribeDataSet", "quicksight:CancelIngestion", "quicksight:DescribeDataSetPermissions", "quicksight:ListIngestions", "quicksight:UpdateDataSetPermissions"]
                } for user in self.users_in_destination
            ],
            "RowLevelPermissionDataSet": describe_data_set["DataSet"].get("RowLevelPermissionDataSet"),
            "Tags": describe_data_set["DataSet"].get("Tags")
        }

        # The code below removes all empty values from `payload` before passing it to the function. This way, only arguments with some real value are sent to the API
        to_func = dict()

        for key, value in payload.items():
            if value:
                to_func.update({
                    key: value
                })
        try:
            action = self.destination.create_data_set(**to_func)
            return action
        
        except self.destination.exceptions.ResourceExistsException:
            # If data set already exists
            return {
                "Arn": f"arn:aws:quicksight:us-west-2:{self.destination_acc}:dataset/{data_set_id}",
                "Name": describe_data_set["DataSet"].get("Name")
            }
       
    def run(self, dashboard_id: str):
        describe_dashboard = self.origin.describe_dashboard(AwsAccountId=self.origin_acc, DashboardId=dashboard_id)
        describe_dashboard_definition = self.origin.describe_dashboard_definition(AwsAccountId=self.origin_acc, DashboardId=dashboard_id)
        print(f"Describing dashboard {dashboard_id}")
        
        data_sets_in_dashboard = [ds.split("/")[-1] for ds in describe_dashboard["Dashboard"]["Version"]["DataSetArns"]]

        data_sets_in_dest = list()
        for data_set_id in data_sets_in_dashboard:
            action = self.create_data_set(data_set_id)
            data_sets_in_dest.append(action)
        print(f"Created {len(data_sets_in_dest)} data sets in destination")

        # Creating the dashboard
        print(f"Creating dashboard")
        for data_set_in_dashboard in describe_dashboard_definition["Definition"]["DataSetIdentifierDeclarations"]:
            data_set_in_dashboard["DataSetArn"] = data_set_in_dashboard["DataSetArn"].replace(self.origin_acc, self.destination_acc)

        payload = {
            "AwsAccountId": self.destination_acc,
            "DashboardId": dashboard_id,
            "Name": describe_dashboard["Dashboard"]["Name"],
            "Permissions": [
                {
                    "Principal": user,
                    "Actions": ["quicksight:DescribeDashboard", "quicksight:ListDashboardVersions", "quicksight:UpdateDashboardPermissions", "quicksight:QueryDashboard", "quicksight:UpdateDashboard", "quicksight:DeleteDashboard", "quicksight:DescribeDashboardPermissions", "quicksight:UpdateDashboardPublishedVersion"]
                } for user in self.users_in_destination
            ],
            "DashboardPublishOptions": describe_dashboard_definition.get("DashboardPublishOptions"),
            "Definition": describe_dashboard_definition.get("Definition"),
            "Parameters": describe_dashboard_definition.get("Parameters"),
            "Tags": describe_dashboard_definition.get("Tags"),
            "VersionDescription": describe_dashboard_definition.get("VersionDescription"),
        }

        # The code below removes all empty values from `payload` before passing it to the function. This way, only arguments with some real value are sent to the API
        to_func = dict()

        for key, value in payload.items():
            if value:
                to_func.update({
                    key: value
                })

        try:
            create_dashboard = self.destination.create_dashboard(**to_func)
            creation_status = create_dashboard["CreationStatus"]

            while creation_status != "CREATION_SUCCESSFUL":
                time.sleep(1)
                is_created = self.destination.describe_dashboard(AwsAccountId=self.destination_acc, DashboardId=dashboard_id)
                creation_status = is_created["Dashboard"]["Version"]["Status"]
            
            print(create_dashboard)
        except self.destination.exceptions.ResourceExistsException:
            print(f"Dashboard {dashboard_id} already exists")
        except Exception as e:
            import json
            print(e)
            with open("log.json", "w") as f:
                f.write(json.dumps(to_func))


if __name__ == "__main__":
    migration = QsMigration() 
    migration.run("377f32cc-85d0-4dce-9303-7eb6b1f92f90") # Spend Under Management
    # migration.run("cf061232-3d8c-4ba5-a0be-a671fa462a33") # Part Attributes Dashboard - UAT