SNOW-1924252: Support internal flag (#1319)

Author: sfc-gh-jzana

connection.go

@@ -258,8 +258,9 @@ func (sc *snowflakeConn) BeginTx(
 		return nil, driver.ErrBadConn
 	}
 	isDesc := isDescribeOnly(ctx)
+	isInternal := isInternal(ctx)
 	if _, err := sc.exec(ctx, "BEGIN", false, /* noResult */
-		false /* isInternal */, isDesc, nil); err != nil {
+		isInternal, isDesc, nil); err != nil {
 		return nil, err
 	}
 	return &snowflakeTx{sc, ctx}, nil
@@ -318,9 +319,9 @@ func (sc *snowflakeConn) ExecContext(
 	}
 	noResult := isAsyncMode(ctx)
 	isDesc := isDescribeOnly(ctx)
-	// TODO handle isInternal
+	isInternal := isInternal(ctx)
 	ctx = setResultType(ctx, execResultType)
-	data, err := sc.exec(ctx, query, noResult, false /* isInternal */, isDesc, args)
+	data, err := sc.exec(ctx, query, noResult, isInternal, isDesc, args)
 	if err != nil {
 		logger.WithContext(ctx).Infof("error: %v", err)
 		if data != nil {
@@ -407,8 +408,8 @@ func (sc *snowflakeConn) queryContextInternal(
 	noResult := isAsyncMode(ctx)
 	isDesc := isDescribeOnly(ctx)
 	ctx = setResultType(ctx, queryResultType)
-	// TODO: handle isInternal
-	data, err := sc.exec(ctx, query, noResult, false /* isInternal */, isDesc, args)
+	isInternal := isInternal(ctx)
+	data, err := sc.exec(ctx, query, noResult, isInternal, isDesc, args)
 	if err != nil {
 		logger.WithContext(ctx).Errorf("error: %v", err)
 		if data != nil {
@@ -475,9 +476,9 @@ func (sc *snowflakeConn) Ping(ctx context.Context) error {
 	}
 	noResult := isAsyncMode(ctx)
 	isDesc := isDescribeOnly(ctx)
-	// TODO: handle isInternal
+	isInternal := isInternal(ctx)
 	ctx = setResultType(ctx, execResultType)
-	_, err := sc.exec(ctx, "SELECT 1", noResult, false, /* isInternal */
+	_, err := sc.exec(ctx, "SELECT 1", noResult, isInternal,
 		isDesc, []driver.NamedValue{})
 	return err
 }
@@ -518,7 +519,8 @@ func (sc *snowflakeConn) QueryArrowStream(ctx context.Context, query string, bin
 	ctx = WithArrowBatches(context.WithValue(ctx, asyncMode, false))
 	ctx = setResultType(ctx, queryResultType)
 	isDesc := isDescribeOnly(ctx)
-	data, err := sc.exec(ctx, query, false, false /* isinternal */, isDesc, bindings)
+	isInternal := isInternal(ctx)
+	data, err := sc.exec(ctx, query, false, isInternal, isDesc, bindings)
 	if err != nil {
 		logger.WithContext(ctx).Errorf("error: %v", err)
 		if data != nil {

connection_util.go

@@ -221,6 +221,15 @@ func isDescribeOnly(ctx context.Context) bool {
 	return ok && d
 }
 
+func isInternal(ctx context.Context) bool {
+	v := ctx.Value(internalQuery)
+	if v == nil {
+		return false
+	}
+	d, ok := v.(bool)
+	return ok && d
+}
+
 func setResultType(ctx context.Context, resType resultType) context.Context {
 	return context.WithValue(ctx, snowflakeResultType, resType)
 }

transaction.go

@@ -46,7 +46,8 @@ func (tx *snowflakeTx) execTxCommand(command txCommand) (err error) {
 	if tx.sc == nil || tx.sc.rest == nil {
 		return driver.ErrBadConn
 	}
-	_, err = tx.sc.exec(tx.ctx, txStr, false /* noResult */, false /* isInternal */, false /* describeOnly */, nil)
+	isInternal := isInternal(tx.ctx)
+	_, err = tx.sc.exec(tx.ctx, txStr, false /* noResult */, isInternal, false /* describeOnly */, nil)
 	if err != nil {
 		return
 	}

util.go

@@ -40,6 +40,7 @@ const (
 
 const (
 	describeOnly        contextKey = "DESCRIBE_ONLY"
+	internalQuery       contextKey = "INTERNAL_QUERY"
 	cancelRetry         contextKey = "CANCEL_RETRY"
 	streamChunkDownload contextKey = "STREAM_CHUNK_DOWNLOAD"
 )
@@ -173,6 +174,11 @@ func WithArrayValuesNullable(ctx context.Context) context.Context {
 	return context.WithValue(ctx, arrayValuesNullable, true)
 }
 
+// WithInternal sets the internal query flag.
+func WithInternal(ctx context.Context) context.Context {
+	return context.WithValue(ctx, internalQuery, true)
+}
+
 // Get the request ID from the context if specified, otherwise generate one
 func getOrGenerateRequestIDFromContext(ctx context.Context) UUID {
 	requestID, ok := ctx.Value(snowflakeRequestIDKey).(UUID)

util_test.go

@@ -433,3 +433,10 @@ func TestFindByPrefix(t *testing.T) {
 	assertEqualE(t, findByPrefix(nonEmpty, "dd"), -1)
 	assertEqualE(t, findByPrefix([]string{}, "dd"), -1)
 }
+
+func TestInternal(t *testing.T) {
+	ctx := context.Background()
+	assertFalseE(t, isInternal(ctx))
+	ctx = WithInternal(ctx)
+	assertTrueE(t, isInternal(ctx))
+}