com.palantir.example.profile.schema.generated.UserPhotosMetadataCleanupTask Maven / Gradle / Ivy

Go to download
package com.palantir.example.profile.schema.generated;

import java.util.Iterator;
import java.util.Map;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

import com.google.common.collect.Sets;
import com.palantir.atlasdb.cleaner.api.OnCleanupTask;
import com.palantir.atlasdb.encoding.PtBytes;
import com.palantir.atlasdb.keyvalue.api.BatchColumnRangeSelection;
import com.palantir.atlasdb.keyvalue.api.Cell;
import com.palantir.atlasdb.keyvalue.api.Namespace;
import com.palantir.atlasdb.protos.generated.StreamPersistence.Status;
import com.palantir.atlasdb.protos.generated.StreamPersistence.StreamMetadata;
import com.palantir.atlasdb.transaction.api.Transaction;

public class UserPhotosMetadataCleanupTask implements OnCleanupTask {

    private final ProfileTableFactory tables;

    public UserPhotosMetadataCleanupTask(Namespace namespace) {
        tables = ProfileTableFactory.of(namespace);
    }

    @Override
    public boolean cellsCleanedUp(Transaction t, Set cells) {
        UserPhotosStreamMetadataTable metaTable = tables.getUserPhotosStreamMetadataTable(t);
        Set rows = Sets.newHashSetWithExpectedSize(cells.size());
        for (Cell cell : cells) {
            rows.add(UserPhotosStreamMetadataTable.UserPhotosStreamMetadataRow.BYTES_HYDRATOR.hydrateFromBytes(cell.getRowName()));
        }
        UserPhotosStreamIdxTable indexTable = tables.getUserPhotosStreamIdxTable(t);
        Set rowsWithNoIndexEntries =
                        getUnreferencedStreamsByIterator(indexTable, rows);
        Set toDelete = new HashSet<>();
        Map currentMetadata =
                metaTable.getMetadatas(rows);
        for (Map.Entry e : currentMetadata.entrySet()) {
            if (e.getValue().getStatus() != Status.STORED || rowsWithNoIndexEntries.contains(e.getKey())) {
                toDelete.add(e.getKey().getId());
            }
        }
        UserPhotosStreamStore.of(tables).deleteStreams(t, toDelete);
        return false;
    }

    private static Set getUnreferencedStreamsByIterator(UserPhotosStreamIdxTable indexTable, Set metadataRows) {
        Set indexRows = metadataRows.stream()
                .map(UserPhotosStreamMetadataTable.UserPhotosStreamMetadataRow::getId)
                .map(UserPhotosStreamIdxTable.UserPhotosStreamIdxRow::of)
                .collect(Collectors.toSet());
        Map> referenceIteratorByStream
                = indexTable.getRowsColumnRangeIterator(indexRows,
                        BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1));
        return referenceIteratorByStream.entrySet().stream()
                .filter(entry -> !entry.getValue().hasNext())
                .map(Map.Entry::getKey)
                .map(UserPhotosStreamIdxTable.UserPhotosStreamIdxRow::getId)
                .map(UserPhotosStreamMetadataTable.UserPhotosStreamMetadataRow::of)
                .collect(Collectors.toSet());
    }
}