2
\$\begingroup\$

I wrote code that stores an AWS Organization tree in a tree data structure in Rust. From the tree, I use a slightly modified preorder traversal algorithm to print all of the accounts in the organization in a table, while keeping track of all ancestors for each account.

AWS Organizations are composed of a root organizational unit (OU), which can contain other OUs or accounts. OUs can contain OUs or Accounts. In tree parlance, accounts are necessarily leaf nodes, while OUs can be branches or leaves. There is no ordering defined on the tree.

In my code, I represent an AWS organization as a tree data structure. I want to be able to output the names and, OU membership of all accounts in the organization. Say we had the following tree:

         R
        / \
       /   \
      OU1  A4
     /  \    
    /    \
  OU11   A3
  / \ 
 /   \
A1   A2

The algorithm would print something like:

A4
OU1:A3
OU1:OU11:A2
OU1:OU11:A1

To achieve that, I use a slightly modified iterative preorder traversal, with an additional variable that I called height_stack. This stack keeps track of the number of children that we have left to pop in this branch. When we push a branch to the stack, we push the number of children that the branch has to the height_stack. We decrement the last element of that height stack every time we pop the stack. At the end of each iteration, we check whether the last element of the height stack is zero, in which case we pop all zero elements of the height stack.

Below is the code I came up with. I show both the recursive construction of the tree (recursively_build_account_tree) and the traversal (list_accounts_per_ou), as the algorithm assumes that OUs come before accounts in the children Vec.

