D-Net

package eu.dnetlib.data.objectstore.s3;

import com.amazonaws.ClientConfiguration;

import com.amazonaws.Protocol;

import com.amazonaws.auth.AWSCredentials;

import com.amazonaws.auth.AWSStaticCredentialsProvider;

import com.amazonaws.auth.BasicAWSCredentials;

import com.amazonaws.client.builder.AwsClientBuilder;

import com.amazonaws.services.s3.AmazonS3;

import com.amazonaws.services.s3.AmazonS3ClientBuilder;

import com.amazonaws.services.s3.model.ListObjectsV2Request;

import com.amazonaws.services.s3.model.ListObjectsV2Result;

import com.amazonaws.services.s3.model.S3Object;

import com.amazonaws.services.s3.model.S3ObjectSummary;

import com.google.common.collect.Lists;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.model.Filters;

import eu.dnetlib.data.objectstore.modular.ObjectStoreRecord;

import eu.dnetlib.data.objectstore.modular.connector.ObjectStore;

import eu.dnetlib.data.objectstore.rmi.MetadataObjectRecord;

import eu.dnetlib.data.objectstore.rmi.ObjectStoreFile;

import eu.dnetlib.data.objectstore.rmi.ObjectStoreServiceException;

import eu.dnetlib.enabling.resultset.ResultSetListener;

import org.apache.commons.lang3.exception.ExceptionUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.Document;

import java.io.ByteArrayInputStream;

import java.io.IOException;

import java.io.InputStream;

import java.nio.file.Files;

import java.nio.file.Path;

import java.nio.file.Paths;

import java.nio.file.StandardCopyOption;

import java.util.UUID;

import java.util.concurrent.atomic.AtomicInteger;

import java.util.regex.Pattern;

public class S3ObjectStore implements ObjectStore {

    //CONSTANT VARIABLE NAME

    private static final String S3_REGION = "eu-west-3";

    private static final String URI_FIELD = "uri";

    private static final String ID_FIELD = "id";

    private static final String MIME_FIELD = "mime";

    private static final String ORIGINAL_OBJECT_FIELD = "originalObject";

    private static final String TIMESTAMP_FIELD = "timestamp";

    private static final String MD5_FIELD = "md5Sum";

    private static final String SIZE_FIELD = "size";

    private final String id;

    private final String interpretation;

    private final String s3AccessKey;

    private final String s3SecretKey;

    private final String s3EndPoint;

    private final String objectStoreBucket;

    private final String basePath;

    private final MongoCollection<Document> mongoCollection;

//    private AmazonS3 client;

    private static final Log log = LogFactory.getLog(S3ObjectStore.class);

    public S3ObjectStore(final String identifier, final String interpretation, final String s3AccessKey, final String s3SecretKey, final String s3EndPoint, final String objectStoreBucket, final MongoCollection<Document> mongoCollection, final String basePath) throws ObjectStoreServiceException {

        this.id = identifier;

        this.interpretation = interpretation;

        this.s3AccessKey = s3AccessKey;

        this.s3SecretKey = s3SecretKey;

        this.mongoCollection = mongoCollection;

        this.s3EndPoint = s3EndPoint;

        this.objectStoreBucket = objectStoreBucket;

        this.basePath = basePath;

    }

