Examples with VisibleForTesting io.crate.common.annotations.VisibleForTesting used on opensource projects

Example 21 with VisibleForTesting

use of io.crate.common.annotations.VisibleForTesting in project crate by crate.

the class RestoreSnapshotPlan method bind.

@VisibleForTesting
public static BoundRestoreSnapshot bind(AnalyzedRestoreSnapshot restoreSnapshot, CoordinatorTxnCtx txnCtx, NodeContext nodeCtx, Row parameters, SubQueryResults subQueryResults, Schemas schemas) {
    Function<? super Symbol, Object> eval = x -> SymbolEvaluator.evaluate(txnCtx, nodeCtx, x, parameters, subQueryResults);
    Settings settings = GenericPropertiesConverter.genericPropertiesToSettings(restoreSnapshot.properties().map(eval), SnapshotSettings.SETTINGS);
    HashSet<BoundRestoreSnapshot.RestoreTableInfo> restoreTables = new HashSet<>(restoreSnapshot.tables().size());
    for (Table<Symbol> table : restoreSnapshot.tables()) {
        var relationName = RelationName.of(table.getName(), txnCtx.sessionContext().searchPath().currentSchema());
        try {
            DocTableInfo docTableInfo = schemas.getTableInfo(relationName, Operation.RESTORE_SNAPSHOT);
            if (table.partitionProperties().isEmpty()) {
                throw new RelationAlreadyExists(relationName);
            }
            var partitionName = toPartitionName(docTableInfo, Lists2.map(table.partitionProperties(), x -> x.map(eval)));
            if (docTableInfo.partitions().contains(partitionName)) {
                throw new PartitionAlreadyExistsException(partitionName);
            }
            restoreTables.add(new BoundRestoreSnapshot.RestoreTableInfo(relationName, partitionName));
        } catch (RelationUnknown | SchemaUnknownException e) {
            if (table.partitionProperties().isEmpty()) {
                restoreTables.add(new BoundRestoreSnapshot.RestoreTableInfo(relationName, null));
            } else {
                var partitionName = toPartitionName(relationName, Lists2.map(table.partitionProperties(), x -> x.map(eval)));
                restoreTables.add(new BoundRestoreSnapshot.RestoreTableInfo(relationName, partitionName));
            }
        }
    }
    return new BoundRestoreSnapshot(restoreSnapshot.repository(), restoreSnapshot.snapshot(), restoreTables, restoreSnapshot.includeTables(), restoreSnapshot.includeCustomMetadata(), restoreSnapshot.customMetadataTypes(), restoreSnapshot.includeGlobalSettings(), restoreSnapshot.globalSettings(), settings);
}

Example 22 with VisibleForTesting

use of io.crate.common.annotations.VisibleForTesting in project crate by crate.

the class RestoreSnapshotPlan method resolveTableFromSnapshots.

@VisibleForTesting
static void resolveTableFromSnapshots(BoundRestoreSnapshot.RestoreTableInfo table, List<SnapshotInfo> snapshots, ResolveIndicesAndTemplatesContext context) throws RelationUnknown {
    String name = table.tableIdent().indexNameOrAlias();
    for (SnapshotInfo snapshot : snapshots) {
        for (String index : snapshot.indices()) {
            if (name.equals(index)) {
                context.addIndex(index);
                return;
            } else if (isIndexPartitionOfTable(index, table.tableIdent())) {
                String templateName = table.partitionTemplate();
                // add a partitions wildcard
                // to match all partitions if a partitioned table was meant
                context.addIndex(templateName + "*");
                context.addTemplate(templateName);
                return;
            }
        }
    }
    context.addTemplate(table.partitionTemplate());
}

Example 23 with VisibleForTesting

use of io.crate.common.annotations.VisibleForTesting in project crate by crate.

the class RestoreSnapshotPlan method resolveIndexNames.

