Add username/password for Mongo login

Author: palfrey

.github/workflows/native-bazel.yaml

@@ -59,3 +59,5 @@ jobs:
             exit 1
           fi
         shell: bash
+        env:
+          NATIVELINK_TEST_MONGO_URL: ${{ secrets.NATIVELINK_TEST_MONGO_URL}}

.github/workflows/native-cargo.yaml

@@ -46,3 +46,5 @@ jobs:
 
       - name: Test on ${{ runner.os }}
         run: cargo test --all --profile=smol
+        env:
+          NATIVELINK_TEST_MONGO_URL: ${{ secrets.NATIVELINK_TEST_MONGO_URL}}

.github/workflows/nix.yaml

@@ -70,6 +70,8 @@ jobs:
         run: >
           nix develop --impure --command
           bash -c "cargo test --all --profile=smol"
+        env:
+          NATIVELINK_TEST_MONGO_URL: ${{ secrets.NATIVELINK_TEST_MONGO_URL  }}
 
   installation:
     strategy:

.github/workflows/sanitizers.yaml

@@ -44,3 +44,5 @@ jobs:
       - name: Run Bazel tests
         run: bazel test --config=${{ matrix.sanitizer }} --verbose_failures //...
         shell: bash
+        env:
+          NATIVELINK_TEST_MONGO_URL: ${{ secrets.NATIVELINK_TEST_MONGO_URL}}

nativelink-config/src/stores.rs

@@ -1174,6 +1174,16 @@ pub struct ExperimentalMongoSpec {
     #[serde(deserialize_with = "convert_string_with_shellexpand")]
     pub connection_string: String,
 
+    // Username to use for connection
+    // Default: not set, or use the one from connection_string
+    #[serde(default)]
+    pub username: String,
+
+    // Password to use for connection
+    // Default: not set, or use the one from connection_string
+    #[serde(default)]
+    pub password: String,
+
     /// The database name to use.
     /// Default: "nativelink"
     #[serde(default, deserialize_with = "convert_string_with_shellexpand")]

nativelink-store/BUILD.bazel

@@ -110,6 +110,11 @@ rust_test_suite(
         "tests/size_partitioning_store_test.rs",
         "tests/verify_store_test.rs",
     ],
+    env_inherit = [
+        "NATIVELINK_TEST_MONGO_URL",
+        "CI",
+        "GITHUB_ACTIONS",
+    ],
     proc_macro_deps = [
         "//nativelink-macro",
         "@crates//:async-trait",

nativelink-store/src/mongo_store.rs

@@ -23,7 +23,9 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use futures::stream::{Stream, StreamExt, TryStreamExt};
 use mongodb::bson::{Bson, Document, doc};
-use mongodb::options::{ClientOptions, FindOptions, IndexOptions, ReturnDocument, WriteConcern};
+use mongodb::options::{
+    ClientOptions, Credential, FindOptions, IndexOptions, ReturnDocument, WriteConcern,
+};
 use mongodb::{Client as MongoClient, Collection, Database, IndexModel};
 use nativelink_config::stores::ExperimentalMongoSpec;
 use nativelink_error::{Code, Error, make_err, make_input_err};
@@ -163,6 +165,21 @@ impl ExperimentalMongoStore {
                 )
             })?;
 
+        if !spec.username.is_empty() || !spec.password.is_empty() {
+            if client_options.credential.is_some() {
+                return Err(make_err!(
+                    Code::InvalidArgument,
+                    "Can't set credentials in both the connection_string and username/password"
+                ));
+            }
+            client_options.credential = Some(
+                Credential::builder()
+                    .username(spec.username)
+                    .password(spec.password)
+                    .build(),
+            );
+        }
+
         client_options.server_selection_timeout =
             Some(Duration::from_millis(spec.connection_timeout_ms));
         client_options.connect_timeout = Some(Duration::from_millis(spec.connection_timeout_ms));

nativelink-store/tests/mongo_store_test.rs

