From a7124535e38f5762e30312ae075935612ae9a1e8 Mon Sep 17 00:00:00 2001
From: uki00a <uki00a@gmail.com>
Date: Thu, 21 May 2020 22:04:25 +0900
Subject: [PATCH] feat: add support for getting result metadata (#96)

---
 connection.ts    | 10 +++++++++
 query.ts         | 28 ++++++++++++++++++++++++
 tests/queries.ts | 55 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+)

diff --git a/connection.ts b/connection.ts
index 94b50360..3fd3c00d 100644
--- a/connection.ts
+++ b/connection.ts
@@ -299,6 +299,8 @@ export class Connection {
       // command complete
       // TODO: this is duplicated in next loop
       case "C":
+        const commandTag = this._readCommandTag(msg);
+        result.handleCommandComplete(commandTag);
         result.done();
         break;
       default:
@@ -316,6 +318,8 @@ export class Connection {
           break;
         // command complete
         case "C":
+          const commandTag = this._readCommandTag(msg);
+          result.handleCommandComplete(commandTag);
           result.done();
           break;
         // ready for query
@@ -505,6 +509,8 @@ export class Connection {
           break;
         // command complete
         case "C":
+          const commandTag = this._readCommandTag(msg);
+          result.handleCommandComplete(commandTag);
           result.done();
           break outerLoop;
         // error response
@@ -569,6 +575,10 @@ export class Connection {
     return row;
   }
 
+  _readCommandTag(msg: Message) {
+    return msg.reader.readString(msg.byteCount);
+  }
+
   async initSQL(): Promise<void> {
     const config: QueryConfig = { text: "select 1;", args: [] };
     const query = new Query(config);
diff --git a/query.ts b/query.ts
index dea4cef8..58cf30d2 100644
--- a/query.ts
+++ b/query.ts
@@ -4,6 +4,18 @@ import { encode, EncodedArg } from "./encode.ts";
 
 import { decode } from "./decode.ts";
 
+const commandTagRegexp = /^([A-Za-z]+)(?: (\d+))?(?: (\d+))?/;
+
+type CommandType = (
+  | "INSERT"
+  | "DELETE"
+  | "UPDATE"
+  | "SELECT"
+  | "MOVE"
+  | "FETCH"
+  | "COPY"
+);
+
 export interface QueryConfig {
   text: string;
   args?: Array<unknown>;
@@ -15,6 +27,8 @@ export class QueryResult {
   public rowDescription!: RowDescription;
   private _done = false;
   public rows: any[] = []; // actual results
+  public rowCount?: number;
+  public command!: CommandType;
 
   constructor(public query: Query) {}
 
@@ -48,6 +62,20 @@ export class QueryResult {
     this.rows.push(parsedRow);
   }
 
+  handleCommandComplete(commandTag: string): void {
+    const match = commandTagRegexp.exec(commandTag);
+    if (match) {
+      this.command = match[1] as CommandType;
+      if (match[3]) {
+        // COMMAND OID ROWS
+        this.rowCount = parseInt(match[3], 10);
+      } else {
+        // COMMAND ROWS
+        this.rowCount = parseInt(match[2], 10);
+      }
+    }
+  }
+
   rowsOfObjects() {
     return this.rows.map((row) => {
       const rv: { [key: string]: any } = {};
diff --git a/tests/queries.ts b/tests/queries.ts
index 89782ebc..30980374 100644
--- a/tests/queries.ts
+++ b/tests/queries.ts
@@ -2,6 +2,7 @@ import { Client } from "../mod.ts";
 import { assertEquals } from "../test_deps.ts";
 import { DEFAULT_SETUP, TEST_CONNECTION_PARAMS } from "./constants.ts";
 import { getTestClient } from "./helpers.ts";
+import { QueryResult } from "../query.ts";
 
 const CLIENT = new Client(TEST_CONNECTION_PARAMS);
 
@@ -136,3 +137,57 @@ testClient(async function multiQueryWithManyQueryTypeArray() {
 
   assertEquals(result[2].rows.length, 2);
 });
+
+testClient(async function resultMetadata() {
+  let result: QueryResult;
+
+  // simple select
+  result = await CLIENT.query("SELECT * FROM ids WHERE id = 100");
+  assertEquals(result.command, "SELECT");
+  assertEquals(result.rowCount, 1);
+
+  // parameterized select
+  result = await CLIENT.query(
+    "SELECT * FROM ids WHERE id IN ($1, $2)",
+    200,
+    300,
+  );
+  assertEquals(result.command, "SELECT");
+  assertEquals(result.rowCount, 2);
+
+  // simple delete
+  result = await CLIENT.query("DELETE FROM ids WHERE id IN (100, 200)");
+  assertEquals(result.command, "DELETE");
+  assertEquals(result.rowCount, 2);
+
+  // parameterized delete
+  result = await CLIENT.query("DELETE FROM ids WHERE id = $1", 300);
+  assertEquals(result.command, "DELETE");
+  assertEquals(result.rowCount, 1);
+
+  // simple insert
+  result = await CLIENT.query("INSERT INTO ids VALUES (4), (5)");
+  assertEquals(result.command, "INSERT");
+  assertEquals(result.rowCount, 2);
+
+  // parameterized insert
+  result = await CLIENT.query("INSERT INTO ids VALUES ($1)", 3);
+  assertEquals(result.command, "INSERT");
+  assertEquals(result.rowCount, 1);
+
+  // simple update
+  result = await CLIENT.query(
+    "UPDATE ids SET id = 500 WHERE id IN (500, 600)",
+  );
+  assertEquals(result.command, "UPDATE");
+  assertEquals(result.rowCount, 2);
+
+  // parameterized update
+  result = await CLIENT.query("UPDATE ids SET id = 400 WHERE id = $1", 400);
+  assertEquals(result.command, "UPDATE");
+  assertEquals(result.rowCount, 1);
+}, [
+  "DROP TABLE IF EXISTS ids",
+  "CREATE UNLOGGED TABLE ids (id integer)",
+  "INSERT INTO ids VALUES (100), (200), (300), (400), (500), (600)",
+]);