From e21e16249ffa23014574e58c5cea27eb25264c0a Mon Sep 17 00:00:00 2001
From: Gordy French <gordyf@meta.com>
Date: Wed, 16 Oct 2024 10:09:38 -0700
Subject: [PATCH] Allow codemod to specify rollout percentage

Reviewed By: captbaritone

Differential Revision: D64425741

fbshipit-source-id: 59f78cde8622d9f4c0be34e570844d5ee652073f
---
 compiler/crates/common/src/rollout.rs        |  2 +-
 compiler/crates/relay-bin/src/main.rs        |  2 +-
 compiler/crates/relay-codemod/src/codemod.rs | 36 ++++++++++++++++----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/compiler/crates/common/src/rollout.rs b/compiler/crates/common/src/rollout.rs
index 278e2413ff86f..92afe103cd1da 100644
--- a/compiler/crates/common/src/rollout.rs
+++ b/compiler/crates/common/src/rollout.rs
@@ -15,7 +15,7 @@ use serde::Serialize;
 /// Can be constructed as the Default which passes or a percentage between 0 and
 /// 100.
 #[derive(Default, Debug, Serialize, Deserialize, Clone, Copy, JsonSchema)]
-pub struct Rollout(Option<u8>);
+pub struct Rollout(pub Option<u8>);
 
 impl Rollout {
     /// Checks some key deterministically and passes on average the given
diff --git a/compiler/crates/relay-bin/src/main.rs b/compiler/crates/relay-bin/src/main.rs
index 51a2ebb864ff3..27cf59245a886 100644
--- a/compiler/crates/relay-bin/src/main.rs
+++ b/compiler/crates/relay-bin/src/main.rs
@@ -82,7 +82,7 @@ struct CodemodCommand {
     config: Option<PathBuf>,
 
     /// The name of the codemod to run
-    #[clap(long, short, value_enum)]
+    #[clap(subcommand)]
     codemod: AvailableCodemod,
 }
 
diff --git a/compiler/crates/relay-codemod/src/codemod.rs b/compiler/crates/relay-codemod/src/codemod.rs
index a701bff3f177c..5b57869da10da 100644
--- a/compiler/crates/relay-codemod/src/codemod.rs
+++ b/compiler/crates/relay-codemod/src/codemod.rs
@@ -8,8 +8,10 @@
 use std::fs;
 use std::sync::Arc;
 
-use clap::ValueEnum;
+use clap::Args;
+use clap::Subcommand;
 use common::FeatureFlag;
+use common::Rollout;
 use log::info;
 use lsp_types::CodeActionOrCommand;
 use lsp_types::TextEdit;
@@ -18,10 +20,17 @@ use relay_compiler::config::Config;
 use relay_transforms::fragment_alias_directive;
 use relay_transforms::Programs;
 
-#[derive(ValueEnum, Debug, Clone)]
+#[derive(Subcommand, Debug, Clone)]
 pub enum AvailableCodemod {
     /// Marks unaliased conditional fragment spreads as @dangerously_unaliased_fixme
-    MarkDangerousConditionalFragmentSpreads,
+    MarkDangerousConditionalFragmentSpreads(MarkDangerousConditionalFragmentSpreadsArgs),
+}
+
+#[derive(Args, Debug, Clone)]
+pub struct MarkDangerousConditionalFragmentSpreadsArgs {
+    /// If using a feature flag, specify the rollout percentage. If omitted, the codemod is fully enabled.
+    #[clap(long, short, value_parser=valid_percent)]
+    pub rollout_percentage: Option<u8>,
 }
 
 pub async fn run_codemod(
@@ -32,9 +41,13 @@ pub async fn run_codemod(
     let diagnostics = programs
         .iter()
         .flat_map(|programs| match &codemod {
-            AvailableCodemod::MarkDangerousConditionalFragmentSpreads => {
-                // TODO: Figure out how to accept FeatureFlag as an optional CLI argument
-                match fragment_alias_directive(&programs.source, &FeatureFlag::Enabled) {
+            AvailableCodemod::MarkDangerousConditionalFragmentSpreads(opts) => {
+                match fragment_alias_directive(
+                    &programs.source,
+                    &FeatureFlag::Rollout {
+                        rollout: Rollout(opts.rollout_percentage),
+                    },
+                ) {
                     Ok(_) => vec![],
                     Err(e) => e,
                 }
@@ -122,3 +135,14 @@ fn sort_changes(url: &Url, changes: &mut Vec<TextEdit>) -> Result<(), std::io::E
     }
     Ok(())
 }
+
+fn valid_percent(s: &str) -> Result<u8, String> {
+    // turn s into a u8
+    let s = s.parse::<u8>().map_err(|_| "not a number".to_string())?;
+    // check if s is less than 100
+    if (0..=100).contains(&s) {
+        Ok(s)
+    } else {
+        Err("number must be between 0 and 100, inclusive".to_string())
+    }
+}