Function queryRetry

backend/api/v1/sql_service.go:640–835

(
	ctx context.Context,
	stores *store.Store,
	user *store.UserMessage,
	instance *store.InstanceMessage,
	database *store.DatabaseMessage,
	driver db.Driver,
	conn *sql.Conn,
	statements []parserbase.Statement,
	originalStatement string,
	queryContext db.QueryContext,
	licenseService *enterprise.LicenseService,
	optionalAccessCheck accessCheckFunc,
	schemaSyncer *schemasync.Syncer,
	multiStatement bool,
)

}

func queryRetry(
	ctx context.Context,
	stores *store.Store,
	user *store.UserMessage,
	instance *store.InstanceMessage,
	database *store.DatabaseMessage,
	driver db.Driver,
	conn *sql.Conn,
	statements []parserbase.Statement,
	originalStatement string,
	queryContext db.QueryContext,
	licenseService *enterprise.LicenseService,
	optionalAccessCheck accessCheckFunc,
	schemaSyncer *schemasync.Syncer,
	multiStatement bool,
) ([]*v1pb.QueryResult, []*parserbase.QuerySpan, time.Duration, error) {
	var spans []*parserbase.QuerySpan
	var sensitivePredicateColumns [][]parserbase.ColumnResource
	var err error
	if !queryContext.Explain {
		spans, err = parserbase.GetQuerySpan(
			ctx,
			parserbase.GetQuerySpanContext{
				InstanceID:                    instance.ResourceID,
				GetDatabaseMetadataFunc:       parsercontext.BuildGetDatabaseMetadataFunc(stores),
				ListDatabaseNamesFunc:         parsercontext.BuildListDatabaseNamesFunc(stores),
				GetLinkedDatabaseMetadataFunc: parsercontext.BuildGetLinkedDatabaseMetadataFunc(stores, instance.Metadata.GetEngine()),
			},
			instance.Metadata.GetEngine(),
			statements,
			database.DatabaseName,
			queryContext.Schema,
			!store.IsObjectCaseSensitive(instance),
		)
		if err != nil {
			return nil, nil, time.Duration(0), err
		}
		// After replacing backup table with source, we can apply the original access check and mask sensitive data for backup table.
		// If err != nil, this function will return the original spans.
		if err := replaceBackupTableWithSource(ctx, stores, instance, database, spans); err != nil {
			slog.Debug("failed to replace backup table with source", log.BBError(err))
		}
		if optionalAccessCheck != nil {
			// Check query access
			if err := optionalAccessCheck(ctx, instance, database, user, spans, queryContext.Explain, statements, queryContext.Schema, multiStatement); err != nil {
				return nil, nil, time.Duration(0), err
			}
			slog.Debug("optional access check", slog.String("instance", instance.ResourceID), slog.String("database", database.DatabaseName))
		}
		if !queryContext.SkipMasking && licenseService.IsFeatureEnabledForInstance(ctx, common.GetWorkspaceIDFromContext(ctx), v1pb.PlanFeature_FEATURE_DATA_MASKING, instance) == nil {
			masker := NewQueryResultMasker(stores)
			sensitivePredicateColumns, err = masker.ExtractSensitivePredicateColumns(ctx, spans, instance, user)
			if err != nil {
				return nil, nil, time.Duration(0), connect.NewError(connect.CodeInternal, errors.New(err.Error()))
			}
			slog.Debug("extract sensitive predicate columns", slog.String("instance", instance.ResourceID), slog.String("database", database.DatabaseName))
		}
	}