From 336d759ec934ac2c285ec6393ed500dda5506319 Mon Sep 17 00:00:00 2001
From: Pierre Chalamet <pierre@chalamet.net>
Date: Sun, 5 Jan 2025 09:16:21 +0100
Subject: [PATCH] fix script loading

---
 .../TestFiles/WORKSPACE                       |  5 +++
 .../Workspace.fs                              |  8 ++++-
 src/Terrabuild/Core/Configuration.fs          | 31 ++++++++++++++-----
 src/Terrabuild/Helpers/Extensions.fs          |  9 +++---
 4 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/src/Terrabuild.Configuration.Tests/TestFiles/WORKSPACE b/src/Terrabuild.Configuration.Tests/TestFiles/WORKSPACE
index 5c1cd9bf..58638012 100644
--- a/src/Terrabuild.Configuration.Tests/TestFiles/WORKSPACE
+++ b/src/Terrabuild.Configuration.Tests/TestFiles/WORKSPACE
@@ -35,6 +35,11 @@ extension dotnet {
 
 extension docker { }
 
+extension npmext {
+    container = "node:20"
+    script = "scripts/npm.fsx"
+}
+
 # multiple comments at the end
 # multiple comments at the end
 
diff --git a/src/Terrabuild.Configuration.Tests/Workspace.fs b/src/Terrabuild.Configuration.Tests/Workspace.fs
index 940acf56..a180d71c 100644
--- a/src/Terrabuild.Configuration.Tests/Workspace.fs
+++ b/src/Terrabuild.Configuration.Tests/Workspace.fs
@@ -36,6 +36,11 @@ let parseWorkspace() =
               Variables = Set.empty
               Script = None
               Defaults = Map.empty }
+        let extNpm =
+            { Container = Some "node:20"
+              Variables = Set.empty
+              Script = Some "scripts/npm.fsx"
+              Defaults = Map.empty }
 
         { Workspace = { Space = Some "magnusopera/default"; Ignores = Set ["**/node_modules"] }
           Targets = Map [ "build", targetBuild
@@ -44,7 +49,8 @@ let parseWorkspace() =
           Configurations = Map [ "release", envRelease
                                  "dummy", envDummy ]
           Extensions = Map [ "dotnet", extDotnet
-                             "docker", extDocker ] }
+                             "docker", extDocker
+                             "npmext", extNpm ] }
 
 
     let content = File.ReadAllText("TestFiles/WORKSPACE")
diff --git a/src/Terrabuild/Core/Configuration.fs b/src/Terrabuild/Core/Configuration.fs
index 4c34e216..71fd4520 100644
--- a/src/Terrabuild/Core/Configuration.fs
+++ b/src/Terrabuild/Core/Configuration.fs
@@ -174,14 +174,29 @@ let read (options: ConfigOptions.Options) =
         { evaluationContext with Eval.Variables = evaluationContext.Variables |> Map.addMap configVariables }
 
 
-    let extensions = 
-        Extensions.systemExtensions
-        |> Map.addMap workspaceConfig.Extensions
-
-    let scripts =
-        extensions
-        |> Map.map (fun _ _ -> None)
-        |> Map.map Extensions.lazyLoadScript
+    let extensions, scripts =
+        // load system extensions
+        let sysScripts =
+            Extensions.systemExtensions
+            |> Map.map (fun _ _ -> None)
+            |> Map.map Extensions.lazyLoadScript
+
+        // load user extension
+        let usrScripts =
+            workspaceConfig.Extensions
+            |> Map.map (fun _ ext ->
+                match ext.Script with
+                | Some script -> script |> FS.workspaceRelative options.Workspace "" |> Some
+                | _ -> None)
+            |> Map.map Extensions.lazyLoadScript
+
+        let extensions =
+            Extensions.systemExtensions
+            |> Map.addMap workspaceConfig.Extensions
+
+        let scripts = sysScripts |> Map.addMap usrScripts
+
+        extensions, scripts
 
 
     // this is the first stage: load project and mostly get dependencies references
diff --git a/src/Terrabuild/Helpers/Extensions.fs b/src/Terrabuild/Helpers/Extensions.fs
index f7bc79d6..c985aad0 100644
--- a/src/Terrabuild/Helpers/Extensions.fs
+++ b/src/Terrabuild/Helpers/Extensions.fs
@@ -29,7 +29,8 @@ let terrabuildExtensibility =
 let lazyLoadScript (name: string) (script: string option) =
     let initScript () =
         match script with
-        | Some script -> loadScript [ terrabuildExtensibility ] script
+        | Some script ->
+            loadScript [ terrabuildExtensibility ] script
         | _ ->
             match Terrabuild.Extensions.Factory.systemScripts |> Map.tryFind name with
             | Some sysTpe -> Script(sysTpe)
@@ -38,9 +39,9 @@ let lazyLoadScript (name: string) (script: string option) =
     lazy(initScript())
 
 let getScript (extension: string) (scripts: Map<string, Lazy<Script>>) =
-    match scripts |> Map.tryFind extension with
-    | None -> None
-    | Some script -> script.Value |> Some
+    scripts
+    |> Map.tryFind extension
+    |> Option.map (fun script -> script.Value)
 
 let invokeScriptMethod<'r> (method: string) (args: Value) (script: Script option) =
     match script with