Issue #144: Branching output of zenith branch

* Add ancestor_id to pg_list->branch_list output of pageserver.
* Display branching point (LSN) for each non-root branch.
* Add tests for `zenith branch`.

README.md

@@ -11,7 +11,7 @@ cd zenith
 make
 ```
 
-2. Start pageserver and postggres on top of it (should be called from repo root):
+2. Start pageserver and postgres on top of it (should be called from repo root):
 ```sh
 # Create ~/.zenith with proper paths to binaries and data
 # Later that would be responsibility of a package install script
@@ -30,7 +30,7 @@ Database initialized
 > ./target/debug/zenith pg start pg1
 
 # look up status and connection info
-> ./target/debug/zenith pg list     
+> ./target/debug/zenith pg list
 NODE		ADDRESS				STATUS
 pg1			127.0.0.1:55432		running
 ```

control_plane/src/storage.rs

@@ -167,7 +167,7 @@ impl PageServerNode {
 
     pub fn branches_list(&self) -> Result<Vec<BranchInfo>> {
         let mut client = self.page_server_psql_client()?;
-        let query_result = client.simple_query("pg_list")?;
+        let query_result = client.simple_query("branch_list")?;
         let branches_json = query_result
             .first()
             .map(|msg| match msg {

pageserver/src/branches.rs

@@ -29,6 +29,8 @@ pub struct BranchInfo {
     pub name: String,
     pub timeline_id: ZTimelineId,
     pub latest_valid_lsn: Option<Lsn>,
+    pub ancestor_id: Option<String>,
+    pub ancestor_lsn: Option<String>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -117,8 +119,10 @@ pub fn init_repo(conf: &PageServerConf, repo_dir: &Path) -> Result<()> {
 }
 
 pub(crate) fn get_branches(repository: &dyn Repository) -> Result<Vec<BranchInfo>> {
-    // adapted from CLI code
+    // Each branch has a corresponding record (text file) in the refs/branches
+    // with timeline_id.
     let branches_dir = std::path::Path::new("refs").join("branches");
+
     std::fs::read_dir(&branches_dir)?
         .map(|dir_entry_res| {
             let dir_entry = dir_entry_res?;
@@ -130,10 +134,26 @@ pub(crate) fn get_branches(repository: &dyn Repository) -> Result<Vec<BranchInfo
                 .map(|timeline| timeline.get_last_valid_lsn())
                 .ok();
 
+            let ancestor_path = std::path::Path::new("timelines")
+                .join(timeline_id.to_string())
+                .join("ancestor");
+            let mut ancestor_id: Option<String> = None;
+            let mut ancestor_lsn: Option<String> = None;
+
+            if ancestor_path.exists() {
+                let ancestor = std::fs::read_to_string(ancestor_path)?;
+                let mut strings = ancestor.split('@');
+
+                ancestor_id = Some(strings.next().unwrap().to_owned());
+                ancestor_lsn = Some(strings.next().unwrap().to_owned());
+            }
+
             Ok(BranchInfo {
                 name,
                 timeline_id,
                 latest_valid_lsn,
+                ancestor_id,
+                ancestor_lsn,
             })
         })
         .collect()
@@ -205,6 +225,8 @@ pub(crate) fn create_branch(
         name: branchname.to_string(),
         timeline_id: newtli,
         latest_valid_lsn: Some(startpoint.lsn),
+        ancestor_id: None,
+        ancestor_lsn: None,
     })
 }
 

pageserver/src/page_service.rs

@@ -762,7 +762,7 @@ impl Connection {
             self.write_message_noflush(&BeMessage::RowDescription)?;
             self.write_message_noflush(&BeMessage::DataRow(Bytes::from(branch)))?;
             self.write_message_noflush(&BeMessage::CommandComplete)?;
-        } else if query_string.starts_with(b"pg_list") {
+        } else if query_string.starts_with(b"branch_list") {
             let branches = crate::branches::get_branches(&*page_cache::get_repository())?;
             let branches_buf = serde_json::to_vec(&branches)?;
 

test_runner/batch_others/test_pageserver_api.py

@@ -13,37 +13,39 @@ def test_status(pageserver):
     assert cur.fetchone() == ('hello world',)
     pg_conn.close()
 
-def test_pg_list(pageserver, zenith_cli):
+def test_branch_list(pageserver, zenith_cli):
 
     # Create a branch for us
-    zenith_cli.run(["branch", "test_pg_list_main", "empty"]);
+    zenith_cli.run(["branch", "test_branch_list_main", "empty"]);
 
     page_server_conn = psycopg2.connect(pageserver.connstr())
     page_server_conn.autocommit = True
     page_server_cur = page_server_conn.cursor()
 
-    page_server_cur.execute('pg_list;')
+    page_server_cur.execute('branch_list;')
     branches = json.loads(page_server_cur.fetchone()[0])
     # Filter out branches created by other tests
-    branches = [x for x in branches if x['name'].startswith('test_pg_list')]
+    branches = [x for x in branches if x['name'].startswith('test_branch_list')]
 
     assert len(branches) == 1
-    assert branches[0]['name'] == 'test_pg_list_main'
+    assert branches[0]['name'] == 'test_branch_list_main'
     assert 'timeline_id' in branches[0]
     assert 'latest_valid_lsn' in branches[0]
+    assert 'ancestor_id' in branches[0]
+    assert 'ancestor_lsn' in branches[0]
 
     # Create another branch, and start Postgres on it
-    zenith_cli.run(['branch', 'test_pg_list_experimental', 'test_pg_list_main'])
-    zenith_cli.run(['pg', 'create', 'test_pg_list_experimental'])
+    zenith_cli.run(['branch', 'test_branch_list_experimental', 'test_branch_list_main'])
+    zenith_cli.run(['pg', 'create', 'test_branch_list_experimental'])
 
-    page_server_cur.execute('pg_list;')
+    page_server_cur.execute('branch_list;')
     new_branches = json.loads(page_server_cur.fetchone()[0])
     # Filter out branches created by other tests
-    new_branches = [x for x in new_branches if x['name'].startswith('test_pg_list')]
+    new_branches = [x for x in new_branches if x['name'].startswith('test_branch_list')]
     assert len(new_branches) == 2
     new_branches.sort(key=lambda k: k['name'])
 
-    assert new_branches[0]['name'] == 'test_pg_list_experimental'
+    assert new_branches[0]['name'] == 'test_branch_list_experimental'
     assert new_branches[0]['timeline_id'] != branches[0]['timeline_id']
 
     # TODO: do the LSNs have to match here?

test_runner/batch_others/test_zenith_cli.py

@@ -0,0 +1,49 @@
+import pytest
+import psycopg2
+import json
+
+pytest_plugins = ("fixtures.zenith_fixtures")
+
+def helper_compare_branch_list(page_server_cur, zenith_cli):
+    """
+    Compare branches list returned by CLI and directly via API.
+    Filters out branches created by other tests.
+    """
+
+    page_server_cur.execute('branch_list;')
+    branches_api = sorted(map(lambda b: b['name'], json.loads(page_server_cur.fetchone()[0])))
+    branches_api = [b for b in branches_api if b.startswith('test_cli_') or b in ('empty', 'main')]
+
+    res = zenith_cli.run(["branch"]);
+    assert(res.stderr == '')
+    branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
+    branches_cli = [b for b in branches_cli if b.startswith('test_cli_') or b in ('empty', 'main')]
+
+    assert(branches_api == branches_cli)
+
+def test_cli_branch_list(pageserver, zenith_cli):
+
+    page_server_conn = psycopg2.connect(pageserver.connstr())
+    page_server_conn.autocommit = True
+    page_server_cur = page_server_conn.cursor()
+
+    # Initial sanity check
+    helper_compare_branch_list(page_server_cur, zenith_cli)
+
+    # Create a branch for us
+    res = zenith_cli.run(["branch", "test_cli_branch_list_main", "main"]);
+    assert(res.stderr == '')
+    helper_compare_branch_list(page_server_cur, zenith_cli)
+
+    # Create a nested branch
+    res = zenith_cli.run(["branch", "test_cli_branch_list_nested", "test_cli_branch_list_main"]);
+    assert(res.stderr == '')
+    helper_compare_branch_list(page_server_cur, zenith_cli)
+
+    # Check that all new branches are visible via CLI
+    res = zenith_cli.run(["branch"]);
+    assert(res.stderr == '')
+    branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
+
+    assert('test_cli_branch_list_main' in branches_cli)
+    assert('test_cli_branch_list_nested' in branches_cli)

test_runner/fixtures/zenith_fixtures.py

@@ -84,11 +84,18 @@ class ZenithCli:
         """ Run "zenith" with the specified arguments.
 
         arguments must be in list form, e.g. ['pg', 'create']
+
+        Return both stdout and stderr, which can be accessed as
+
+        result = zenith_cli.run(...)
+        assert(result.stderr == "")
+        print(result.stdout)
+
         """
         assert type(arguments) == list
         args = [self.bin_zenith] + arguments
         print('Running command "{}"'.format(' '.join(args)))
-        subprocess.run(args, env=self.env, check=True)
+        return subprocess.run(args, env=self.env, check=True, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
 
 @zenfixture

zenith/src/main.rs

@@ -12,6 +12,11 @@ use std::str::FromStr;
 use pageserver::{branches::BranchInfo, ZTimelineId};
 use zenith_utils::lsn::Lsn;
 
+struct BranchTreeEl {
+    pub info: BranchInfo,
+    pub children: Vec<String>,
+}
+
 // Main entry point for the 'zenith' CLI utility
 //
 // This utility helps to manage zenith installation. That includes following:
@@ -159,6 +164,80 @@ fn main() -> Result<()> {
     Ok(())
 }
 
+// Print branches list as a tree-like structure.
+fn print_branches_tree(branches: Vec<BranchInfo>) {
+    let mut branches_hash: HashMap<String, BranchTreeEl> = HashMap::new();
+
+    // Form a hash table of branch timeline_id -> BranchTreeEl.
+    for branch in &branches {
+        branches_hash.insert(
+            branch.timeline_id.to_string(),
+            BranchTreeEl {
+                info: branch.clone(),
+                children: Vec::new(),
+            },
+        );
+    }
+
+    // Memorize all direct children of each branch.
+    for branch in &branches {
+        if let Some(name) = &branch.ancestor_id {
+            branches_hash
+                .get_mut(name)
+                .unwrap()
+                .children
+                .push(branch.timeline_id.to_string());
+        }
+    }
+
+    // Sort children by name to bring some order.
+    for (_tid, branch) in &mut branches_hash {
+        branch.children.sort();
+    }
+
+    for (_tid, branch) in &branches_hash {
+        // Start with root branches (no ancestors) first.
+        // Now it is 'main' branch only, but things may change.
+        if branch.info.ancestor_id.is_none() {
+            print_branch(0, 0, branch, &branches_hash);
+        }
+    }
+}
+
+// Recursively print branch info with all its children.
+fn print_branch(
+    nesting_level: usize,
+    padding: usize,
+    branch: &BranchTreeEl,
+    branches: &HashMap<String, BranchTreeEl>,
+) {
+    let new_padding: usize;
+
+    if nesting_level > 0 {
+        // Six extra chars for graphics, spaces and so on in addition to LSN length.
+        new_padding = padding + 6 + branch.info.ancestor_lsn.as_ref().unwrap().chars().count();
+
+        print!(
+            "{}└─ @{}:",
+            " ".repeat(padding),
+            branch.info.ancestor_lsn.as_ref().unwrap()
+        );
+    } else {
+        new_padding = 1;
+    }
+
+    print!(" {}\n", branch.info.name);
+
+    for child in &branch.children {
+        print_branch(
+            nesting_level + 1,
+            new_padding,
+            branches.get(child).unwrap(),
+            branches,
+        );
+    }
+}
+
 /// Returns a map of timeline IDs to branch_name@lsn strings.
 /// Connects to the pageserver to query this information.
 fn get_branch_infos(env: &local_env::LocalEnv) -> Result<HashMap<ZTimelineId, BranchInfo>> {
@@ -188,9 +267,8 @@ fn handle_branch(branch_match: &ArgMatches, env: &local_env::LocalEnv) -> Result
         }
     } else {
         // No arguments, list branches
-        for branch in pageserver.branches_list()? {
-            println!(" {}", branch.name);
-        }
+        let branches = pageserver.branches_list().unwrap();
+        print_branches_tree(branches);
     }
 
     Ok(())