diff --git a/README.md b/README.md index 80ea755778..65b290918d 100644 --- a/README.md +++ b/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 ``` diff --git a/control_plane/src/storage.rs b/control_plane/src/storage.rs index 92e9eafaba..ff85e67e80 100644 --- a/control_plane/src/storage.rs +++ b/control_plane/src/storage.rs @@ -167,7 +167,7 @@ impl PageServerNode { pub fn branches_list(&self) -> Result> { 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 { diff --git a/pageserver/src/branches.rs b/pageserver/src/branches.rs index 00702ed7fe..aed8c33c9f 100644 --- a/pageserver/src/branches.rs +++ b/pageserver/src/branches.rs @@ -29,6 +29,8 @@ pub struct BranchInfo { pub name: String, pub timeline_id: ZTimelineId, pub latest_valid_lsn: Option, + pub ancestor_id: Option, + pub ancestor_lsn: Option, } #[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> { - // 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 = None; + let mut ancestor_lsn: Option = 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, }) } diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index 491bf1f5d3..8bd52eb553 100644 --- a/pageserver/src/page_service.rs +++ b/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)?; diff --git a/test_runner/batch_others/test_pageserver_api.py b/test_runner/batch_others/test_pageserver_api.py index 5ea65fae51..af1c6723db 100644 --- a/test_runner/batch_others/test_pageserver_api.py +++ b/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? diff --git a/test_runner/batch_others/test_zenith_cli.py b/test_runner/batch_others/test_zenith_cli.py new file mode 100644 index 0000000000..510205c003 --- /dev/null +++ b/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) diff --git a/test_runner/fixtures/zenith_fixtures.py b/test_runner/fixtures/zenith_fixtures.py index f3ecfa4eb2..cfba28d03c 100644 --- a/test_runner/fixtures/zenith_fixtures.py +++ b/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 diff --git a/zenith/src/main.rs b/zenith/src/main.rs index efcbca25b7..4622aedf5b 100644 --- a/zenith/src/main.rs +++ b/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, +} + // 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) { + let mut branches_hash: HashMap = 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, +) { + 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> { @@ -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(())