ClickHouseRetentionQueryExecutor.java

package org.rakam.clickhouse.analysis;

import com.facebook.presto.sql.RakamSqlFormatter;
import com.facebook.presto.sql.tree.Expression;
import com.google.inject.Inject;
import org.rakam.analysis.RetentionQueryExecutor;
import org.rakam.analysis.metadata.Metastore;
import org.rakam.config.ProjectConfig;
import org.rakam.report.DelegateQueryExecution;
import org.rakam.report.QueryExecution;
import org.rakam.report.QueryExecutor;
import org.rakam.util.ValidationUtil;

import java.time.LocalDate;
import java.time.ZoneId;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static com.facebook.presto.sql.RakamExpressionFormatter.formatIdentifier;
import static java.lang.String.format;
import static java.time.format.DateTimeFormatter.ISO_DATE;
import static org.rakam.util.ValidationUtil.checkCollection;
import static org.rakam.util.ValidationUtil.checkTableColumn;

public class ClickHouseRetentionQueryExecutor
        implements RetentionQueryExecutor
{
    private final QueryExecutor executor;
    private final Metastore metastore;
    private final ProjectConfig projectConfig;

    @Inject
    public ClickHouseRetentionQueryExecutor(ProjectConfig projectConfig, QueryExecutor executor, Metastore metastore)
    {
        this.projectConfig = projectConfig;
        this.executor = executor;
        this.metastore = metastore;
    }

    private static String formatExpression(Expression value)
    {
        return RakamSqlFormatter.formatExpression(value,
                name -> name.getParts().stream().map(e -> formatIdentifier(e, '`')).collect(Collectors.joining(".")),
                name -> name.getParts().stream()
                        .map(e -> formatIdentifier(e, '`')).collect(Collectors.joining(".")), '`');
    }

    @Override
    public QueryExecution query(String project,
            Optional<RetentionAction> firstAction,
            Optional<RetentionAction> returningAction,
            DateUnit dateUnit,
            Optional<String> dimension,
            Optional<Integer> period,
            LocalDate startDate,
            LocalDate endDate, ZoneId zoneId,
            boolean approximate)
    {
        int startEpoch = (int) startDate.toEpochDay();
        int endEpoch = (int) endDate.toEpochDay();

        String firstActionQuery = firstAction.map(action -> format("SELECT `$date`, %s, %s %s FROM %s.%s %s",
                checkTableColumn(projectConfig.getUserColumn(), '`'), checkTableColumn(projectConfig.getTimeColumn(), '`'),
                dimension.map(e -> "," + e).orElse(""), project, ValidationUtil.checkCollection(action.collection(), '`'),
                action.filter().map(f -> "WHERE " + formatExpression(f)).orElse("")))
                .orElseGet(() -> metastore.getCollectionNames(project).stream()
                        .map(collection -> format("SELECT `$date`, %s, %s %s FROM %s.%s",
                                checkTableColumn(projectConfig.getUserColumn(), '`'), checkTableColumn(projectConfig.getTimeColumn(), '`'),
                                dimension.map(e -> "," + e).orElse(""), project, ValidationUtil.checkCollection(collection, '`')))
                        .collect(Collectors.joining(" UNION ALL ")));

        String query = IntStream.range(startEpoch, endEpoch).boxed().flatMap(epoch -> {
            LocalDate beginDate = LocalDate.ofEpochDay(epoch);
            String date = startDate.format(ISO_DATE);
            String endDateStr = period.map(e -> beginDate.plus(e, dateUnit.getTemporalUnit()))
                    .map(e -> e.isAfter(endDate) ? endDate : e)
                    .orElse(LocalDate.ofEpochDay(endEpoch)).format(ISO_DATE);

            return Stream.of(format("select %s, CAST(-1 AS Int64) as lead, uniq(_user) users from" +
                            " (%s) where `$date` between toDate('%s')  and toDate('%s') " +
                            " group by `$date` %s",
                    dimension.map(e -> checkTableColumn(e, '`')).orElse(format("toDate('%s') as date", date)),
                    firstActionQuery,
                    date, endDateStr, dimension.map(e -> "," + checkCollection(e, '`')).orElse("")),
                    format("select %s, (`$date` - toDate('%s')) - 1 as lead, sum(_user IN (select _user from (%s) WHERE `$date` = toDate('%s'))) as users from" +
                                    " (%s) where `$date` between toDate('%s') and toDate('%s') " +
                                    " group by `$date` %s order by %s",
                            dimension.map(e -> checkTableColumn(e, '`')).orElse(format("toDate('%s') as date", date)),
                            date,
                            firstActionQuery,
                            date,
                            firstActionQuery,
                            date, endDateStr,
                            dimension.map(e -> ", " + checkCollection(e, '`')).orElse(""),
                            dimension.map(e -> checkCollection(e, '`')).orElse("`$date`")));
        }).collect(Collectors.joining(" UNION ALL \n"));

        return new DelegateQueryExecution(executor.executeRawQuery(query), (result) -> {
            if (result.isFailed()) {
                return result;
            }

            Long uniqUser = new Long(-1);

            result.getResult().stream()
                    .filter(objects -> objects.get(1).equals(uniqUser))
                    .forEach(objects -> objects.set(1, null));

            return result;
        });
    }
}