From b97651c9438011631849dad7249af79876abf948 Mon Sep 17 00:00:00 2001
From: Nathan Whitaker <nathan@deno.com>
Date: Mon, 9 Dec 2024 13:03:58 -0800
Subject: [PATCH] add trait to allow more control over version resolution

---
 lib/lib.rs              |  1 +
 src/graph.rs            | 26 ++++++++++++++++++++++++++
 src/jsr.rs              | 12 ++++++++++++
 src/packages.rs         | 10 ++++++++++
 tests/ecosystem_test.rs |  1 +
 5 files changed, 50 insertions(+)

diff --git a/lib/lib.rs b/lib/lib.rs
index 392235d5e..c9646b908 100644
--- a/lib/lib.rs
+++ b/lib/lib.rs
@@ -282,6 +282,7 @@ pub async fn js_create_graph(
         imports,
         reporter: None,
         executor: Default::default(),
+        jsr_version_resolver: None,
       },
     )
     .await;
diff --git a/src/graph.rs b/src/graph.rs
index d75c712f9..038fbfc43 100644
--- a/src/graph.rs
+++ b/src/graph.rs
@@ -24,6 +24,7 @@ use crate::module_specifier::SpecifierError;
 use crate::packages::resolve_version;
 use crate::packages::JsrPackageInfo;
 use crate::packages::JsrPackageVersionInfo;
+use crate::packages::JsrVersionResolver;
 use crate::packages::PackageSpecifiers;
 use crate::rt::Executor;
 
@@ -1212,6 +1213,7 @@ pub struct BuildOptions<'a> {
   pub npm_resolver: Option<&'a dyn NpmResolver>,
   pub reporter: Option<&'a dyn Reporter>,
   pub resolver: Option<&'a dyn Resolver>,
+  pub jsr_version_resolver: Option<&'a dyn JsrVersionResolver>,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -3403,6 +3405,7 @@ struct Builder<'a, 'graph> {
   fill_pass_mode: FillPassMode,
   executor: &'a dyn Executor,
   resolved_roots: BTreeSet<ModuleSpecifier>,
+  jsr_version_resolver: Option<&'a dyn JsrVersionResolver>,
 }
 
 impl<'a, 'graph> Builder<'a, 'graph> {
@@ -3433,6 +3436,7 @@ impl<'a, 'graph> Builder<'a, 'graph> {
       fill_pass_mode,
       executor: options.executor,
       resolved_roots: Default::default(),
+      jsr_version_resolver: options.jsr_version_resolver,
     };
     builder.fill(roots, options.imports).await;
   }
@@ -3641,6 +3645,7 @@ impl<'a, 'graph> Builder<'a, 'graph> {
                     executor: self.executor,
                     jsr_url_provider: self.jsr_url_provider,
                     loader: self.loader,
+                    jsr_version_resolver: self.jsr_version_resolver,
                   },
                 );
                 pending_resolutions.push_front(pending_resolution);
@@ -3951,6 +3956,25 @@ impl<'a, 'graph> Builder<'a, 'graph> {
     package_info: &JsrPackageInfo,
     package_req: &PackageReq,
   ) -> Option<Version> {
+    if let Some(version_resolver) = self.jsr_version_resolver {
+      let version =
+        version_resolver.resolve_package_version(package_info, package_req);
+      if let Some(version) = version {
+        let is_yanked = package_info
+          .versions
+          .get(version)
+          .map(|i| i.yanked)
+          .unwrap_or(false);
+        let version = version.clone();
+        if is_yanked {
+          self.graph.packages.add_used_yanked_package(PackageNv {
+            name: package_req.name.clone(),
+            version: version.clone(),
+          });
+        }
+        return Some(version);
+      }
+    }
     // 1. try to resolve with the list of existing versions
     if let Some(existing_versions) =
       self.graph.packages.versions_by_name(&package_req.name)
@@ -4715,6 +4739,7 @@ impl<'a, 'graph> Builder<'a, 'graph> {
         executor: self.executor,
         jsr_url_provider: self.jsr_url_provider,
         loader: self.loader,
+        jsr_version_resolver: self.jsr_version_resolver,
       },
     );
   }
@@ -4728,6 +4753,7 @@ impl<'a, 'graph> Builder<'a, 'graph> {
         executor: self.executor,
         jsr_url_provider: self.jsr_url_provider,
         loader: self.loader,
+        jsr_version_resolver: self.jsr_version_resolver,
       },
     );
   }
diff --git a/src/jsr.rs b/src/jsr.rs
index 83a3335c3..e4ced1d08 100644
--- a/src/jsr.rs
+++ b/src/jsr.rs
@@ -12,6 +12,7 @@ use futures::FutureExt;
 use crate::graph::JsrLoadError;
 use crate::packages::JsrPackageInfo;
 use crate::packages::JsrPackageVersionInfo;
+use crate::packages::JsrVersionResolver;
 use crate::rt::spawn;
 use crate::rt::JoinHandle;
 use crate::source::CacheSetting;
@@ -37,6 +38,7 @@ pub struct JsrMetadataStoreServices<'a> {
   pub loader: &'a dyn Loader,
   pub executor: &'a dyn Executor,
   pub jsr_url_provider: &'a dyn JsrUrlProvider,
+  pub jsr_version_resolver: Option<&'a dyn JsrVersionResolver>,
 }
 
 #[derive(Debug, Default)]
@@ -82,6 +84,16 @@ impl JsrMetadataStore {
       .url()
       .join(&format!("{}/meta.json", package_name))
       .unwrap();
+    let cache_setting =
+      if let Some(version_resolver) = services.jsr_version_resolver {
+        if version_resolver.force_fetch_metadata(package_name) {
+          CacheSetting::Reload
+        } else {
+          cache_setting
+        }
+      } else {
+        cache_setting
+      };
     let fut = self.load_data(
       specifier,
       services,
diff --git a/src/packages.rs b/src/packages.rs
index da8baeadd..5489bdc85 100644
--- a/src/packages.rs
+++ b/src/packages.rs
@@ -16,6 +16,16 @@ use serde::Serialize;
 use crate::analyzer::module_graph_1_to_2;
 use crate::ModuleInfo;
 
+pub trait JsrVersionResolver {
+  fn resolve_package_version<'a>(
+    &self,
+    package_info: &'a JsrPackageInfo,
+    package_req: &PackageReq,
+  ) -> Option<&'a Version>;
+
+  fn force_fetch_metadata<'a>(&self, package_name: &str) -> bool;
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct JsrPackageInfo {
   pub versions: HashMap<Version, JsrPackageInfoVersion>,
diff --git a/tests/ecosystem_test.rs b/tests/ecosystem_test.rs
index e156bce6a..8af320910 100644
--- a/tests/ecosystem_test.rs
+++ b/tests/ecosystem_test.rs
@@ -287,6 +287,7 @@ async fn test_version(
         passthrough_jsr_specifiers: true,
         executor: Default::default(),
         imports: vec![],
+        jsr_version_resolver: None,
       },
     )
     .await;