Log execution status/state and fix flaw in cancel logic.

Logging status/state returned by executing a statement.
The logic for handling context cancel/timeout was not working correctly in the case where there are no direct results returned from execute statement.
Signed-off-by: Raymond Cypher <[email protected]>

Author: rcypher-databricks

connection.go

@@ -317,7 +317,7 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 			return false
 		}
 		hasHandle := resp.OperationHandle != nil
-		isOpen := resp.DirectResults != nil && resp.DirectResults.CloseOperation == nil
+		isOpen := resp.DirectResults == nil || resp.DirectResults.CloseOperation == nil
 		return hasHandle && isOpen
 	}
 

internal/client/client.go

@@ -165,6 +165,7 @@ func (tsc *ThriftServiceClient) ExecuteStatement(ctx context.Context, req *cli_s
 	// We use context.Background to fix a problem where on context done the query would not be cancelled.
 	resp, err := tsc.TCLIServiceClient.ExecuteStatement(context.Background(), req)
 	log, ctx = LoggerAndContext(ctx, resp)
+	logExecStatementState(resp, log)
 	defer log.Duration(msg, start)
 	if err != nil {
 		err = handleClientMethodError(ctx, err)
@@ -431,6 +432,20 @@ func guidFromHasSessionHandle(c any) (guid string) {
 	return
 }
 
+func logExecStatementState(resp *cli_service.TExecuteStatementResp, log *logger.DBSQLLogger) {
+	if resp != nil {
+		if resp.DirectResults != nil {
+			state := resp.DirectResults.GetOperationStatus().GetOperationState()
+			log.Debug().Msgf("execute statement state: %s", state)
+			status := resp.DirectResults.GetOperationStatus().GetStatus().StatusCode
+			log.Debug().Msgf("execute statement status: %s", status)
+		} else {
+			status := resp.GetStatus().StatusCode
+			log.Debug().Msgf("execute statement status: %s", status)
+		}
+	}
+}
+
 var retryableStatusCodes = map[int]any{http.StatusTooManyRequests: struct{}{}, http.StatusServiceUnavailable: struct{}{}}
 
 func isRetryableServerResponse(resp *http.Response) bool {