    private AmazonS3 initializeClient() throws ObjectStoreServiceException {

        try {

            final AWSCredentials credentials = new BasicAWSCredentials(this.s3AccessKey, this.s3SecretKey);

            final ClientConfiguration cfg = new ClientConfiguration().withProtocol(Protocol.HTTPS);

            final AmazonS3 s3 = AmazonS3ClientBuilder.standard().withCredentials(new AWSStaticCredentialsProvider(credentials))

                    .withClientConfiguration(cfg)

                    .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(this.s3EndPoint, S3_REGION))

                    .build();

            if (s3 == null)

                throw new ObjectStoreServiceException("Cannot initialize s3 client because is NULL");

            return s3;

        } catch (Throwable e) {

            log.error("An Error happen on initialize client ", e);

            throw new ObjectStoreServiceException("Cannot initialize s3 client", e);

        }

    }

    @Override

    public String getId() {

        return this.id;

    }

    @Override

    public String getInterpretation() {

        return this.interpretation;

    }

    @Override

    public int feed(Iterable<ObjectStoreRecord> iterable, boolean upsert) throws ObjectStoreServiceException {

        final AtomicInteger count = new AtomicInteger();

        iterable.forEach(objectStoreRecord -> {

            try {

                feedObjectRecord(objectStoreRecord);

                count.incrementAndGet();

            } catch (ObjectStoreServiceException e) {

                log.error("Error on saving file in a temporary Folder");

            }

        });

        return count.intValue();

    }

    @Override

    public int feedMetadataRecord(Iterable<MetadataObjectRecord> iterable, boolean b) throws ObjectStoreServiceException {

        final AtomicInteger count = new AtomicInteger();

        iterable.forEach(mor -> {

            final ObjectStoreRecord r = new ObjectStoreRecord();

            r.setInputStream(new ByteArrayInputStream(mor.getRecord().getBytes()));

            final ObjectStoreFile fileMetadata = new ObjectStoreFile();

            fileMetadata.setObjectID(mor.getId());

            fileMetadata.setMimeType(mor.getMime());

            r.setFileMetadata(fileMetadata);

            try {

                feedObjectRecord(r);

                count.incrementAndGet();

            } catch (ObjectStoreServiceException e) {

                log.error("Unable to store record r", e);

            }

        });

        return count.intValue();

    }

    @Override

    public String feedObjectRecord(ObjectStoreRecord objectStoreRecord) throws ObjectStoreServiceException {

        final AmazonS3 client = initializeClient();

        try {

            long start = System.currentTimeMillis();

            final Path result = saveAndGenerateMD5(objectStoreRecord.getInputStream(), basePath);

            log.debug("Total time to download into fileSystem " + (System.currentTimeMillis() - start));

            if (result!= null) {

                log.debug("Saving object with ID: " + objectStoreRecord.getFileMetadata().getObjectID() + " on s3 ");

                client.putObject(objectStoreBucket, id + "/" + objectStoreRecord.getFileMetadata().getObjectID(), result.toFile());

                log.debug("Total time to put into the ObjectStore " + (System.currentTimeMillis() - start));

                log.debug("Deleting file");

                boolean deleteSuccess = Files.deleteIfExists(result);

                log.debug("Temporary file deleting success " + deleteSuccess);

                log.debug("Saved object on s3 ");

                final S3Object s3Object = client.getObject(objectStoreBucket ,id+ "/" + objectStoreRecord.getFileMetadata().getObjectID());

                double timestamp = System.currentTimeMillis();

                Document metadata = new Document()

                        .append(ID_FIELD, objectStoreRecord.getFileMetadata().getObjectID())

                        .append(MIME_FIELD, objectStoreRecord.getFileMetadata().getMimeType())

                        .append(ORIGINAL_OBJECT_FIELD, objectStoreRecord.getFileMetadata().toJSON())

                        .append(TIMESTAMP_FIELD, timestamp)

                        .append(MD5_FIELD, s3Object.getObjectMetadata().getETag())

                        .append(SIZE_FIELD, s3Object.getObjectMetadata().getContentLength())

                        .append(URI_FIELD, String.format("s3://%s/%s/%s", objectStoreBucket, id, objectStoreRecord.getFileMetadata().getObjectID()));

                log.debug("saving metadata object to the collection: " + metadata.toString());

                start = System.currentTimeMillis();

                mongoCollection.insertOne(metadata);

                log.debug("Total time to save in Mongo " + (System.currentTimeMillis() - start));

            }

        } catch (Throwable e) {

            log.error("Error on put file in the objectStore", e);

            log.error(ExceptionUtils.getStackTrace(e));

            throw new ObjectStoreServiceException(e);

        }

        return null;

    }

    @Override

    public ResultSetListener deliver(final Long from, final Long until) throws ObjectStoreServiceException {

        S3ObjectStoreResultSetListener resultSet = new S3ObjectStoreResultSetListener();

        resultSet.setMongoCollection(this.mongoCollection);

        resultSet.setObjectStoreID(getId());

        resultSet.setFromDate(from);

        resultSet.setUntilDate(until);

        return resultSet;

    }

    @Override

    public ResultSetListener deliverIds(final Iterable<String> ids) throws ObjectStoreServiceException {

        S3ObjectStoreResultSetListener resultSet = new S3ObjectStoreResultSetListener();

        resultSet.setMongoCollection(this.mongoCollection);

        resultSet.setObjectStoreID(getId());

        resultSet.setRecords(Lists.newArrayList(ids));

        return resultSet;

    }

    @Override

    public ObjectStoreFile deliverObject(String objectId) throws ObjectStoreServiceException {

        Document resultQuery = this.mongoCollection.find(Filters.eq("id", objectId)).first();

        if (resultQuery!= null)

            return ObjectStoreS3Utility.build(resultQuery);

        else return null;

    }

    @Override

    public int getSize() throws ObjectStoreServiceException {

        return (int) this.mongoCollection.count();

    }

    @Override

    public void deleteObject(String objectId) throws ObjectStoreServiceException {

        final Document response =this.mongoCollection.findOneAndDelete(Filters.eq("id", objectId));

        if (response == null)

            throw new ObjectStoreServiceException("Error document not found with objectId: "+objectId);

        final AmazonS3 client = initializeClient();

        client.deleteObject(this.objectStoreBucket, String.format("%s/%s", this.id, response.get(ID_FIELD)));

    }

    @Override

    public String getObject(String objectId) throws ObjectStoreServiceException {

        Document response = this.mongoCollection.find(Filters.eq("id", objectId)).first();

        if (response == null || !response.containsKey(URI_FIELD))

            return null;

        return response.getString(URI_FIELD);

    }

    @Override

    public boolean existIDStartsWith(String startId) throws ObjectStoreServiceException {

        return this.mongoCollection.count(Filters.regex("id", Pattern.compile(startId))) > 0;

    }

    @Override

    public boolean dropContent() throws ObjectStoreServiceException {

        final AmazonS3 client = initializeClient();

        final ListObjectsV2Request req = new ListObjectsV2Request().withBucketName(objectStoreBucket).withPrefix(id);

        ListObjectsV2Result result;

        do {

            result = client.listObjectsV2(req);

            for (S3ObjectSummary objectSummary : result.getObjectSummaries()) {

                log.debug(String.format(" - %s (size: %d)\n", objectSummary.getKey(), objectSummary.getSize()));

                client.deleteObject(objectStoreBucket, objectSummary.getKey());

                log.debug("Object Deleted");

            }

            String token = result.getNextContinuationToken();

            log.debug("Next Continuation Token: " + token);

            req.setContinuationToken(token);

        } while (result.isTruncated());

        this.mongoCollection.drop();

        return true;

    }

    private static Path saveAndGenerateMD5(final InputStream inputStream, final String basePath) throws IOException {

        if (inputStream == null) return null;

        Path tempFile = Paths.get(basePath,  UUID.randomUUID().toString());

        log.debug("Temp file On Default Location: " + tempFile.toString());

        long size = 0;

        try {

            Files.copy(inputStream, tempFile, StandardCopyOption.REPLACE_EXISTING);

        } catch (IOException e1) {

            log.error(e1);

        }

        return tempFile;

    }

}