feat: io streams and csv

plmercereau · Apr 24, 2023 · ef37622 · ef37622
1 parent 6b4ac2a
commit ef37622
Show file tree

Hide file tree

Showing 23 changed files with 224 additions and 102 deletions.
diff --git a/.changeset/fuzzy-humans-swim.md b/.changeset/fuzzy-humans-swim.md
@@ -0,0 +1,5 @@
+---
+"chat-dbt": minor
+---
+
+select output streams for sql query, results and information: stdout, stderr, or none
diff --git a/.changeset/light-mails-think.md b/.changeset/light-mails-think.md
@@ -0,0 +1,5 @@
+---
+"chat-dbt": minor
+---
+
+output to csv
diff --git a/.changeset/ninety-candles-leave.md b/.changeset/ninety-candles-leave.md
@@ -0,0 +1,5 @@
+---
+"chat-dbt": minor
+---
+
+execute query from stdin
diff --git a/README.md b/README.md
@@ -103,6 +103,8 @@ chat-dbt --history-mode=[all|none|queries]
 
 Please note that the previous query will however always be sent when you asked to retry a query that failed.
 
+<!-- TODO explain a bit further why this feature is powerful, but why it costs an arm -->
+
 ### Change the format of the result
 
 By default, Chat-DBT renders the results as a table. To obtain the results in JSON format, pass the following option:
@@ -111,6 +113,8 @@ By default, Chat-DBT renders the results as a table. To obtain the results in JS
 chat-dbt --format json
 ```
 
+<!-- TODO IO streams -->
+
 ## Development
 
 ```sh

diff --git a/ROADMAP.md b/ROADMAP.md
@@ -23,10 +23,5 @@
 -   Keep the modifications of the edited SQL query in the history
 -   Dark/light mode toggle
 -   Add CI tests for both the CLI and the web interface
--   `--hide-sql`
--   `--hide-result`
--   Input from stdin/output to stdout (combined with the above)
--   Output sql (on success) to file
--   Output result to file
 -   Verbose level, timestamps etc
 -   Comment the code
diff --git a/cli/cli.ts b/cli/cli.ts
@@ -1,17 +1,22 @@
 import chalk from 'chalk'
+import { stringify } from 'csv-stringify/sync'
 import inquirer from 'inquirer'
-import { Configuration, OpenAIApi } from 'openai'
+import { OpenAIApi } from 'openai'
 import ora from 'ora'
+import readline from 'readline'
 import table from 'tty-table'
 
 import {
+    OutputStream,
+    ResultFormat,
+    ERROR_PROMPT,
+    getIntrospection,
+    initOpenAI,
     GptSqlResponse,
     GptSqlResultItem,
     getSqlQuery,
     runSqlQuery
-} from '@/shared/chat-gpt'
-import { ERROR_PROMPT } from '@/shared/error'
-import { getIntrospection } from '@/shared/introspection'
+} from '@/shared'
 
 import { CommonOptions } from './index'
 import InputHistoryPrompt from './input-history'
