From f20f276ce68cb36eda775176307df9d71cdd9ac5 Mon Sep 17 00:00:00 2001
From: Rasmus Porsager <rasmus@porsager.com>
Date: Mon, 26 Jun 2023 00:39:02 +0200
Subject: [PATCH] Add query stats

---
 src/connection.js |  2 ++
 src/index.js      |  4 ++++
 src/query.js      | 44 +++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/connection.js b/src/connection.js
index e8e4881d..897ed3f3 100644
--- a/src/connection.js
+++ b/src/connection.js
@@ -165,6 +165,8 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
         : (query = q, query.active = true)
 
       build(q)
+      q.statistics && (q.statistics.executed = performance.now())
+      q.handler.onquery && (q.handler.onquery = q.handler.onquery(q))
       return write(toBuffer(q))
         && !q.describeFirst
         && !q.cursorFn
diff --git a/src/index.js b/src/index.js
index ff990586..9ceefc59 100644
--- a/src/index.js
+++ b/src/index.js
@@ -85,6 +85,8 @@ function Postgres(a, b) {
 
   function Sql(handler) {
     handler.debug = options.debug
+    handler.stats = options.stats
+    handler.onquery = options.onquery
 
     Object.entries(options.types).reduce((acc, [name, type]) => {
       acc[name] = (x) => new Parameter(x, type.to)
@@ -486,6 +488,8 @@ function parseOptions(a, b) {
     onclose         : o.onclose,
     onparameter     : o.onparameter,
     socket          : o.socket,
+    stats           : o.stats,
+    onquery         : o.onquery,
     transform       : parseTransform(o.transform || { undefined: undefined }),
     parameters      : {},
     shared          : { retries: 0, typeArrayMap: {} },
diff --git a/src/query.js b/src/query.js
index 0d44a15c..9725285d 100644
--- a/src/query.js
+++ b/src/query.js
@@ -13,6 +13,10 @@ export class Query extends Promise {
       reject = b
     })
 
+    this.resolver = resolve
+    this.rejecter = reject
+
+    this.statistics = handler.stats || handler.debug ? { started: -1, executed: -1 } : undefined
     this.tagged = Array.isArray(strings.raw)
     this.strings = strings
     this.args = args
@@ -23,19 +27,30 @@ export class Query extends Promise {
     this.state = null
     this.statement = null
 
-    this.resolve = x => (this.active = false, resolve(x))
-    this.reject = x => (this.active = false, reject(x))
-
     this.active = false
     this.cancelled = null
     this.executed = false
     this.signature = ''
 
-    this[originError] = this.handler.debug
+    this[originError] = handler.debug
       ? new Error()
       : this.tagged && cachedError(this.strings)
   }
 
+  resolve(x) {
+    this.active = false
+    this.statistics && addStats(this, x)
+    this.handler.onquery && (this.handler.onquery = this.handler.onquery(x))
+    this.resolver(x)
+  }
+
+  reject(x) {
+    this.active = false
+    this.statistics && addStats(this, x)
+    this.handler.onquery && (this.handler.onquery = this.handler.onquery(x))
+    this.rejecter(x)
+  }
+
   get origin() {
     return (this.handler.debug
       ? this[originError].stack
@@ -131,13 +146,25 @@ export class Query extends Promise {
     return this
   }
 
+  stats() {
+    this.statistics = { started: -1, executed: -1 }
+    return this
+  }
+
   values() {
     this.isRaw = 'values'
     return this
   }
 
   async handle() {
-    !this.executed && (this.executed = true) && await 1 && this.handler(this)
+    if (this.executed)
+      return
+
+    this.executed = true
+    await 1
+    this.statistics && (this.statistics.started = performance.now())
+    this.handler.onquery && (this.handler.onquery = this.handler.onquery(this))
+    this.handler(this)
   }
 
   execute() {
@@ -171,3 +198,10 @@ function cachedError(xs) {
   Error.stackTraceLimit = x
   return originCache.get(xs)
 }
+
+
+function addStats(query, result) {
+  result.waiting = query.statistics.executed - query.statistics.started
+  result.duration = performance.now() - query.statistics.started
+  result.execution = performance.now() - query.statistics.executed
+}