override NO_COLOR env

Author: Shubham8287

crates/schema/src/auto_migrate.rs

@@ -28,6 +28,12 @@ pub enum MigratePlan<'def> {
     Auto(AutoMigratePlan<'def>),
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub enum PrettyPrintStyle {
+    AnsiColor,
+    NoColor,
+}
+
 impl<'def> MigratePlan<'def> {
     /// Get the old `ModuleDef` for this migration plan.
     pub fn old_def(&self) -> &'def ModuleDef {
@@ -45,25 +51,25 @@ impl<'def> MigratePlan<'def> {
         }
     }
 
-    pub fn pretty_print(&self, no_color: bool) -> anyhow::Result<String> {
+    pub fn pretty_print(&self, style: PrettyPrintStyle) -> anyhow::Result<String> {
+        use PrettyPrintStyle::*;
+
         match self {
             MigratePlan::Manual(_) => {
                 anyhow::bail!("Manual migration plans are not yet supported for pretty printing.")
             }
-            MigratePlan::Auto(plan) => {
-                if no_color {
+
+            MigratePlan::Auto(plan) => match style {
+                NoColor => {
                     let mut fmt = plain_formatter::PlainFormatter::new(1024);
-                    format_plan(&mut fmt, plan)
-                        .map_err(|e| anyhow::anyhow!("Failed to format migration plan: {e}"))
-                        .map(|_| fmt.to_string())
-                } else {
-                    // Use the ANSI formatter with colors.
+                    format_plan(&mut fmt, plan).map(|_| fmt.to_string())
+                }
+                AnsiColor => {
                     let mut fmt = AnsiFormatter::new(1024, ColorScheme::default());
-                    format_plan(&mut fmt, plan)
-                        .map_err(|e| anyhow::anyhow!("Failed to format migration plan: {e}"))
-                        .map(|_| fmt.to_string())
+                    format_plan(&mut fmt, plan).map(|_| fmt.to_string())
                 }
             }
+            .map_err(|e| anyhow::anyhow!("Failed to format migration plan: {e}")),
         }
     }
 }
@@ -1433,7 +1439,8 @@ mod tests {
 
         insta::assert_snapshot!(
             "empty_to_populated_migration",
-            plan.pretty_print(false).expect("should pretty print")
+            plan.pretty_print(PrettyPrintStyle::AnsiColor)
+                .expect("should pretty print")
         );
     }
 
@@ -1445,7 +1452,8 @@ mod tests {
 
         insta::assert_snapshot!(
             "updated pretty print",
-            plan.pretty_print(false).expect("should pretty print")
+            plan.pretty_print(PrettyPrintStyle::AnsiColor)
+                .expect("should pretty print")
         );
     }
 
@@ -1457,7 +1465,8 @@ mod tests {
 
         insta::assert_snapshot!(
             "updated pretty print no color",
-            plan.pretty_print(true).expect("should pretty print")
+            plan.pretty_print(PrettyPrintStyle::NoColor)
+                .expect("should pretty print")
         );
     }
 }

crates/schema/src/auto_migrate/ansi_formatter.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 
-use colored::{Color, Colorize as _};
+use colored::{control::set_override, Color, Colorize as _};
 use spacetimedb_lib::{db::raw_def::v9::TableAccess, AlgebraicType};
 use spacetimedb_sats::algebraic_type::fmt::fmt_algebraic_type;
 
@@ -55,6 +55,8 @@ impl fmt::Display for AnsiFormatter {
 impl AnsiFormatter {
     /// Create a new formatter with custom colors
     pub fn new(cap: usize, colors: ColorScheme) -> Self {
+        // This overrides `NO_COLOR` as `ANSIFormatter` should always be colored.
+        set_override(true);
         Self {
             buffer: String::with_capacity(cap),
             colors,