SxssfStreamingToS3Example.java

import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;

import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.streaming.SXSSFWorkbook;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.core.async.AsyncRequestBody;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3AsyncClient;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.S3Exception;

public class SxssfStreamingToS3Example {

    public static void main(String[] args) {
        // Specify your S3 bucket name and key
        String bucketName = "your-bucket-name";
        String key = "example.xlsx";

        // Create a stream generator to write the workbook to S3
        Consumer<OutputStream> streamGenerator = outputStream -> {
            try {
                generateWorkbookInChunks(outputStream);
            } catch (IOException e) {
                // Handle the exception
                e.printStackTrace();
            }
        };

        // Upload the workbook to S3
        uploadWorkbookToS3(bucketName, key, streamGenerator);
    }

    private static void generateWorkbookInChunks(OutputStream outputStream) throws IOException {
        WorkbookGenerator workbookGenerator = new WorkbookGenerator();
        SXSSFWorkbook workbook = new SXSSFWorkbook();

        workbookGenerator.generate(workbook::createSheet);

        workbook.write(outputStream);
        outputStream.flush();

        workbook.dispose();
    }

    private static void uploadWorkbookToS3(String bucketName, String key, Consumer<OutputStream> streamGenerator) {
        // Create an S3AsyncClient using the default credentials provider chain
        S3AsyncClient s3Client = S3AsyncClient.builder()
                .region(Region.US_EAST_1)
                .credentialsProvider(DefaultCredentialsProvider.create())
                .build();

        try {
            // Create a PutObjectRequest to upload the workbook to S3
            PutObjectRequest request = PutObjectRequest.builder()
                    .bucket(bucketName)
                    .key(key)
                    .contentType("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
                    .build();

            // Upload the workbook to S3 using an asynchronous request
            s3Client.putObject(request, AsyncRequestBody.fromPublisher(streamGenerator::accept))
                    .join(); // Wait for the upload to complete

            System.out.println("Workbook uploaded to S3 successfully!");
        } catch (S3Exception e) {
            // Handle the exception
            e.printStackTrace();
        } finally {
            // Close the S3 client
            s3Client.close();
        }
    }

    private static class WorkbookGenerator {
        private static final int ROW_CHUNK_SIZE = 100;
        private static final int CELL_CHUNK_SIZE = 1000;

        private final AtomicBoolean done = new AtomicBoolean(false);

        public void generate(Supplier<Sheet> sheetSupplier) {
            Sheet sheet = sheetSupplier.get();

            // Generate the workbook content in chunks
            for (int rowChunk = 0; !done.get(); rowChunk++) {
                for (int cellChunk = 0; !done.get() && cellChunk < CELL_CHUNK_SIZE; cellChunk++) {
                    int rowIndex = rowChunk * ROW_CHUNK_SIZE + cellChunk / sheet.getRowBreak();
                    int cellIndex = cellChunk % sheet.getRowBreak();

                    Row row = sheet.getRow(rowIndex);
                    if (row == null) {
                        row = sheet.createRow(rowIndex);
                    }

                    Cell cell = row.createCell(cellIndex);
                    cell.setCellValue("Value"); // Replace with your logic to generate cell values
                }
                try {
                    Thread.sleep(100); // Simulate some processing time between chunks
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }
        }

        public void finish() {
            done.set(true);
        }
    }
}