Is there a better algorithm to do this for a tree with an unknown number of children by node? (I know there is a lot of Rusoto boilerplate, hopefully it doesn't make the code too hard to read.)

use comfy_table::{Attribute, Cell, ContentArrangement, Row, Table, TableComponent};
use rusoto_organizations::{
    Account, DescribeAccountRequest, DescribeOrganizationalUnitRequest, ListChildrenRequest,
    ListRootsRequest, OrganizationalUnit, Organizations, OrganizationsClient, Root,
};
use std::{thread, time};

/// In an AWS Organizations, each node of the account tree can either be an OU or an account.
#[derive(Debug)]
enum OrgNode {
    Ou(OrganizationalUnit),
    Account(Account),
}

/// m-ary tree to represent the AWS organization.
#[derive(Debug)]
struct OrgTree {
    root: OrgNode,
    children: Option<Vec<Box<OrgTree>>>,
}

impl OrgTree {
    pub fn new(root: OrgNode) -> OrgTree {
        OrgTree {
            root: root,
            children: None,
        }
    }
}

/// Starting from a root of the AWS Organization, we first request all of its child OUs, and then
/// all of its child accounts. This ensures that preorder traversal list all accounts that belong
/// to an OU before listing accounts that belong to nested OUs.
///
/// The `thread::sleep` in the loops are necessary to not hit AWS Organizations' API limits.
///
/// # Example
///
/// Say we had the following tree:
/// ```ascii
///
///                            R
///                           / \
///                          /   \
///                         OU    A
///                        / | \
///                       /  |  \
///                      /   |   \
///                     OU  A1    A2
/// ```
/// Preorder traversal would print nodes in this order:
///  * R:A
///  * R:OU:A2
///  * R:OU:A1
fn recursively_build_account_tree(client: &OrganizationsClient, node: &mut OrgTree) {
    match &node.root {
        OrgNode::Ou(v) => {
            let list_children_request = ListChildrenRequest {
                parent_id: v.id.as_ref().unwrap().to_string(),
                child_type: "ORGANIZATIONAL_UNIT".to_string(),
                max_results: None,
                next_token: None,
            };

            let list_children_response = client
                .list_children(list_children_request)
                .sync()
                .unwrap()
                .children
                .unwrap();

            if list_children_response.len() > 0 {
                node.children = Some(Vec::new());

                for element in list_children_response.iter() {
                    // Describe the OU.
                    let describe_org_unit_request = DescribeOrganizationalUnitRequest {
                        organizational_unit_id: element.id.as_ref().unwrap().to_string(),
                    };

                    let describe_org_unit_response = client
                        .describe_organizational_unit(describe_org_unit_request)
                        .sync()
                        .unwrap()
                        .organizational_unit
                        .unwrap();
                    if let Some(v) = &mut node.children {
                        v.push(Box::new(OrgTree::new(OrgNode::Ou(
                            describe_org_unit_response,
                        ))));

                        thread::sleep(time::Duration::from_millis(200));
                    }
                }
            }

            // Request accounts in this OU.
            let list_children_request = ListChildrenRequest {
                parent_id: v.id.as_ref().unwrap().to_string(),
                child_type: "ACCOUNT".to_string(),
                max_results: None,
                next_token: None,
            };

            let list_children_response = client
                .list_children(list_children_request)
                .sync()
                .unwrap()
                .children
                .unwrap();

            if list_children_response.len() > 0 {
                match &mut node.children {
                    Some(_) => (),
                    None => {
                        node.children = Some(Vec::new());
                    }
                }

                for element in list_children_response {
                    let describe_account_request = DescribeAccountRequest {
                        account_id: element.id.unwrap(),
                    };

                    let describe_account_response = client
                        .describe_account(describe_account_request)
                        .sync()
                        .unwrap()
                        .account
                        .unwrap();

                    node.children
                        .as_mut()
                        .unwrap()
                        .push(Box::new(OrgTree::new(OrgNode::Account(
                            describe_account_response,
                        ))));

                    thread::sleep(time::Duration::from_millis(200));
                }
            }

            // Recursively build the tree.
            match &mut node.children {
                Some(v) => {
                    for element in v {
                        recursively_build_account_tree(client, &mut *element);
                    }
                }

                None => (),
            };
        }

        OrgNode::Account(_) => (),
    }
}

/// Fetches all of the accounts in the AWS Organizations and outputs
/// them in a Markdown-compatible table.
pub fn list_accounts_per_ou(client: &OrganizationsClient) -> Table {
    let root_request = ListRootsRequest {
        max_results: None,
        next_token: None,
    };

    let root: Root = client
        .list_roots(root_request)
        .sync()
        .unwrap()
        .roots
        .unwrap()[0]
        .clone();

    // Coerce the root into an OU (kinda cheating).
    let root_as_ou = OrganizationalUnit {
        arn: root.arn,
        id: root.id,
        name: root.name,
    };

    let mut org_tree = OrgTree::new(OrgNode::Ou(root_as_ou));

    // Recursively build the account tree.
    match &org_tree.root {
        OrgNode::Ou(_) => {
            recursively_build_account_tree(&client, &mut org_tree);
        }

        _ => panic!(
            "The root node should be an OU, no whathever it is: {:?}",
            org_tree
        ),
    }

    let mut table = Table::new();
    const MARKDOWN: &str = "||  |-|||           ";
    table.load_preset(MARKDOWN).set_header(vec![
        Cell::new("Account name").add_attribute(Attribute::Bold),
        Cell::new("Account ID").add_attribute(Attribute::Bold),
        Cell::new("OU name").add_attribute(Attribute::Bold),
        Cell::new("Email").add_attribute(Attribute::Bold),
    ]);

    // When printin the accounts, I print the account's parent OUs' names.
    // To do that, I need to keep of how deep I am in the tree struture, and
    // what the parent OUs are.
    //
    // I use the a typical stack-based preorder traversal algorithm.
    // To keep track of the OUs,, I define what I can an height stack that is vec that stores
    // the degree of the node that we pop from the stack. Each time we push an OU, we push
    // degree of that OU to the height stack. When we pop any node, we decrement the
    // degree counter in the vec until it reaches 0, then we pop the stack. We pop
    // until there are no non-zero counters.
    let mut stack: Vec<&OrgTree> = Vec::new();
    let mut height_stack: Vec<usize> = Vec::new();
    let mut res: Vec<&OrgTree> = Vec::new();
    let mut ou_prefix: Vec<String> = Vec::new();

    stack.push(&org_tree);
    height_stack.push(1);
    while !stack.is_empty() {
        let node = stack.pop().unwrap();
        res.push(node);

        println!("{:#?}", ou_prefix);
        println!("{:#?}", height_stack);

        if let Some(last) = height_stack.last_mut() {
            *last -= 1;
        }

        match &node.root {
            OrgNode::Ou(ou) => {
                if let Some(ref children) = node.children {
                    ou_prefix.push(ou.name.as_ref().unwrap().to_string());
                    height_stack.push(children.len());
                    for elem in children {
                        stack.push(elem);
                    }
                }
            }
            OrgNode::Account(account) => {
                table.add_row(Row::from(vec![
                    account.name.as_ref().unwrap().to_string(),
                    account.id.as_ref().unwrap().to_string(),
                    build_ou_prefix(&ou_prefix),
                    account.email.as_ref().unwrap().to_string(),
                ]));
            }
        }

        while let Some(0) = height_stack.last() {
            height_stack.pop();
            ou_prefix.pop();
        }
    }

    table
}
```
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.