diff --git a/databuild/graph/analyze.rs b/databuild/graph/analyze.rs
index 3935ba8..527805b 100644
--- a/databuild/graph/analyze.rs
+++ b/databuild/graph/analyze.rs
@@ -569,8 +569,9 @@ async fn main() {
         .arg(
             Arg::new("partitions")
                 .help("Partition references to analyze")
-                .required(true)
-                .num_args(1..)
+                .required(false)
+                .num_args(0..)
+                .value_name("PARTITIONS")
         )
         .get_matches();
     
@@ -579,6 +580,23 @@ async fn main() {
         .cloned()
         .collect();
     
+    // Validate arguments based on mode
+    match mode.as_str() {
+        "plan" | "mermaid" => {
+            if args.is_empty() {
+                error!("Error: Partition references are required for {} mode", mode);
+                eprintln!("Error: Partition references are required for {} mode", mode);
+                exit(1);
+            }
+        }
+        "import_test" => {
+            // No partition arguments needed for test mode
+        }
+        _ => {
+            // Unknown mode, will be handled later
+        }
+    }
+    
     // Get build event log configuration from environment variables
     let build_event_log_uri = env::var("DATABUILD_BUILD_EVENT_LOG").ok();
     let build_request_id = env::var("DATABUILD_BUILD_REQUEST_ID")