3 files changed, 206 insertions, 9 deletions
diff --git a/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/TraceConverter.kt b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/ConvertCommand.kt
index c71035d4..970de0f4 100644
--- a/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/TraceConverter.kt
+++ b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/ConvertCommand.kt
@@ -20,7 +20,6 @@
  * SOFTWARE.
  */
 
-@file:JvmName("TraceConverter")
 package org.opendc.trace.tools
 
 import com.github.ajalt.clikt.core.CliktCommand
@@ -44,14 +43,9 @@ import kotlin.math.max
 import kotlin.math.min
 
 /**
- * A script to convert a trace in text format into a Parquet trace.
+ * A [CliktCommand] that can convert between workload trace formats.
  */
-fun main(args: Array<String>): Unit = TraceConverterCli().main(args)
-
-/**
- * Represents the command for converting traces
- */
-internal class TraceConverterCli : CliktCommand(name = "trace-converter") {
+internal class ConvertCommand : CliktCommand(name = "convert", help = "Convert between workload trace formats") {
     /**
      * The logger instance for the converter.
      */
@@ -73,7 +67,7 @@ internal class TraceConverterCli : CliktCommand(name = "trace-converter") {
     /**
      * The input format of the trace.
      */
-    private val inputFormat by option("-f", "--input-format", help = "format of output trace")
+    private val inputFormat by option("-f", "--input-format", help = "format of input trace")
         .required()
 
     /**
diff --git a/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/QueryCommand.kt b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/QueryCommand.kt
new file mode 100644
index 00000000..b0f95de2
--- /dev/null
+++ b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/QueryCommand.kt
@@ -0,0 +1,159 @@
+/*
+ * Copyright (c) 2022 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.trace.tools
+
+import com.github.ajalt.clikt.core.CliktCommand
+import com.github.ajalt.clikt.parameters.arguments.argument
+import com.github.ajalt.clikt.parameters.options.option
+import com.github.ajalt.clikt.parameters.options.required
+import com.github.ajalt.clikt.parameters.types.file
+import org.apache.calcite.jdbc.CalciteConnection
+import org.jline.builtins.Styles
+import org.jline.console.Printer
+import org.jline.console.impl.DefaultPrinter
+import org.jline.terminal.Terminal
+import org.jline.terminal.TerminalBuilder
+import org.jline.utils.AttributedStringBuilder
+import org.opendc.trace.Trace
+import org.opendc.trace.calcite.TraceSchema
+import java.nio.charset.StandardCharsets
+import java.sql.DriverManager
+import java.sql.ResultSet
+import java.sql.ResultSetMetaData
+import java.util.*
+
+/**
+ * A [CliktCommand] that allows users to query workload traces using SQL.
+ */
+internal class QueryCommand : CliktCommand(name = "query", help = "Query workload traces") {
+    /**
+     * The trace to open.
+     */
+    private val input by option("-i", "--input")
+        .file(mustExist = true)
+        .required()
+
+    /**
+     * The input format of the trace.
+     */
+    private val inputFormat by option("-f", "--format", help = "format of the trace")
+        .required()
+
+    /**
+     * The query to execute.
+     */
+    private val query by argument()
+
+    /**
+     * Access to the terminal.
+     */
+    private val terminal = TerminalBuilder.builder()
+        .system(false)
+        .streams(System.`in`, System.out)
+        .encoding(StandardCharsets.UTF_8)
+        .build()
+
+    /**
+     * Helper class to print results to console.
+     */
+    private val printer = QueryPrinter(terminal)
+
+    override fun run() {
+        val inputTrace = Trace.open(input, format = inputFormat)
+        val info = Properties().apply { this["lex"] = "JAVA" }
+        val connection = DriverManager.getConnection("jdbc:calcite:", info).unwrap(CalciteConnection::class.java)
+        connection.rootSchema.add("trace", TraceSchema(inputTrace))
+        connection.schema = "trace"
+
+        val stmt = connection.createStatement()
+        stmt.executeQuery(query)
+
+        val start = System.currentTimeMillis()
+        val hasResults = stmt.execute(query)
+
+        try {
+            if (hasResults) {
+                do {
+                    stmt.resultSet.use { rs ->
+                        val count: Int = printResults(rs)
+                        val duration = (System.currentTimeMillis() - start) / 1000.0
+                        printer.println("$count rows selected (${"%.3f".format(duration)} seconds)")
+                    }
+                } while (stmt.moreResults)
+            } else {
+                val count: Int = stmt.updateCount
+                val duration = (System.currentTimeMillis() - start) / 1000.0
+
+                printer.println("$count rows affected (${"%0.3f".format(duration)} seconds)")
+            }
+        } finally {
+            stmt.close()
+            connection.close()
+        }
+    }
+
+    /**
+     * Helper function to print the results to console.
+     */
+    private fun printResults(rs: ResultSet): Int {
+        var count = 0
+        val meta: ResultSetMetaData = rs.metaData
+
+        val options = mapOf(
+            Printer.COLUMNS to List(meta.columnCount) { meta.getColumnName(it + 1) },
+            Printer.BORDER to "|",
+        )
+        val data = mutableListOf<Map<String, Any>>()
+
+        while (rs.next()) {
+            val row = mutableMapOf<String, Any>()
+            for (i in 1..meta.columnCount) {
+                row[meta.getColumnName(i)] = rs.getObject(i)
+            }
+            data.add(row)
+
+            count++
+        }
+
+        printer.println(options, data)
+
+        return count
+    }
+
+    /**
+     * Helper class to print the results of the query.
+     */
+    private class QueryPrinter(private val terminal: Terminal) : DefaultPrinter(null) {
+        override fun terminal(): Terminal = terminal
+
+        override fun highlightAndPrint(options: MutableMap<String, Any>, exception: Throwable) {
+            if (options.getOrDefault("exception", "stack") == "stack") {
+                exception.printStackTrace()
+            } else {
+                val asb = AttributedStringBuilder()
+                asb.append(exception.message, Styles.prntStyle().resolve(".em"))
+                asb.toAttributedString().println(terminal())
+            }
+        }
+    }
+}
diff --git a/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/TraceTools.kt b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/TraceTools.kt
new file mode 100644
index 00000000..b480484b
--- /dev/null
+++ b/opendc-trace/opendc-trace-tools/src/main/kotlin/org/opendc/trace/tools/TraceTools.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2022 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+@file:JvmName("TraceTools")
+package org.opendc.trace.tools
+
+import com.github.ajalt.clikt.core.CliktCommand
+import com.github.ajalt.clikt.core.subcommands
+
+/**
+ * A script for querying and manipulating workload traces supported by OpenDC.
+ */
+fun main(args: Array<String>): Unit = TraceToolsCli().main(args)
+
+/**
+ * The primary [CliktCommand] for the trace tools offered by OpenDC.
+ */
+class TraceToolsCli : CliktCommand(name = "trace-tools") {
+    init {
+        subcommands(QueryCommand())
+        subcommands(ConvertCommand())
+    }
+
+    override fun run() {}
+}