@VisibleForTesting
static CompletableFuture<ResolveIndicesAndTemplatesContext> resolveIndexNames(String repositoryName, Set<BoundRestoreSnapshot.RestoreTableInfo> restoreTables, boolean ignoreUnavailable, TransportGetSnapshotsAction transportGetSnapshotsAction) {
    ResolveIndicesAndTemplatesContext context = new ResolveIndicesAndTemplatesContext();
    ArrayList<BoundRestoreSnapshot.RestoreTableInfo> toResolveFromSnapshot = new ArrayList<>();
    for (var table : restoreTables) {
        if (table.hasPartitionInfo()) {
            context.addIndex(table.partitionName().asIndexName());
            context.addTemplate(table.partitionTemplate());
        } else if (ignoreUnavailable) {
            // If ignoreUnavailable is true, it's cheaper to simply
            // return indexName and the partitioned wildcard instead
            // checking if it's a partitioned table or not
            context.addIndex(table.tableIdent().indexNameOrAlias());
            // For the case its a partitioned table we restore all partitions and the templates
            String templateName = table.partitionTemplate();
            context.addIndex(templateName + "*");
            context.addTemplate(templateName);
        } else {
            // index name needs to be resolved from snapshot
            toResolveFromSnapshot.add(table);
        }
    }
    if (toResolveFromSnapshot.isEmpty()) {
        return CompletableFuture.completedFuture(context);
    } else {
        FutureActionListener<GetSnapshotsResponse, ResolveIndicesAndTemplatesContext> listener = new FutureActionListener<>(response -> {
            resolveTablesFromSnapshots(toResolveFromSnapshot, response.getSnapshots(), context);
            return context;
        });
        transportGetSnapshotsAction.execute(new GetSnapshotsRequest(repositoryName), listener);
        return listener;
    }
}

Example 24 with VisibleForTesting

use of io.crate.common.annotations.VisibleForTesting in project crate by crate.

the class WindowAgg method create.

@VisibleForTesting
public static LogicalPlan create(LogicalPlan source, List<WindowFunction> windowFunctions) {
    if (windowFunctions.isEmpty()) {
        return source;
    }
    LinkedHashMap<WindowDefinition, ArrayList<WindowFunction>> groupedFunctions = new LinkedHashMap<>();
    for (WindowFunction windowFunction : windowFunctions) {
        WindowDefinition windowDefinition = windowFunction.windowDefinition();
        ArrayList<WindowFunction> functions = groupedFunctions.computeIfAbsent(windowDefinition, w -> new ArrayList<>());
        functions.add(windowFunction);
    }
    LogicalPlan lastWindowAgg = source;
    for (Map.Entry<WindowDefinition, ArrayList<WindowFunction>> entry : groupedFunctions.entrySet()) {
        /*
             * Pass along the source outputs as standalone symbols as they might be required in cases like:
             *      select x, avg(x) OVER() from t;
             */
        ArrayList<WindowFunction> functions = entry.getValue();
        WindowDefinition windowDefinition = entry.getKey();
        OrderBy orderBy = windowDefinition.orderBy();
        if (orderBy == null || lastWindowAgg.outputs().containsAll(orderBy.orderBySymbols())) {
            lastWindowAgg = new WindowAgg(lastWindowAgg, windowDefinition, functions, lastWindowAgg.outputs());
        } else {
            // ``WindowProjector.createUpdateProbeValueFunction` expects that all OrderBY symbols are `InputColumn`
            // Here we have a case where there is a function or something in the orderBy expression that is *not*
            // already provided by the source.
            // -> Inject `eval` so that the `orderBy` of the window-function will turn into a `InputColumn`
            Eval eval = new Eval(lastWindowAgg, Lists2.concatUnique(lastWindowAgg.outputs(), orderBy.orderBySymbols()));
            lastWindowAgg = new WindowAgg(eval, windowDefinition, functions, eval.outputs());
        }
    }
    return lastWindowAgg;
}

Example 25 with VisibleForTesting

use of io.crate.common.annotations.VisibleForTesting in project crate by crate.

the class S3FileInput method toPreGlobUri.

@VisibleForTesting
@Nullable
static S3URI toPreGlobUri(S3URI uri) {
    Matcher hasGlobMatcher = HAS_GLOBS_PATTERN.matcher(uri.toString());
    S3URI preGlobUri = null;
    if (hasGlobMatcher.matches()) {
        preGlobUri = S3URI.toS3URI(URI.create(hasGlobMatcher.group(1)));
    }
    return preGlobUri;
}