@@ -53,6 +53,8 @@ fn create_test_spec_with_key_prefix(key_prefix: Option<String>) -> ExperimentalM
 
     ExperimentalMongoSpec {
         connection_string: std::env::var("NATIVELINK_TEST_MONGO_URL").unwrap_or(default_connection),
+        username: String::new(),
+        password: String::new(),
         database: format!(
             "nltest_{}_{}",
             timestamp,
@@ -83,11 +85,13 @@ pub struct TestMongoHelper {
     pub database_name: String,
 }
 
+fn in_ci() -> bool {
+    std::env::var("CI").is_ok() || std::env::var("GITHUB_ACTIONS").is_ok()
+}
+
 fn non_ci_test_store_access() {
     // Check if this is a local development environment without credentials
-    let is_ci = std::env::var("CI").is_ok() || std::env::var("GITHUB_ACTIONS").is_ok();
-
-    if !is_ci {
+    if !in_ci() {
         eprintln!("\n🔒 MongoDB tests require access to the NativeLink test database.");
         eprintln!("   For local development access, please email: [email protected]");
         eprintln!("   ");
@@ -177,18 +181,68 @@ async fn create_test_store() -> Result<Arc<ExperimentalMongoStore>, Error> {
     ExperimentalMongoStore::new(spec).await
 }
 
+enum MongoDBState {
+    Ok(TestMongoHelper),
+    Broken(Error),
+    Skipping,
+}
+
 /// Utility to check if `MongoDB` is available for testing.
 /// Returns an error that can be used to skip tests when `MongoDB` is not available.
-async fn check_mongodb_available() -> Result<(), Error> {
-    TestMongoHelper::new_or_skip().await.map(|_| ())
+async fn check_mongodb_available() -> MongoDBState {
+    match TestMongoHelper::new_or_skip().await {
+        Ok(helper) => MongoDBState::Ok(helper),
+        Err(err) if err.code == Code::Unavailable && !in_ci() => {
+            eprintln!("Skipping MongoDB test - MongoDB not available");
+            MongoDBState::Skipping
+        }
+        Err(err) => {
+            eprintln!("Mongo test error: {err}");
+            MongoDBState::Broken(err)
+        }
+    }
+}
+
+#[nativelink_test]
+async fn connect_with_username_and_password() -> Result<(), Error> {
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
+    }
+
+    let mut spec = create_test_spec();
+    // These are incorrect, but should allow us to connect and fail with bad login
+    // v.s. "cannot do login at all" or "config failure"
+    spec.username = "foo".to_string();
+    spec.password = "bar".to_string();
+    let store = ExperimentalMongoStore::new(spec)
+        .await
+        .expect("Working store");
+    let digest = DigestInfo::try_new(VALID_HASH1, 2)?;
+    let result = store.has(digest).await?;
+    assert!(
+        result.is_some(),
+        "Expected mongo store to have hash: {VALID_HASH1}",
+    );
+    Ok(())
 }
 
 #[nativelink_test]
 async fn upload_and_get_data() -> Result<(), Error> {
     // Create test helper with automatic cleanup
-    let Ok(helper) = TestMongoHelper::new_or_skip().await else {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    let helper = match check_mongodb_available().await {
+        MongoDBState::Ok(helper) => helper,
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     };
 
     // Construct the data we want to send.
@@ -252,9 +306,14 @@ async fn upload_and_get_data_without_prefix() -> Result<(), Error> {
 #[nativelink_test]
 async fn upload_empty_data() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     let data = Bytes::from_static(b"");
@@ -277,9 +336,14 @@ async fn upload_empty_data() -> Result<(), Error> {
 #[allow(clippy::items_after_statements)]
 async fn test_large_downloads_are_chunked() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     const READ_CHUNK_SIZE: usize = 1024;
@@ -316,9 +380,14 @@ async fn test_large_downloads_are_chunked() -> Result<(), Error> {
 #[allow(clippy::items_after_statements)]
 async fn yield_between_sending_packets_in_update() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     const READ_CHUNK_SIZE: usize = 1024;
@@ -372,9 +441,14 @@ async fn yield_between_sending_packets_in_update() -> Result<(), Error> {
 #[nativelink_test]
 async fn zero_len_items_exist_check() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     let digest = DigestInfo::try_new(VALID_HASH1, 0)?;
@@ -427,9 +501,14 @@ async fn test_empty_connection_string() -> Result<(), Error> {
 #[nativelink_test]
 async fn test_default_values() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     let mut spec = create_test_spec();
@@ -464,9 +543,14 @@ async fn test_default_values() -> Result<(), Error> {
 #[nativelink_test]
 async fn dont_loop_forever_on_empty() -> Result<(), Error> {
     // Skip test if MongoDB is not available
-    if check_mongodb_available().await.is_err() {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    match check_mongodb_available().await {
+        MongoDBState::Ok(_helper) => {}
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     }
 
     let store = create_test_store().await?;
@@ -502,9 +586,14 @@ async fn dont_loop_forever_on_empty() -> Result<(), Error> {
 #[nativelink_test]
 async fn test_partial_reads() -> Result<(), Error> {
     // Create test helper with automatic cleanup
-    let Ok(helper) = TestMongoHelper::new_or_skip().await else {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    let helper = match check_mongodb_available().await {
+        MongoDBState::Ok(helper) => helper,
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     };
 
     let data = Bytes::from_static(b"Hello, MongoDB World!");
@@ -620,9 +709,14 @@ async fn test_database_lifecycle() -> Result<(), Error> {
 #[allow(clippy::use_debug)]
 async fn create_ten_cas_entries() -> Result<(), Error> {
     // Create test helper with automatic cleanup
-    let Ok(helper) = TestMongoHelper::new_or_skip().await else {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    let helper = match check_mongodb_available().await {
+        MongoDBState::Ok(helper) => helper,
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     };
 
     eprintln!(
@@ -848,9 +942,14 @@ impl SchedulerStoreDecodeTo for TestSchedulerKey {
 #[nativelink_test]
 async fn test_scheduler_store_operations() -> Result<(), Error> {
     // Create test helper
-    let Ok(helper) = TestMongoHelper::new_or_skip().await else {
-        eprintln!("Skipping MongoDB test - MongoDB not available");
-        return Ok(());
+    let helper = match check_mongodb_available().await {
+        MongoDBState::Ok(helper) => helper,
+        MongoDBState::Broken(err) => {
+            return Err(err);
+        }
+        MongoDBState::Skipping => {
+            return Ok(());
+        }
     };
 
     eprintln!(
@@ -1062,6 +1161,8 @@ async fn test_non_w_config() -> Result<(), Error> {
         Error::new(Code::InvalidArgument, "write_concern_w not set, but j and/or timeout set. Please set 'write_concern_w' to a non-default value. See https://www.mongodb.com/docs/manual/reference/write-concern/#w-option for options.".to_string()),
         ExperimentalMongoStore::new(ExperimentalMongoSpec {
             connection_string: "mongodb://dummy".to_string(),
+            username: String::new(),
+            password: String::new(),
             database: "dummy".to_string(),
             cas_collection: "test_cas".to_string(),
             scheduler_collection: "test_scheduler".to_string(),