Small refactoring to support 'HttpError'.

Azure · Apr 28, 2022 · ca8b845 · ca8b845
1 parent 5936442
commit ca8b845
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 27 deletions.
diff --git a/kusto/conn.go b/kusto/conn.go
@@ -183,12 +183,8 @@ func (c *conn) execute(ctx context.Context, execType int, db string, query Stmt,
 		return execResp{}, err
 	}
 
-	if resp.StatusCode == http.StatusTooManyRequests {
-		return execResp{}, errors.HTTPErrorCode(op, resp.StatusCode, body, fmt.Sprintf("request got throttled for query %q: ", query.String()))
-	}
-
 	if resp.StatusCode != http.StatusOK {
-		return execResp{}, errors.HTTP(op, resp.Status, body, fmt.Sprintf("error from Kusto endpoint for query %q: ", query.String()))
+		return execResp{}, errors.HTTP(op, resp.Status, resp.StatusCode, body, fmt.Sprintf("error from Kusto endpoint for query %q: ", query.String()))
 	}
 
 	var dec frames.Decoder

diff --git a/kusto/data/errors/errors.go b/kusto/data/errors/errors.go
@@ -18,7 +18,6 @@ import (
 	"io/ioutil"
 	"net/http"
 	"runtime"
-	"strconv"
 	"strings"
 )
 
@@ -76,8 +75,9 @@ type Error struct {
 	inner *Error
 }
 
+type KustoError = Error
 type HttpError struct {
-	err        Error
+	KustoError
 	StatusCode int
 }
 
@@ -224,29 +224,24 @@ func ES(o Op, k Kind, s string, args ...interface{}) *Error {
 }
 
 // HTTP constructs an *Error from an *http.Response and a prefix to the error message.
-func HTTP(o Op, status string, body io.ReadCloser, prefix string) *Error {
+func HTTP(o Op, status string, statusCode int, body io.ReadCloser, prefix string) *HttpError {
 	bodyBytes, err := ioutil.ReadAll(body)
 	if err != nil {
 		bodyBytes = []byte(fmt.Sprintf("Failed to read body: %v", err))
 	}
 
-	e := &Error{
-		Op:         o,
-		Kind:       KHTTPError,
-		restErrMsg: bodyBytes,
-		Err:        fmt.Errorf("%s(%s):\n%s", prefix, status, string(bodyBytes)),
+	e := HttpError{
+		KustoError: KustoError{
+			Op:         o,
+			Kind:       KHTTPError,
+			restErrMsg: bodyBytes,
+			Err:        fmt.Errorf("%s(%s):\n%s", prefix, status, string(bodyBytes)),
+		},
+		StatusCode: statusCode,
 	}
 	e.UnmarshalREST()
-	return e
-}
 
-func HTTPErrorCode(o Op, status int, body io.ReadCloser, prefix string) *HttpError {
-	err := HTTP(o, strconv.Itoa(status), body, prefix)
-	httpError := &HttpError{
-		StatusCode: status,
-		err:        *err,
-	}
-	return httpError
+	return &e
 }
 
 // e constructs an Error. You may pass in an Op, Kind, string or error.  This will strip an *Error if you
@@ -385,7 +380,3 @@ func oneToErr(m map[string]interface{}, err *Error, op Op) *Error {
 func (e *HttpError) IsThrottled() bool {
 	return e != nil && (e.StatusCode == http.StatusTooManyRequests)
 }
-
-func (e *HttpError) Error() string {
-	return e.err.Error()
-}
diff --git a/kusto/ingest/internal/conn/conn.go b/kusto/ingest/internal/conn/conn.go
@@ -167,7 +167,7 @@ func (c *Conn) StreamIngest(ctx context.Context, db, table string, payload io.Re
 		if err != nil {
 			return err
 		}
-		return errors.HTTP(writeOp, resp.Status, body, "streaming ingest issue")
+		return errors.HTTP(writeOp, resp.Status, resp.StatusCode, body, "streaming ingest issue")
 	}
 	return nil
 }