github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/hadoopfs/src/main/java/io/lakefs/storage/LakeFSFileSystemOutputStream.java (about)

     1  package io.lakefs.storage;
     2  
     3  import java.io.ByteArrayOutputStream;
     4  import java.io.IOException;
     5  import java.io.OutputStream;
     6  import java.net.HttpURLConnection;
     7  import java.util.concurrent.atomic.AtomicBoolean;
     8  
     9  import io.lakefs.LakeFSLinker;
    10  import org.apache.commons.codec.binary.Base64;
    11  import org.apache.commons.codec.binary.Hex;
    12  import org.apache.commons.lang3.StringUtils;
    13  
    14  /**
    15   * Handle writes into a storage URL. Will set the request Content-Length header. When closed, links
    16   * the address in lakeFS.
    17   * <p>
    18   * TODO(johnnyaug): do not hold everything in memory
    19   * TODO(johnnyaug): support multipart uploads
    20   */
    21  public class LakeFSFileSystemOutputStream extends OutputStream {
    22      private final HttpURLConnection connection;
    23      private final ByteArrayOutputStream buffer;
    24      private final AtomicBoolean isClosed = new AtomicBoolean(false);
    25      private final LakeFSLinker linker;
    26  
    27      public LakeFSFileSystemOutputStream(HttpURLConnection connection, LakeFSLinker linker)
    28              throws IOException {
    29          this.connection = connection;
    30          this.buffer = new ByteArrayOutputStream();
    31          this.linker = linker;
    32      }
    33  
    34      @Override
    35      public void write(int b) throws IOException {
    36          buffer.write(b);
    37      }
    38  
    39      @Override
    40      public void close() throws IOException {
    41          if (isClosed.getAndSet(true)) {
    42              return;
    43          }
    44          connection.setRequestProperty("Content-Length", String.valueOf(buffer.size()));
    45          connection.setRequestProperty("x-ms-blob-type", "BlockBlob");
    46          OutputStream out = connection.getOutputStream();
    47          out.write(buffer.toByteArray());
    48          out.close();
    49          String eTag = extractETag();
    50          if (eTag == null) {
    51              throw new IOException("Failed to finish writing to presigned link. No ETag found.");
    52          }
    53          linker.link(eTag, buffer.size());
    54          if (connection.getResponseCode() > 299) {
    55              throw new IOException("Failed to finish writing to presigned link. Response code: "
    56                      + connection.getResponseCode());
    57          }
    58      }
    59  
    60      /**
    61       * Extracts the ETag from the response headers. Use Content-MD5 if present, otherwise use ETag.
    62       * @return the ETag of the uploaded object, or null if not found
    63       */
    64      private String extractETag() {
    65          String contentMD5 = connection.getHeaderField("Content-MD5");
    66          if (contentMD5 != null) {
    67              //noinspection VulnerableCodeUsages
    68              byte[] dataMD5 = Base64.decodeBase64(contentMD5);
    69              return Hex.encodeHexString(dataMD5);
    70          }
    71  
    72          String eTag = connection.getHeaderField("ETag");
    73          if (eTag != null) {
    74              return StringUtils.strip(eTag, "\" ");
    75          }
    76          return null;
    77      }
    78  }