Merge pull request #28 from Automattic/simplify-exports

Simplify buildkite job script generation

Author: crazytonyli

Package.resolved

@@ -16,6 +16,9 @@ let package = Package(
         .package(url: "https://github.com/ebraraktas/swift-tqdm.git", from: "0.1.2"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(name: "kcpassword", url: "https://github.com/jkmassel/kcpassword-swift.git", from: "1.0.0"),
+        .package(url: "https://github.com/swiftpackages/DotEnv.git", from: "3.0.0"),
+        .package(url: "https://github.com/apple/swift-tools-support-core", from: "0.2.5")
+
     ],
     targets: [
         // Targets are the basic building blocks of a package. A target can define a module or a test suite.
@@ -37,14 +40,23 @@ let package = Package(
             dependencies: [
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
                 .product(name: "SotoS3", package: "soto"),
+                .product(name: "TSCBasic", package: "swift-tools-support-core")
             ]
         ),
         .testTarget(
             name: "libhostmgrTests",
-            dependencies: ["libhostmgr"],
+            dependencies: [
+                "libhostmgr",
+                .product(name: "DotEnv", package: "DotEnv"),
+            ],
             resources: [
                 .copy("resources/configurations/0.6.0.json"),
                 .copy("resources/configurations/defaults.json"),
+                .copy("resources/buildkite-environment-variables-basic-expected-output.txt"),
+                .copy("resources/buildkite-environment-variables-basic.env"),
+                .copy("resources/buildkite-environment-variables-with-code-quotes.env"),
+                .copy("resources/buildkite-commit-message-original.txt"),
+                .copy("resources/buildkite-commit-message-expected.txt"),
             ]
         ),
     ]

Package.swift

@@ -11,101 +11,34 @@ struct GenerateBuildkiteJobScript: ParsableCommand {
     )
 
     func run() throws {
-        let path = try FileManager.default.createTemporaryFile(containing: generateBuildScript()).path
-        print(path)
-    }
-
-    func generateBuildScript() throws -> String {
-        let exports = try generateExports()
-            .map { "export \($0.key)=\"\($0.value)\"" }
-            .joined(separator: "\n")
-
-        return [
-            "source ~/.circ",               // Need to source .circ first in order to set up the SSH session properly
-            exports,                        // Declare all of our environment variables
-            "buildkite-agent bootstrap"    // Then let's go!
-        ].joined(separator: "\n")
-    }
+        var scriptBuilder = BuildkiteScriptBuilder()
 
-    // swiftlint:disable:next function_body_length
-    func generateExports(
-        from environment: [String: String] = ProcessInfo.processInfo.environment
-    ) throws -> [String: String] {
+        scriptBuilder.addDependency(atPath: "~/.circ")
+        scriptBuilder.copyEnvironmentVariables(prefixedBy: "BUILDKITE_")
+        scriptBuilder.addCommand("buildkite-agent bootstrap")
 
-        // See https://buildkite.com/docs/pipelines/environment-variables for
-        // and up-to-date list of environment variables that Buildkite exports
-        let copyableExports = [
-            "BUILDKITE",
-            "BUILDKITE_AGENT_DEBUG",
-            "BUILDKITE_AGENT_DEBUG_HTTP",
-            "BUILDKITE_AGENT_ENDPOINT",
-            "BUILDKITE_AGENT_INCLUDE_RETRIED_JOBS",
-            "BUILDKITE_AGENT_META_DATA_QUEUE",
-            "BUILDKITE_AGENT_NAME",
-            "BUILDKITE_AGENT_NO_COLOR",
-            "BUILDKITE_AGENT_PROFILE",
-            "BUILDKITE_ARTIFACT_PATHS",
-            "BUILDKITE_ARTIFACT_UPLOAD_DESTINATION",
-            "BUILDKITE_BOOTSTRAP_PHASES",
-            "BUILDKITE_BRANCH",
-            "BUILDKITE_BUILD_NUMBER",
-            "BUILDKITE_BUILD_URL",
-            "BUILDKITE_CLEAN_CHECKOUT",
-            "BUILDKITE_COMMIT",
-            "BUILDKITE_GIT_CLONE_FLAGS",
-            "BUILDKITE_HOOKS_PATH",
-            "BUILDKITE_JOB_ID",
-            "BUILDKITE_LABEL",
-            "BUILDKITE_MESSAGE",
-            "BUILDKITE_NO_HTTP2",
-            "BUILDKITE_ORGANIZATION_SLUG",
-            "BUILDKITE_PIPELINE_PROVIDER",
-            "BUILDKITE_PIPELINE_SLUG",
-            "BUILDKITE_PULL_REQUEST",
-            "BUILDKITE_PULL_REQUEST_REPO",
-            "BUILDKITE_REFSPEC",
-            "BUILDKITE_REPO",
-            "BUILDKITE_TAG",
+        // Manually specify the build path to keep them nice and clean in the output
+        scriptBuilder.addEnvironmentVariable(
+            named: "BUILDKITE_BUILD_PATH",
+            value: "/usr/local/var/buildkite-agent/builds"
+        )
 
-            // These ones aren't printed as part of the default list – we're copying them so that `bootstrap` works
-            "BUILDKITE_AGENT_ACCESS_TOKEN",
-            "BUILDKITE_BUILD_ID",
-            "BUILDKITE_PLUGINS_PATH"
+        // Keep the agent name simple for better folder paths
+        scriptBuilder.addEnvironmentVariable(named: "BUILDKITE_AGENT_NAME", value: "builder")
 
-        ].reduce([String: String]()) { dictionary, key in
-            var mutableDictionary = dictionary
-            mutableDictionary[key] = environment[key]
-            return mutableDictionary
-        }
+        // Required to convince `fastlane` that we're running in CI
+        scriptBuilder.addEnvironmentVariable(named: "CI", value: "true")
 
-        return [
-            // Manually specify the build path to keep them nice and clean in the output
-            "BUILDKITE_BUILD_PATH": "/usr/local/var/buildkite-agent/builds",
+        // Used by the S3 Git Mirror plugin
+        scriptBuilder.addEnvironmentVariable(
+            named: "GIT_MIRROR_SERVER_ROOT",
+            value: "http://\(try getIpAddress()):\( Configuration.shared.gitMirrorPort)"
+        )
 
-            // Keep the agent name simple for better folder paths
-            "BUILDKITE_AGENT_NAME": "builder",
+        let scriptText = scriptBuilder.build()
 
-            // Required to convince `fastlane` that we're running in CI
-            "CI": "true",
-
-            // We need to escape double slashes in the command / script, otherwise we can't pass it via command line
-            "BUILDKITE_COMMAND": environment["BUILDKITE_COMMAND"]?.escapingQuotes(),
-            "BUILDKITE_SCRIPT_PATH": environment["BUILDKITE_SCRIPT_PATH"]?.escapingQuotes(),
-
-            // We need to escape double slashes in the Buildkite Plugins JSON,
-            // otherwise it causes an inscrutable error like:
-            //
-            // ```
-            // Error: Failed to parse a plugin definition: invalid character 'g'
-            // looking for beginning of object key string
-            // ```
-            "BUILDKITE_PLUGINS": environment["BUILDKITE_PLUGINS"]?.escapingQuotes(),
-
-            // Used by the S3 Git Mirror plugin
-            "GIT_MIRROR_SERVER_ROOT": "http://\(try getIpAddress()):\( Configuration.shared.gitMirrorPort)"
-        ]
-        .merging(copyableExports, uniquingKeysWith: { lhs, _ in lhs })
-        .compactMapValues { $0 }
+        let path = try FileManager.default.createTemporaryFile(containing: scriptText).path
+        print(path)
     }
 
     // A somewhat hack-ey way to get the device's IP address, but it should continue

Sources/hostmgr/commands/generate/GenerateBuildkiteJobScript.swift

@@ -35,10 +35,6 @@ extension URL: ExpressibleByArgument {
 }
 
 extension String {
-    func escapingQuotes() -> Self {
-        self.replacingOccurrences(of: "\"", with: "\\\"")
-    }
-
     var expandingTildeInPath: String {
         return NSString(string: self).expandingTildeInPath
     }

Sources/hostmgr/helpers/Foundation.swift

@@ -0,0 +1,163 @@
+import Foundation
+import TSCBasic
+
+public struct BuildkiteScriptBuilder {
+
+    /// A list of files to run `source` against prior to executing the rest of the script.
+    var dependencies = [String]()
+
+    /// A list of key/value pairs representing environment variables that should be defined at the start of the script.
+    var environmentVariables: [String: Value] = [:]
+
+    /// A list of commands to run in the script.
+    var commands = [Command]()
+
+    public init() {
+        self.environmentVariables["BUILDKITE"] = Value(wrapping: "true")
+    }
+
+    /// Add another dependency to the build script.
+    ///
+    /// Each dependency will be placed in a `source $DEPENDENCY` block at the top of the emitted build script
+    public mutating func addDependency(atPath path: String) {
+        self.dependencies.append(path)
+    }
+
+    /// Add an environment variable pair to the build script.
+    ///
+    /// If there's an existing environment variable with the same name, it will be overwritten.
+    public mutating func addEnvironmentVariable(named key: String, value: String) {
+        self.environmentVariables[key] = Value(wrapping: value)
+    }
+
+    /// Copy environment variables from the existing environment into the build script based on their prefix.
+    public mutating func copyEnvironmentVariables(
+        prefixedBy prefix: String,
+        from environment: [String: String] = ProcessInfo.processInfo.environment
+    ) {
+        for (key, value) in environment where key.starts(with: prefix) {
+            environmentVariables[key] = Value(wrapping: value)
+        }
+    }
+
+    /// Add a line to the build script.
+    ///
+    /// This typically takes the form of a single command (like `cp foo bar`).
+    public mutating func addCommand(_ command: String, _ arguments: String...) {
+        self.commands.append(Command(command: command, arguments: arguments))
+    }
+
+    /// Compile the build script into a single string
+    public func build() -> String {
+        return [
+            dependencies
+                .map(convertDependencyToSource)
+                .joined(separator: "\n"),
+            environmentVariables
+                .sorted { $0.0 < $1.0 }
+                .filter { !$0.value.rawValue.isEmpty }
+                .map(convertEnvironmentVariableToExport)
+                .joined(separator: "\n"),
+            commands
+                .map(escapeCommand)
+                .joined(separator: "\n")
+        ].joined(separator: "\n")
+    }
+
+    /// Helper that takes a path like `~/.bashrc` and make it into a bash `source` command.
+    ///
+    /// Escapes spaces in paths automatically.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// # Given a path of `~/.bashrc`:
+    /// source ~/.bashrc
+    /// ```
+    func convertDependencyToSource(_ path: String) -> String {
+        "source \(path.escapingSpaces)".trimmingWhitespace
+    }
+
+    /// Helper that takes an environment variable key/value pair to an `export` statement.
+    ///
+    /// Automatically quote-wraps the value and escapes quotes as needed.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// # Given `foo:bar`
+    /// export foo="bar"
+    /// ```
+    func convertEnvironmentVariableToExport(_ pair: (String, Value)) -> String {
+        return "export \(pair.0)=\"\(pair.1.escapedRepresentation)\"".trimmingWhitespace
+    }
+
+    /// Helper that wraps command escape logic for shorthand use in a `map` statement.
+    func escapeCommand(_ command: Command) -> String {
+        command.escapedText
+    }
+
+    /// An object representing the `value` in an environment variable's key/value pair.
+    ///
+    /// Mostly just a way to organize escaping
+    struct Value: Equatable {
+        let rawValue: String
+
+        init(wrapping: String) {
+            self.rawValue = wrapping
+        }
+
+        /// An version of this value suitable for placement in a shell script
+        var escapedRepresentation: String {
+            rawValue
+                .escapingCodeQuotes
+                .escapingDoubleQuotes
+        }
+    }
+
+    /// An object representing one command in a shell script
+    struct Command {
+        /// The underlying command. If you want to handle escaping yourself, put the entire command here.
+        let command: String
+
+        /// The arguments for the command.
+        let arguments: [String]
+
+        init(command: String, arguments: [String] = []) {
+            self.command = command
+            self.arguments = arguments
+        }
+
+        init(_ command: String, _ arguments: String...) {
+            self.init(command: command, arguments: arguments)
+        }
+
+        /// The escaped command text, sutible for placement in a shell script
+        var escapedText: String {
+            "\(command) \(escapedArguments.joined(separator: " "))".trimmingWhitespace
+        }
+
+        /// A helper to print only the escaped arguments
+        var escapedArguments: [String] {
+            arguments.map { $0.spm_shellEscaped() }
+        }
+    }
+}
+
+extension String {
+    var escapingSpaces: String {
+        replacingOccurrences(of: " ", with: "\\ ")
+    }
+
+    var escapingCodeQuotes: String {
+        replacingOccurrences(of: "`", with: "\\`")
+    }
+
+    var escapingDoubleQuotes: String {
+        replacingOccurrences(of: "\"", with: "\\\"")
+    }
+
+    var trimmingWhitespace: String {
+        trimmingCharacters(in: .whitespacesAndNewlines)
+    }
+}