Allow to specify env vars to be set via the connection file (#653)

almond-sh · Sep 30, 2020 · 372f310 · 372f310
1 parent 6a547f9
commit 372f310
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 10 deletions.
diff --git a/modules/scala/scala-kernel/src/main/scala/almond/ScalaKernel.scala b/modules/scala/scala-kernel/src/main/scala/almond/ScalaKernel.scala
@@ -44,12 +44,14 @@ object ScalaKernel extends CaseApp[Options] {
             .getContextClassLoader
             .getResource("almond/scala-logo-64x64.png")
         ),
+        connectionFileArgs = Install.defaultConnectionFileArgs,
         interruptMode = {
           if (options.installOptions.interruptViaMessage)
             Some("message")
           else
             None
-        }
+        },
+        env = options.installOptions.envMap()
       ) match {
         case Left(e) =>
           log.debug("Cannot install kernel", e)

diff --git a/modules/shared/kernel/src/main/scala/almond/kernel/install/Install.scala b/modules/shared/kernel/src/main/scala/almond/kernel/install/Install.scala
@@ -197,15 +197,38 @@ object Install {
         None
     }
 
+  def defaultConnectionFileArgs: Seq[String] =
+    Seq("--connection-file", "{connection_file}")
 
   def install(
     defaultId: String,
     defaultDisplayName: String,
     language: String,
     options: Options,
     defaultLogoOpt: Option[URL] = None,
-    connectionFileArgs: Seq[String] = Seq("--connection-file", "{connection_file}"),
+    connectionFileArgs: Seq[String] = defaultConnectionFileArgs,
     interruptMode: Option[String] = None
+  ): Path =
+    install(
+      defaultId,
+      defaultDisplayName,
+      language,
+      options,
+      defaultLogoOpt,
+      connectionFileArgs,
+      interruptMode,
+      Map.empty
+    )
+
+  def install(
+    defaultId: String,
+    defaultDisplayName: String,
+    language: String,
+    options: Options,
+    defaultLogoOpt: Option[URL],
+    connectionFileArgs: Seq[String],
+    interruptMode: Option[String],
+    env: Map[String, String]
   ): Path = {
 
     val path =
@@ -248,10 +271,11 @@ object Install {
     Install.installIn(
       options.id.getOrElse(defaultId),
       KernelSpec(
-        (cmd ++ connectionFileArgs).toList,
-        options.displayName.getOrElse(defaultDisplayName),
-        language,
-        interrupt_mode = interruptMode
+        argv = (cmd ++ connectionFileArgs).toList,
+        display_name = options.displayName.getOrElse(defaultDisplayName),
+        language = language,
+        interrupt_mode = interruptMode,
+        env = env
       ),
       path,
       logo64PngOpt = logoOpt,
@@ -266,8 +290,29 @@ object Install {
     language: String,
     options: Options,
     defaultLogoOpt: Option[URL] = None,
-    connectionFileArgs: Seq[String] = Seq("--connection-file", "{connection_file}"),
+    connectionFileArgs: Seq[String] = defaultConnectionFileArgs,
     interruptMode: Option[String] = None
+  ): Either[InstallException, Path] =
+    installOrError(
+      defaultId,
+      defaultDisplayName,
+      language,
+      options,
+      defaultLogoOpt,
+      connectionFileArgs,
+      interruptMode,
+      Map.empty
+    )
+
+  def installOrError(
+    defaultId: String,
+    defaultDisplayName: String,
+    language: String,
+    options: Options,
+    defaultLogoOpt: Option[URL],
+    connectionFileArgs: Seq[String],
+    interruptMode: Option[String],
+    env: Map[String, String]
   ): Either[InstallException, Path] =
     try {
       val dir = install(
@@ -277,7 +322,8 @@ object Install {
         options,
         defaultLogoOpt,
         connectionFileArgs,
-        interruptMode
+        interruptMode,
+        env
       )
       Right(dir)
     } catch {

diff --git a/modules/shared/kernel/src/main/scala/almond/kernel/install/Options.scala b/modules/shared/kernel/src/main/scala/almond/kernel/install/Options.scala
@@ -1,6 +1,6 @@
 package almond.kernel.install
 
-import caseapp.{HelpMessage, Name}
+import caseapp.{HelpMessage, Name, ValueDescription}
 
 final case class Options(
   @HelpMessage("erase any previously existing kernel with the same id")
@@ -25,10 +25,22 @@ final case class Options(
   @HelpMessage("whether to request frontends to interrupt this kernel via a message")
     interruptViaMessage: Boolean = false,
   @HelpMessage("Whether to copy the kernel launcher in the kernelspec directory (default: false if --arg or --command specified, true else)")
-    copyLauncher: Option[Boolean] = None
+    copyLauncher: Option[Boolean] = None,
+  @HelpMessage("Environment variables to pass to the kernel via its connection file")
+  @ValueDescription("name=value")
+    env: List[String] = Nil
 ) {
   def copyLauncher0: Boolean =
     copyLauncher.getOrElse {
       arg.isEmpty && command.isEmpty
     }
+  def envMap(): Map[String, String] =
+    env
+      .map { input =>
+        input.split("=", 2) match {
+          case Array(k, v) => (k, v)
+          case _ => sys.error(s"Malformed --env value '$input' (expected 'name=value')")
+        }
+      }
+      .toMap
 }