Introduce cluster support (#299)

Author: farost

BUILD

@@ -29,6 +29,7 @@ rust_binary(
         "@crates//:futures",
         "@crates//:glob",
         "@crates//:home",
+        "@crates//:itertools",
         "@crates//:rpassword",
         "@crates//:rustyline",
         "@crates//:sentry",

WORKSPACE

@@ -75,8 +75,8 @@ rust_register_toolchains(
         "x86_64-pc-windows-msvc",
         "x86_64-unknown-linux-gnu",
     ],
-    rust_analyzer_version = "1.81.0",
-    versions = ["1.81.0"],
+    rust_analyzer_version = "1.84.0",
+    versions = ["1.84.0"],
 )
 
 rust_analyzer_toolchain_tools_repository(

dependencies/typedb/repositories.bzl

@@ -5,17 +5,19 @@
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 def typedb_dependencies():
-    git_repository(
-        name = "typedb_dependencies",
-        remote = "https://github.com/typedb/typedb-dependencies",
-        commit = "f6e710f9857b1c30ad1764c1c41afce4e4e02981", # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_dependencies
-    )
+     # TODO: Return ref after merge to master
+     git_repository(
+         name = "typedb_dependencies",
+         remote = "https://github.com/typedb/typedb-dependencies",
+         commit = "19a70bcad19b9a28814016f183ac3e3a23c1ff0d",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_dependencies
+     )
 
 def typedb_driver():
+     # TODO: Return ref after merge to master
     git_repository(
         name = "typedb_driver",
         remote = "https://github.com/typedb/typedb-driver",
-        tag = "3.7.0",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_driver
+        commit = "a804da767e68154cc3ee5a7531feeaba5dbf0c17",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_driver
     )
 
 def typeql():

src/cli.rs

@@ -24,10 +24,28 @@ pub struct Args {
     #[arg(long, value_name = "path to script file")]
     pub script: Vec<String>,
 
-    /// TypeDB address to connect to. If using TLS encryption, this must start with "https://"
-    #[arg(long, value_name = ADDRESS_VALUE_NAME)]
+    /// TypeDB address to connect to (host:port). If using TLS encryption, this must start with "https://".
+    #[arg(long, value_name = ADDRESS_VALUE_NAME, conflicts_with_all = ["addresses", "address_translation"])]
     pub address: Option<String>,
 
+    /// A comma-separated list of TypeDB replica addresses of a single cluster to connect to.
+    #[arg(long, value_name = "host1:port1,host2:port2", conflicts_with_all = ["address", "address_translation"])]
+    pub addresses: Option<String>,
+
+    /// A comma-separated list of public=private address pairs. Public addresses are the user-facing
+    /// addresses of the replicas, and private addresses are the addresses used for the server-side
+    /// connection between replicas.
+    #[arg(long, value_name = "public=private,...", conflicts_with_all = ["address", "addresses"])]
+    pub address_translation: Option<String>,
+
+    /// If used in a Cluster environment (Cloud or Enterprise), disables attempts to redirect 
+    /// requests to server replicas, limiting Console to communicate only with the single address 
+    /// specified in the `address` argument.
+    /// Use for administrative / debug purposes to test a specific replica only: this option will
+    /// lower the success rate of Console's operations in production.
+    #[arg(long = "replication-disabled", default_value = "false")]
+    pub replication_disabled: bool,
+
     /// Username for authentication
     #[arg(long, value_name = USERNAME_VALUE_NAME)]
     pub username: Option<String>,
@@ -48,8 +66,8 @@ pub struct Args {
 
     /// Disable error reporting. Error reporting helps TypeDB improve by reporting
     /// errors and crashes to the development team.
-    #[arg(long = "diagnostics-disable", default_value = "false")]
-    pub diagnostics_disable: bool,
+    #[arg(long = "diagnostics-disabled", default_value = "false")]
+    pub diagnostics_disabled: bool,
 
     /// Print the Console binary version
     #[arg(long = "version")]

src/completion_cache.rs

@@ -0,0 +1,33 @@
+use std::collections::HashSet;
+use std::hash::Hash;
+use std::sync::{Arc, RwLock};
+
+#[derive(Clone)]
+pub struct CompletionCache<T> {
+    inner: Arc<RwLock<HashSet<T>>>,
+}
+
+impl<T> CompletionCache<T>
+where
+    T: Eq + Hash + Clone,
+{
+    pub fn new() -> Self {
+        Self { inner: Arc::new(RwLock::new(HashSet::new())) }
+    }
+
+    pub fn replace_all<I>(&self, items: impl IntoIterator<Item = T>) {
+        *self.inner.write().unwrap() = items.into_iter().collect();
+    }
+
+    pub fn snapshot(&self) -> Vec<T> {
+        self.inner.read().unwrap().iter().cloned().collect()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.inner.read().unwrap().is_empty()
+    }
+
+    pub fn clear(&self) {
+        self.inner.write().unwrap().clear();
+    }
+}

src/main.rs

@@ -5,6 +5,7 @@
  */
 
 use std::{
+    collections::HashMap,
     env,
     env::temp_dir,
     error::Error,
@@ -22,15 +23,16 @@ use clap::Parser;
 use home::home_dir;
 use rustyline::error::ReadlineError;
 use sentry::ClientOptions;
-use typedb_driver::{Credentials, DriverOptions, Transaction, TransactionType, TypeDBDriver};
+use typedb_driver::{Addresses, Credentials, DriverOptions, Transaction, TransactionType, TypeDBDriver};
 
 use crate::{
     cli::{Args, ADDRESS_VALUE_NAME, USERNAME_VALUE_NAME},
     completions::{database_name_completer_fn, file_completer},
     operations::{
         database_create, database_create_init, database_delete, database_export, database_import, database_list,
-        database_schema, transaction_close, transaction_commit, transaction_query, transaction_read,
-        transaction_rollback, transaction_schema, transaction_source, transaction_write, user_create, user_delete,
+        database_schema, replica_deregister, replica_list, replica_primary, replica_register, server_version,
+        transaction_close, transaction_commit, transaction_query, transaction_read, transaction_rollback,
+        transaction_schema, transaction_source, transaction_write, user_create, user_delete,
         user_list, user_update_password,
     },
     repl::{
@@ -126,13 +128,19 @@ fn main() {
         println!("{}", VERSION);
         exit(ExitCode::Success as i32);
     }
-    let address = match args.address {
-        Some(address) => address,
-        None => {
-            println_error!("missing server address ('{}').", format_argument!("--address <{ADDRESS_VALUE_NAME}>"));
-            exit(ExitCode::UserInputError as i32);
-        }
-    };
+    let address_info = parse_addresses(&args);
+    if !args.tls_disabled && !address_info.only_https {
+        println_error!(
+            "\
+            TLS connections can only be enabled when connecting to HTTPS endpoints. \
+            For example, using 'https://<ip>:port'.\n\
+            Please modify the address or disable TLS ('{}'). {}\
+        ",
+            format_argument!("--tls-disabled"),
+            format_warning!("WARNING: this will send passwords over plaintext!"),
+        );
+        exit(ExitCode::UserInputError as i32);
+    }
     let username = match args.username {
         Some(username) => username,
         None => {
@@ -146,34 +154,28 @@ fn main() {
     if args.password.is_none() {
         args.password = Some(LineReaderHidden::new().readline(&format!("password for '{username}': ")));
     }
-    if !args.diagnostics_disable {
+    if !args.diagnostics_disabled {
         init_diagnostics()
     }
-    if !args.tls_disabled && !address.starts_with("https:") {
-        println_error!(
-            "\
-            TLS connections can only be enabled when connecting to HTTPS endpoints. \
-            For example, using 'https://<ip>:port'.\n\
-            Please modify the address or disable TLS ('{}'). {}\
-        ",
-            format_argument!("--tls-disabled"),
-            format_warning!("WARNING: this will send passwords over plaintext!"),
-        );
-        exit(ExitCode::UserInputError as i32);
-    }
     let tls_root_ca_path = args.tls_root_ca.as_ref().map(|value| Path::new(value));
-
     let runtime = BackgroundRuntime::new();
+    let driver_options = DriverOptions::new()
+        .use_replication(!args.replication_disabled)
+        .is_tls_enabled(!args.tls_disabled)
+        .tls_root_ca(tls_root_ca_path)
+        .unwrap();
     let driver = match runtime.run(TypeDBDriver::new(
-        address,
+        address_info.addresses,
         Credentials::new(&username, args.password.as_ref().unwrap()),
-        DriverOptions::new(!args.tls_disabled, tls_root_ca_path).unwrap(),
+        driver_options,
     )) {
         Ok(driver) => Arc::new(driver),
         Err(err) => {
             let tls_error =
                 if args.tls_disabled { "" } else { "\nVerify that the server is also configured with TLS encryption." };
-            println_error!("Failed to create driver connection to server. {err}{tls_error}");
+            let replication_error =
+                if args.replication_disabled { "\nVerify that the connection address is **exactly** the same as the server address specified in its config." } else { "" };
+            println_error!("Failed to create driver connection to server. {err}{tls_error}{replication_error}");
             exit(ExitCode::ConnectionError as i32);
         }
     };
@@ -332,6 +334,28 @@ fn execute_commands(context: &mut ConsoleContext, mut input: &str, must_log_comm
 }
 
 fn entry_repl(driver: Arc<TypeDBDriver>, runtime: BackgroundRuntime) -> Repl<ConsoleContext> {
+    let server_commands =
+        Subcommand::new("server").add(CommandLeaf::new("version", "Retrieve server version.", server_version));
+
+    let replica_commands = Subcommand::new("replica")
+        .add(CommandLeaf::new("list", "List replicas.", replica_list))
+        .add(CommandLeaf::new("primary", "Get current primary replica.", replica_primary))
+        .add(CommandLeaf::new_with_inputs(
+            "register",
+            "Register new replica. Requires a clustering address, not a connection address.",
+            vec![
+                CommandInput::new_required("replica id", get_word, None),
+                CommandInput::new_required("clustering address", get_word, None),
+            ],
+            replica_register,
+        ))
+        .add(CommandLeaf::new_with_input(
+            "deregister",
+            "Deregister existing replica.",
+            CommandInput::new_required("replica id", get_word, None),
+            replica_deregister,
+        ));
+
     let database_commands = Subcommand::new("database")
         .add(CommandLeaf::new("list", "List databases on the server.", database_list))
         .add(CommandLeaf::new_with_input(
@@ -451,8 +475,10 @@ fn entry_repl(driver: Arc<TypeDBDriver>, runtime: BackgroundRuntime) -> Repl<Con
     let history_path = home_dir().unwrap_or_else(|| temp_dir()).join(ENTRY_REPL_HISTORY);
 
     let repl = Repl::new(PROMPT.to_owned(), history_path, false, None)
+        .add(server_commands)
         .add(database_commands)
         .add(user_commands)
+        .add(replica_commands)
         .add(transaction_commands);
 
     repl
@@ -508,6 +534,42 @@ fn transaction_type_str(transaction_type: TransactionType) -> &'static str {
     }
 }
 
+struct AddressInfo {
+    only_https: bool,
+    addresses: Addresses,
+}
+
+fn parse_addresses(args: &Args) -> AddressInfo {
+    if let Some(address) = &args.address {
+        AddressInfo {
+            only_https: is_https_address(address),
+            addresses: Addresses::try_from_address_str(address).unwrap(),
+        }
+    } else if let Some(addresses) = &args.addresses {
+        let split = addresses.split(',').map(str::to_string).collect::<Vec<_>>();
+        let only_https = split.iter().all(|address| is_https_address(address));
+        AddressInfo { only_https, addresses: Addresses::try_from_addresses_str(split).unwrap() }
+    } else if let Some(translation) = &args.address_translation {
+        let mut map = HashMap::new();
+        let mut only_https = true;
+        for pair in translation.split(',') {
+            let (public_address, private_address) = pair
+                .split_once('=')
+                .unwrap_or_else(|| panic!("Invalid address pair: {pair}. Must be of form public=private"));
+            only_https = only_https && is_https_address(public_address);
+            map.insert(public_address.to_string(), private_address.to_string());
+        }
+        println!("Translation map:: {map:?}"); // TODO: Remove
+        AddressInfo { only_https, addresses: Addresses::try_from_translation_str(map).unwrap() }
+    } else {
+        panic!("At least one of --address, --addresses, or --address-translation must be provided.");
+    }
+}
+
+fn is_https_address(address: &str) -> bool {
+    address.starts_with("https:")
+}
+
 fn init_diagnostics() {
     let _ = sentry::init((
         DIAGNOSTICS_REPORTING_URI,