MVP working dataconnect:sql:shell

CHANGELOG.md

@@ -1,3 +1,4 @@
 - Fixed Flutter web apps that might require the --no-tree-shake-icons flag in order to build. (#7724)
 - Removed outdated dependency on `rimraf`.
 - Fixed an issue where the Extensions emulator would fail silently if started with a non-existant project without the `demo-` prefix. (#7779)
+- Added new command dataconnect:sql:shell which allows users to directly connect and run queries against their dataconnect cloudsql instance (#7778).

src/commands/dataconnect-sql-shell.ts

@@ -0,0 +1,139 @@
+import * as pg from "pg";
+import chalk from "chalk";
+import { Connector, IpAddressTypes, AuthTypes } from "@google-cloud/cloud-sql-connector";
+
+import { Command } from "../command";
+import { Options } from "../options";
+import { needProjectId } from "../projectUtils";
+import { ensureApis } from "../dataconnect/ensureApis";
+import { requirePermissions } from "../requirePermissions";
+import { pickService } from "../dataconnect/fileUtils";
+import { getIdentifiers } from "../dataconnect/schemaMigration";
+import { requireAuth } from "../requireAuth";
+import { getIAMUser } from "../gcp/cloudsql/connect";
+import * as cloudSqlAdminClient from "../gcp/cloudsql/cloudsqladmin";
+import { prompt, Question } from "../prompt";
+import { logger } from "../logger";
+import { FirebaseError } from "../error";
+import { FBToolsAuthClient } from "../gcp/cloudsql/fbToolsAuthClient";
+import { confirmDangerousQuery, interactiveExecuteQuery } from "../gcp/cloudsql/interactive";
+
+// Not a comprehensive list, used for keyword coloring.
+const sqlKeywords = [
+  "SELECT",
+  "FROM",
+  "WHERE",
+  "INSERT",
+  "UPDATE",
+  "DELETE",
+  "JOIN",
+  "GROUP",
+  "ORDER",
+  "LIMIT",
+  "GRANT",
+  "CREATE",
+  "DROP",
+];
+
+async function promptForQuery(): Promise<string> {
+  let query = "";
+  let line = "";
+
+  do {
+    const question: Question = {
+      type: "input",
+      name: "line",
+      message: query ? "> " : "Enter your SQL query (or '.exit'):",
+      transformer: (input: string) => {
+        // Highlight SQL keywords
+        return input
+          .split(" ")
+          .map((word) => (sqlKeywords.includes(word.toUpperCase()) ? chalk.cyan(word) : word))
+          .join(" ");
+      },
+    };
+
+    ({ line } = await prompt({ nonInteractive: false }, [question]));
+    line = line.trimEnd();
+
+    if (line.toLowerCase() === ".exit") {
+      return ".exit";
+    }
+
+    query += (query ? "\n" : "") + line;
+  } while (line !== "" && !query.endsWith(";"));
+  return query;
+}
+
+async function mainShellLoop(conn: pg.PoolClient) {
+  while (true) {
+    const query = await promptForQuery();
+    if (query.toLowerCase() === ".exit") {
+      break;
+    }
+
+    if (query === "") {
+      continue;
+    }
+
+    if (await confirmDangerousQuery(query)) {
+      await interactiveExecuteQuery(query, conn);
+    } else {
+      logger.info(chalk.yellow("Query cancelled."));
+    }
+  }
+}
+
+export const command = new Command("dataconnect:sql:shell [serviceId]")
+  .description("Starts a shell connected directly to your dataconnect cloudsql instance.")
+  .before(requirePermissions, ["firebasedataconnect.services.list", "cloudsql.instances.connect"])
+  .before(requireAuth)
+  .action(async (serviceId: string, options: Options) => {
+    const projectId = needProjectId(options);
+    await ensureApis(projectId);
+    const serviceInfo = await pickService(projectId, options.config, serviceId);
+    const { instanceId, databaseId } = getIdentifiers(serviceInfo.schema);
+    const { user: username } = await getIAMUser(options);
+    const instance = await cloudSqlAdminClient.getInstance(projectId, instanceId);
+
+    // Setup the connection
+    const connectionName = instance.connectionName;
+    if (!connectionName) {
+      throw new FirebaseError(
+        `Could not get instance connection string for ${options.instanceId}:${options.databaseId}`,
+      );
+    }
+    const connector: Connector = new Connector({
+      auth: new FBToolsAuthClient(),
+    });
+    const clientOpts = await connector.getOptions({
+      instanceConnectionName: connectionName,
+      ipType: IpAddressTypes.PUBLIC,
+      authType: AuthTypes.IAM,
+    });
+    const pool: pg.Pool = new pg.Pool({
+      ...clientOpts,
+      user: username,
+      database: databaseId,
+    });
+    const conn: pg.PoolClient = await pool.connect();
+
+    logger.info(`Logged in as ${username}`);
+    logger.info(chalk.cyan("Welcome to Data Connect Cloud SQL Shell"));
+    logger.info(
+      chalk.gray(
+        "Type your your SQL query or '.exit' to quit, queries should end with ';' or add empty line to execute.",
+      ),
+    );
+
+    // Start accepting queries
+    await mainShellLoop(conn);
+
+    // Cleanup after exit
+    logger.info(chalk.yellow("Exiting shell..."));
+    conn.release();
+    await pool.end();
+    connector.close();
+
+    return { projectId, serviceId };
+  });

src/commands/index.ts

@@ -2,10 +2,10 @@
 /**
  * Loads all commands for our parser.
  */
 export function load(client: any): any {
   function loadCommand(name: string) {
     const t0 = process.hrtime.bigint();
     const { command: cmd } = require(`./${name}`);
     cmd.register(client);
     const t1 = process.hrtime.bigint();
     const diffMS = (t1 - t0) / BigInt(1e6);
@@ -215,6 +215,7 @@
   client.dataconnect.sql.diff = loadCommand("dataconnect-sql-diff");
   client.dataconnect.sql.migrate = loadCommand("dataconnect-sql-migrate");
   client.dataconnect.sql.grant = loadCommand("dataconnect-sql-grant");
+  client.dataconnect.sql.shell = loadCommand("dataconnect-sql-shell");
   client.dataconnect.sdk = {};
   client.dataconnect.sdk.generate = loadCommand("dataconnect-sdk-generate");
   client.target = loadCommand("target");

src/dataconnect/schemaMigration.ts

@@ -282,7 +282,7 @@ function setSchemaValidationMode(schema: Schema, schemaValidation: SchemaValidat
   }
 }
 
-function getIdentifiers(schema: Schema): {
+export function getIdentifiers(schema: Schema): {
   instanceName: string;
   instanceId: string;
   databaseId: string;

src/gcp/cloudsql/interactive.ts

@@ -0,0 +1,55 @@
+import * as pg from "pg";
+import * as ora from "ora";
+import chalk from "chalk";
+import { logger } from "../../logger";
+import { confirm } from "../../prompt";
+
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const Table = require("cli-table");
+
+// Not comprehensive list, used for best offer prompting.
+const destructiveSqlKeywords = ["DROP", "DELETE"];
+
+function checkIsDestructiveSql(query: string): boolean {
+  const upperCaseQuery = query.toUpperCase();
+  return destructiveSqlKeywords.some((keyword) => upperCaseQuery.includes(keyword.toUpperCase()));
+}
+
+export async function confirmDangerousQuery(query: string): Promise<boolean> {
+  if (checkIsDestructiveSql(query)) {
+    return await confirm({
+      message: chalk.yellow("This query may be destructive. Are you sure you want to proceed?"),
+      default: false,
+    });
+  }
+  return true;
+}
+
+// Pretty query execution display such as spinner and actual returned content for `SELECT` query.
+export async function interactiveExecuteQuery(query: string, conn: pg.PoolClient) {
+  const spinner = ora("Executing query...").start();
+  try {
+    const results = await conn.query(query);
+    spinner.succeed(chalk.green("Query executed successfully"));
+
+    if (Array.isArray(results.rows) && results.rows.length > 0) {
+      const table: any[] = new Table({
+        head: Object.keys(results.rows[0]).map((key) => chalk.cyan(key)),
+        style: { head: [], border: [] },
+      });
+
+      for (const row of results.rows) {
+        table.push(Object.values(row) as any);
+      }
+
+      logger.info(table.toString());
+    } else {
+      // If nothing is returned and the query was select, let the user know there was no results.
+      if (query.toUpperCase().includes("SELECT")) {
+        logger.info(chalk.yellow("No results returned"));
+      }
+    }
+  } catch (err) {
+    spinner.fail(chalk.red(`Failed executing query: ${err}`));
+  }
+}