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

     1  package io.lakefs.storage;
     2  
     3  import java.io.IOException;
     4  import java.io.OutputStream;
     5  import java.net.URI;
     6  import java.util.concurrent.atomic.AtomicBoolean;
     7  import com.amazonaws.services.s3.model.ObjectMetadata;
     8  import io.lakefs.LakeFSLinker;
     9  
    10  /**
    11   * Wraps a FSDataOutputStream to link file on staging when done writing
    12   */
    13  class LinkOnCloseOutputStream extends OutputStream {
    14      private final URI physicalUri;
    15      private final MetadataClient metadataClient;
    16      private final OutputStream out;
    17      private final AtomicBoolean isLinked = new AtomicBoolean(false);
    18      private final LakeFSLinker linker;
    19  
    20      /**
    21       * @param physicalUri translated physical location of object data for underlying FileSystem.
    22       * @param metadataClient client used to request metadata information from the underlying FileSystem.
    23       * @param out stream on underlying filesystem to wrap.
    24       * @param linker {@link LakeFSLinker} for the given object.
    25       */
    26      LinkOnCloseOutputStream(URI physicalUri, MetadataClient metadataClient, OutputStream out,
    27              LakeFSLinker linker) {
    28          this.physicalUri = physicalUri;
    29          this.metadataClient = metadataClient;
    30          this.out = out;
    31          this.linker = linker;
    32      }
    33  
    34      @Override
    35      public void flush() throws IOException {
    36          out.flush();
    37      }
    38  
    39      @Override
    40      public void write(byte[] b) throws IOException {
    41          out.write(b);
    42      }
    43  
    44      @Override
    45      public void write(byte[] b, int off, int len) throws IOException {
    46          out.write(b, off, len);
    47      }
    48  
    49      @Override
    50      public void write(int b) throws IOException {
    51          out.write(b);
    52      }
    53  
    54      @Override
    55      public void close() throws IOException {
    56          out.close();
    57          // Now the object is on the underlying store, find its parameters (sadly lost by
    58          // the underlying Hadoop FileSystem) so we can link it on lakeFS.
    59          if (!this.isLinked.getAndSet(true)) {
    60              ObjectMetadata objectMetadata = metadataClient.getObjectMetadata(physicalUri);
    61              linker.link(objectMetadata.getETag(), objectMetadata.getContentLength());
    62          }
    63      }
    64  }