From 383767908cceb15a5322e5582bb331fafe026a1c Mon Sep 17 00:00:00 2001
From: Mihai Claudiu Toader <mihai@sigmacomputing.com>
Date: Thu, 29 Dec 2022 00:46:04 -0800
Subject: [PATCH] Don't cache non success responses (#109)

Don't cache non success responses and debug the body response
---
 monitoring.go | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/monitoring.go b/monitoring.go
index ada13669e..2e125bdd0 100644
--- a/monitoring.go
+++ b/monitoring.go
@@ -3,10 +3,12 @@
 package gosnowflake
 
 import (
+	"bytes"
 	"context"
 	"database/sql/driver"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/url"
 	"runtime"
 	"strconv"
@@ -207,14 +209,13 @@ func (sc *snowflakeConn) getQueryResultResp(
 	resultPath string,
 ) (*execResponse, error) {
 	var cachedResponse *execResponse
-	cachedResponse = nil
-	if respd, ok := sc.execRespCache.load(resultPath); ok {
-		cachedResponse = respd
+	if res, ok := sc.execRespCache.load(resultPath); ok {
 		// return the cached response, unless we pass the flag saying to
 		// bypass the cache
-		if !shouldSkipCache(ctx) {
-			return respd, nil
+		if res.Success && !shouldSkipCache(ctx) {
+			return res, nil
 		}
+		cachedResponse = res
 	}
 
 	headers := getHeaders()
@@ -237,13 +238,25 @@ func (sc *snowflakeConn) getQueryResultResp(
 		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
 		return nil, err
 	}
-	defer res.Body.Close()
+	if res.Body != nil {
+		defer func() { _ = res.Body.Close() }()
+	}
+
+	bodyBytes, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+
 	var respd *execResponse
-	if err = json.NewDecoder(res.Body).Decode(&respd); err != nil {
+	if err = json.NewDecoder(bytes.NewReader(bodyBytes)).Decode(&respd); err != nil {
 		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
 		return nil, err
 	}
 
+	if !respd.Success && respd.Code == "" && respd.Message == "" {
+		logger.WithContext(ctx).Errorf("failed to build a proper exec response. received body: %s", string(bodyBytes))
+	}
+
 	// if we are skipping the cache, log difference between cached and non cached result
 	if shouldSkipCache(ctx) {
 		qid := respd.Data.QueryID
@@ -251,7 +264,6 @@ func (sc *snowflakeConn) getQueryResultResp(
 		// if there was no response in the cache anyway, log that and dont try to log anything else
 		if cachedResponse == nil {
 			logger.WithContext(ctx).Errorf("cached queryId: %v did not use cache", qid)
-
 		} else {
 			// log if there are any differences in the arrow encooded first chunk
 			arrowCached := cachedResponse.Data.RowSetBase64
@@ -285,7 +297,10 @@ func (sc *snowflakeConn) getQueryResultResp(
 		}
 	}
 
-	sc.execRespCache.store(resultPath, respd)
+	if respd.Success {
+		sc.execRespCache.store(resultPath, respd)
+	}
+
 	return respd, nil
 }
 
@@ -340,9 +355,10 @@ func (sc *snowflakeConn) waitForCompletedQueryResultResp(
 
 	if !response.Success {
 		logEverything(ctx, qid, response, startTime)
+	} else {
+		sc.execRespCache.store(resultPath, response)
 	}
 
-	sc.execRespCache.store(resultPath, response)
 	return response, nil
 }