diff --git a/clients/databricks/store.go b/clients/databricks/store.go
index daa39f7b5..ac07b5863 100644
--- a/clients/databricks/store.go
+++ b/clients/databricks/store.go
@@ -119,7 +119,7 @@ func (s Store) PrepareTemporaryTable(ctx context.Context, tableData *optimizatio
 	}()
 
 	// Upload the local file to DBFS
-	ctx = driverctx.NewContextWithStagingInfo(context.Background(), []string{"/var"})
+	ctx = driverctx.NewContextWithStagingInfo(ctx, []string{"/var"})
 
 	castedTempTableID, isOk := tempTableID.(TableIdentifier)
 	if !isOk {