@@ -22,37 +27,55 @@ type YesNo = 'yes' | 'no'
 inquirer.registerPrompt('input-history', InputHistoryPrompt)
 
 export const startCLI = async (options: CommonOptions) => {
-    const { key, org } = options
-    const openai = new OpenAIApi(
-        new Configuration({ apiKey: key, organization: org })
-    )
+    const openai = initOpenAI(options.key, options.org)
     const history: GptSqlResponse[] = []
-    // * The history of queries is not calculated from the full his
-    // * as the history may not be saved between queries when using
-    // * the option --history=queries or --history=none
-    const promptHistory: string[] = []
 
-    while (true) {
-        const { query } = await inquirer.prompt([
-            {
-                type: 'input-history',
-                name: 'query',
-                message: 'Describe your query',
-                validate: value => !!value || 'Query cannot be empty',
-                history: promptHistory
-            }
-        ])
-        promptHistory.push(query)
-        await executeQueryAndShowResult({
-            openai,
-            query,
-            history,
-            ...options
+    if (process.stdin.isTTY) {
+        // * Interactive mode
+        // * The history of queries is not calculated from the full his
+        // * as the history may not be saved between queries when using
+        // * the option --history=queries or --history=none
+        const promptHistory: string[] = []
+        while (true) {
+            const { query } = await inquirer.prompt([
+                {
+                    type: 'input-history',
+                    name: 'query',
+                    message: 'Describe your query',
+                    validate: value => !!value || 'Query cannot be empty',
+                    history: promptHistory
+                }
+            ])
+            promptHistory.push(query)
+            await executeQueryAndShowResult({
+                openai,
+                query,
+                history,
+                ...options
+            })
+        }
+    } else {
+        // * Pipe mode
+        const rl = readline.createInterface({
+            input: process.stdin,
+            terminal: false
         })
+        for await (const query of rl) {
+            await executeQueryAndShowResult({
+                openai,
+                query,
+                history,
+                ...options
+            })
+        }
     }
 }
 
-const printResult = async (result: GptSqlResultItem[], format: string) => {
+const printResult = async (
+    result: GptSqlResultItem[],
+    format: ResultFormat,
+    output: OutputStream
+) => {
     switch (format) {
         case 'table':
             result.forEach(item => {
@@ -67,32 +90,54 @@ const printResult = async (result: GptSqlResultItem[], format: string) => {
                         item.rows,
                         { defaultValue: '' }
                     )
-                    console.log(t.render())
+                    println(t.render(), output)
                 } else {
                     // * No data is expected e.g. INSERT statement with no RETURNING clause
                     const t = table(
                         [{ value: 'count' }],
                         [{ count: item.count }]
                     )
-                    console.log(t.render())
+                    println(t.render(), output)
                 }
             })
             break
         case 'json':
-            console.log(
+            println(
                 JSON.stringify(
                     result.map(({ count, rows, columns }) =>
                         columns ? rows : count === null ? true : count
                     ),
                     null,
                     2
-                )
+                ),
+                output
             )
             break
+        case 'csv':
+            result.forEach(({ count, rows, columns }) => {
+                println(
+                    stringify(
+                        columns
+                            ? rows
+                            : count === null
+                            ? [{ success: true }]
+                            : [{ count }],
+                        { header: true }
+                    ),
+                    output
+                )
+            })
+            break
+    }
+}
+const println = (message: string, output: OutputStream) => {
+    if (output === 'none') {
+        return
     }
+    output === 'stdout' ? console.log(message) : console.error(message)
 }
 
-const executeQueryAndShowResult = async ({
+export const executeQueryAndShowResult = async ({
     query,
     history = [],
     ...options
@@ -101,8 +146,10 @@ const executeQueryAndShowResult = async ({
     query: string
     openai: OpenAIApi
     history?: GptSqlResponse[]
+    stdin?: boolean
 }) => {
-    const spinner = ora()
+    const tty = process.stdin.isTTY
+    const spinner = ora({ isSilent: !tty })
     spinner.start()
     let sqlQuery: string = ''
     try {
@@ -119,13 +166,16 @@ const executeQueryAndShowResult = async ({
                 introspection
             })
             sqlQuery = result.sqlQuery
-            console.log(`${result.usage?.total_tokens} tokens used`)
-            if (options.confirm) {
+            println(
+                `${result.usage?.total_tokens} tokens used`,
+                options.outputInfo
+            )
+            if (options.confirm && tty) {
                 // * Ask the user's confirmation before executing the SQL query
                 spinner.stop()
                 // * Confirm the SQL query, with the possibility to edit it
                 const confirmPrompt = async (): Promise<void> => {
-                    console.log(chalk.greenBright(sqlQuery))
+                    println(chalk.greenBright(sqlQuery), options.outputSql)
                     const { confirm } = await inquirer.prompt<{
                         confirm: YesNo | 'edit'
                     }>([
@@ -170,46 +220,52 @@ const executeQueryAndShowResult = async ({
         spinner.stop()
         if (!options.confirm) {
             // * Print the SQL query, but only if it's not already printed
-            console.log(chalk.dim(sqlQuery))
+            println(chalk.dim(sqlQuery), options.outputSql)
         }
         spinner.succeed('Success')
         if (options.historyMode === 'none') {
             history.length = 0
         } else {
             history.push({ query, sqlQuery, result })
         }
-        printResult(result, options.format)
+        printResult(result, options.format, options.outputResult)
     } catch (e) {
         const error = e as Error
         spinner.stop()
         if (sqlQuery) {
-            console.log(chalk.dim(sqlQuery))
+            println(chalk.dim(sqlQuery), options.outputSql)
         }
-        spinner.fail(error.message)
+        if (tty) {
+            spinner.fail(error.message)
+        } else {
+            println(error.message, options.outputInfo)
+        }
+
         type Retry = 'yes' | 'no' | 'editPrompt' | 'editSql'
-        let retry: Retry | 'editPrompt' | 'editSql' = 'no'
-        if (
-            typeof options.autoCorrect === 'number' &&
-            options.autoCorrect > 0
-        ) {
+        let retry: Retry = 'no'
+        if (options.autoCorrect > 0) {
             retry = 'yes'
             options.autoCorrect--
-        } else {
+        } else if (tty) {
             const choices = [
                 { key: 'y', name: 'Yes', value: 'yes' },
                 { key: 'n', name: 'No', value: 'no' },
                 { key: 'e', name: 'Edit prompt', value: 'editPrompt' }
             ]
             if (sqlQuery) {
-                choices.push({ key: 'q', name: 'Edit SQL', value: 'editSql' })
+                choices.push({
+                    key: 'q',
+                    name: 'Edit SQL',
+                    value: 'editSql'
+                })
             }
             const prompt = await inquirer.prompt<{ retry: Retry }>([
                 {
                     type: 'list',
                     name: 'retry',
                     message: 'Ask correction?',
                     choices,
-                    default: options.autoCorrect === undefined ? 'yes' : 'no'
+                    default: 'yes'
                 }
             ])
             retry = prompt.retry
@@ -242,14 +298,17 @@ const executeQueryAndShowResult = async ({
                         database: options.database
                     })
                     spinner.stop()
-                    console.log(chalk.dim(sqlQuery))
+                    println(chalk.dim(sqlQuery), options.outputSql)
                     spinner.succeed('Success')
-                    printResult(result, options.format)
+                    printResult(result, options.format, options.outputResult)
                 } finally {
                     return
                 }
             }
-            console.log(chalk.blue('!'), 'Retrying with:', chalk.bold(query))
+            println(
+                `${chalk.blue('!')} Retrying with: ${chalk.bold(query)}`,
+                options.outputInfo
+            )
             await executeQueryAndShowResult({
                 query,
                 history,