Examples with QualifiedObjectName io.prestosql.spi.connector.QualifiedObjectName used on opensource projects

Example 26 with QualifiedObjectName

use of io.prestosql.spi.connector.QualifiedObjectName in project hetu-core by openlookeng.

the class SimpleNodeSelector method computeAssignments.

@Override
public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTask> existingTasks, Optional<SqlStageExecution> stage) {
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    NodeMap nodeMapSlice = this.nodeMap.get().get();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMapSlice, existingTasks);
    ResettableRandomizedIterator<InternalNode> randomCandidates = randomizedNodes(nodeMapSlice, ImmutableSet.of());
    Set<InternalNode> blockedExactNodes = new HashSet<>();
    boolean splitWaitingForAnyNode = false;
    // splitsToBeRedistributed becomes true only when splits go through locality-based assignment
    boolean splitsToBeRedistributed = false;
    Set<Split> remainingSplits = new HashSet<>();
    // Check if the current stage has a TableScanNode which is reading the table for the 2nd time or beyond
    if (stage.isPresent() && stage.get().getStateMachine().getConsumerScanNode() != null) {
        try {
            // if node exists, get the TableScanNode and cast it as consumer
            TableScanNode consumer = stage.get().getStateMachine().getConsumerScanNode();
            // all tables part of this stage
            Map<PlanNodeId, TableInfo> tables = stage.get().getStageInfo().getTables();
            QualifiedObjectName tableName;
            for (Map.Entry<PlanNodeId, TableInfo> entry : tables.entrySet()) {
                tableName = entry.getValue().getTableName();
                if (tableSplitAssignmentInfo.getReuseTableScanMappingIdSplitAssignmentMap().containsKey(consumer.getReuseTableScanMappingId())) {
                    // compare splitkey using equals and then assign nodes accordingly.
                    HashMap<SplitKey, InternalNode> splitKeyNodeAssignment = tableSplitAssignmentInfo.getSplitKeyNodeAssignment(consumer.getReuseTableScanMappingId());
                    Set<SplitKey> splitKeySet = splitKeyNodeAssignment.keySet();
                    assignment.putAll(createConsumerScanNodeAssignment(tableName, splits, splitKeySet, splitKeyNodeAssignment));
                    for (Map.Entry<InternalNode, Split> nodeAssignmentEntry : assignment.entries()) {
                        InternalNode node = nodeAssignmentEntry.getKey();
                        assignmentStats.addAssignedSplit(node);
                    }
                }
            }
            log.debug("Consumer:: Assignment size is " + assignment.size() + " ,Assignment is " + assignment + " ,Assignment Stats is " + assignmentStats);
        } catch (NotImplementedException e) {
            log.error("Not a Hive Split! Other Connector Splits not supported currently. Error: " + e);
            throw new UnsupportedOperationException("Not a Hive Split! Other Connector Splits not supported currently. Error: " + e);
        }
    } else {
        // optimizedLocalScheduling enables prioritized assignment of splits to local nodes when splits contain locality information
        if (optimizedLocalScheduling) {
            // should not hit for consumer case
            for (Split split : splits) {
                if (split.isRemotelyAccessible() && !split.getAddresses().isEmpty()) {
                    List<InternalNode> candidateNodes = selectExactNodes(nodeMapSlice, split.getAddresses(), includeCoordinator);
                    Optional<InternalNode> chosenNode = candidateNodes.stream().filter(ownerNode -> assignmentStats.getTotalSplitCount(ownerNode) < maxSplitsPerNode).min(comparingInt(assignmentStats::getTotalSplitCount));
                    if (chosenNode.isPresent()) {
                        assignment.put(chosenNode.get(), split);
                        // check later
                        assignmentStats.addAssignedSplit(chosenNode.get());
                        splitsToBeRedistributed = true;
                        continue;
                    }
                }
                remainingSplits.add(split);
            }
        } else {
            remainingSplits = splits;
        }
        for (Split split : remainingSplits) {
            randomCandidates.reset();
            List<InternalNode> candidateNodes;
            if (!split.isRemotelyAccessible()) {
                candidateNodes = selectExactNodes(nodeMapSlice, split.getAddresses(), includeCoordinator);
            } else {
                candidateNodes = selectNodes(minCandidates, randomCandidates);
            }
            if (candidateNodes.isEmpty()) {
                log.debug("No nodes available to schedule %s. Available nodes %s", split, nodeMapSlice.getNodesByHost().keys());
                throw new PrestoException(NO_NODES_AVAILABLE, "No nodes available to run query");
            }
            InternalNode chosenNode = null;
            int min = Integer.MAX_VALUE;
            for (InternalNode node : candidateNodes) {
                int totalSplitCount = assignmentStats.getTotalSplitCount(node);
                if (totalSplitCount < min && totalSplitCount < maxSplitsPerNode) {
                    chosenNode = node;
                    min = totalSplitCount;
                }
            }
            if (chosenNode == null) {
                // min is guaranteed to be MAX_VALUE at this line
                for (InternalNode node : candidateNodes) {
                    int totalSplitCount = assignmentStats.getQueuedSplitCountForStage(node);
                    if (totalSplitCount < min && totalSplitCount < maxPendingSplitsPerTask) {
                        chosenNode = node;
                        min = totalSplitCount;
                    }
                }
            }
            if (chosenNode != null) {
                assignment.put(chosenNode, split);
                assignmentStats.addAssignedSplit(chosenNode);
            } else {
                if (split.isRemotelyAccessible()) {
                    splitWaitingForAnyNode = true;
                } else // Exact node set won't matter, if a split is waiting for any node
                if (!splitWaitingForAnyNode) {
                    blockedExactNodes.addAll(candidateNodes);
                }
            }
        }
    }
    ListenableFuture<?> blocked;
    if (splitWaitingForAnyNode) {
        blocked = toWhenHasSplitQueueSpaceFuture(existingTasks, calculateLowWatermark(maxPendingSplitsPerTask));
    } else {
        blocked = toWhenHasSplitQueueSpaceFuture(blockedExactNodes, existingTasks, calculateLowWatermark(maxPendingSplitsPerTask));
    }
    if (!stage.isPresent() || stage.get().getStateMachine().getConsumerScanNode() == null) {
        if (splitsToBeRedistributed) {
            // skip for consumer
            equateDistribution(assignment, assignmentStats, nodeMapSlice);
        }
    }
    // Check if the current stage has a TableScanNode which is reading the table for the 1st time
    if (stage.isPresent() && stage.get().getStateMachine().getProducerScanNode() != null) {
        // if node exists, get the TableScanNode and annotate it as producer
        saveProducerScanNodeAssignment(stage, assignment, assignmentStats);
    }
    // Check if its CTE node and its feeder
    if (stage.isPresent() && stage.get().getFragment().getFeederCTEId().isPresent()) {
        updateFeederNodeAndSplitCount(stage.get(), assignment);
    }
    return new SplitPlacementResult(blocked, assignment);
}

Example 27 with QualifiedObjectName

use of io.prestosql.spi.connector.QualifiedObjectName in project hetu-core by openlookeng.

the class ExtractSpatialJoins method loadKdbTree.

private static KdbTree loadKdbTree(String tableName, Session session, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager, PlanNodeId nodeId) {
    QualifiedObjectName name = toQualifiedObjectName(tableName, session.getCatalog().get(), session.getSchema().get());
    TableHandle tableHandle = metadata.getTableHandle(session, name).orElseThrow(() -> new PrestoException(INVALID_SPATIAL_PARTITIONING, format("Table not found: %s", name)));
    Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
    List<ColumnHandle> visibleColumnHandles = columnHandles.values().stream().filter(handle -> !metadata.getColumnMetadata(session, tableHandle, handle).isHidden()).collect(toImmutableList());
    checkSpatialPartitioningTable(visibleColumnHandles.size() == 1, "Expected single column for table %s, but found %s columns", name, columnHandles.size());
    ColumnHandle kdbTreeColumn = Iterables.getOnlyElement(visibleColumnHandles);
    Optional<KdbTree> kdbTree = Optional.empty();
    try (SplitSource splitSource = splitManager.getSplits(session, tableHandle, UNGROUPED_SCHEDULING, null, Optional.empty(), Collections.emptyMap(), ImmutableSet.of(), false, nodeId)) {
        while (!Thread.currentThread().isInterrupted()) {
            SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000));
            List<Split> splits = splitBatch.getSplits();
            for (Split split : splits) {
                try (ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, tableHandle, ImmutableList.of(kdbTreeColumn), Optional.empty())) {
                    do {
                        getFutureValue(pageSource.isBlocked());
                        Page page = pageSource.getNextPage();
                        if (page != null && page.getPositionCount() > 0) {
                            checkSpatialPartitioningTable(!kdbTree.isPresent(), "Expected exactly one row for table %s, but found more", name);
                            checkSpatialPartitioningTable(page.getPositionCount() == 1, "Expected exactly one row for table %s, but found %s rows", name, page.getPositionCount());
                            String kdbTreeJson = VARCHAR.getSlice(page.getBlock(0), 0).toStringUtf8();
                            try {
                                kdbTree = Optional.of(KdbTreeUtils.fromJson(kdbTreeJson));
                            } catch (IllegalArgumentException e) {
                                checkSpatialPartitioningTable(false, "Invalid JSON string for KDB tree: %s", e.getMessage());
                            }
                        }
                    } while (!pageSource.isFinished());
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
            }
            if (splitBatch.isLastBatch()) {
                break;
            }
        }
    }
    checkSpatialPartitioningTable(kdbTree.isPresent(), "Expected exactly one row for table %s, but got none", name);
    return kdbTree.get();
}

Example 28 with QualifiedObjectName

use of io.prestosql.spi.connector.QualifiedObjectName in project hetu-core by openlookeng.

the class TestAccessControlManager method testReadOnlySystemAccessControl.

@Test
public void testReadOnlySystemAccessControl() {
    Identity identity = new Identity(USER_NAME, Optional.of(PRINCIPAL));
    QualifiedObjectName tableName = new QualifiedObjectName("catalog", "schema", "table");
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControlManager = new AccessControlManager(transactionManager);
    accessControlManager.setSystemAccessControl(ReadOnlySystemAccessControl.NAME, ImmutableMap.of());
    accessControlManager.checkCanSetUser(Optional.of(PRINCIPAL), USER_NAME);
    accessControlManager.checkCanSetSystemSessionProperty(identity, "property");
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanSetCatalogSessionProperty(transactionId, identity, "catalog", "property");
        accessControlManager.checkCanShowSchemas(transactionId, identity, "catalog");
        accessControlManager.checkCanShowTablesMetadata(transactionId, identity, new CatalogSchemaName("catalog", "schema"));
        accessControlManager.checkCanSelectFromColumns(transactionId, identity, tableName, ImmutableSet.of("column"));
        accessControlManager.checkCanCreateViewWithSelectFromColumns(transactionId, identity, tableName, ImmutableSet.of("column"));
        Set<String> catalogs = ImmutableSet.of("catalog");
        assertEquals(accessControlManager.filterCatalogs(identity, catalogs), catalogs);
        Set<String> schemas = ImmutableSet.of("schema");
        assertEquals(accessControlManager.filterSchemas(transactionId, identity, "catalog", schemas), schemas);
        Set<SchemaTableName> tableNames = ImmutableSet.of(new SchemaTableName("schema", "table"));
        assertEquals(accessControlManager.filterTables(transactionId, identity, "catalog", tableNames), tableNames);
    });
    try {
        transaction(transactionManager, accessControlManager).execute(transactionId -> {
            accessControlManager.checkCanInsertIntoTable(transactionId, identity, tableName);
        });
        fail();
    } catch (AccessDeniedException expected) {
    }
}

Example 29 with QualifiedObjectName

use of io.prestosql.spi.connector.QualifiedObjectName in project hetu-core by openlookeng.

the class TestAccessControlManager method testDenyCatalogAccessControl.

@Test(expectedExceptions = PrestoException.class, expectedExceptionsMessageRegExp = "Access Denied: Cannot select from columns \\[column\\] in table or view schema.table")
public void testDenyCatalogAccessControl() {
    CatalogManager catalogManager = new CatalogManager();
    TransactionManager transactionManager = createTestTransactionManager(catalogManager);
    AccessControlManager accessControlManager = new AccessControlManager(transactionManager);
    TestSystemAccessControlFactory accessControlFactory = new TestSystemAccessControlFactory("test");
    accessControlManager.addSystemAccessControlFactory(accessControlFactory);
    accessControlManager.setSystemAccessControl("test", ImmutableMap.of());
    CatalogName catalogName = registerBogusConnector(catalogManager, transactionManager, accessControlManager, "catalog");
    accessControlManager.addCatalogAccessControl(catalogName, new DenyConnectorAccessControl());
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanSelectFromColumns(transactionId, new Identity(USER_NAME, Optional.of(PRINCIPAL)), new QualifiedObjectName("catalog", "schema", "table"), ImmutableSet.of("column"));
    });
}

Example 30 with QualifiedObjectName

use of io.prestosql.spi.connector.QualifiedObjectName in project hetu-core by openlookeng.

the class TestAccessControlManager method testColumnMaskOrdering.

@Test
public void testColumnMaskOrdering() {
    CatalogManager catalogManager = new CatalogManager();
    TransactionManager transactionManager = createTestTransactionManager(catalogManager);
    AccessControlManager accessControlManager = new AccessControlManager(transactionManager);
    accessControlManager.addSystemAccessControlFactory(new SystemAccessControlFactory() {

        @Override
        public String getName() {
            return "test";
        }

        @Override
        public SystemAccessControl create(Map<String, String> config) {
            return new SystemAccessControl() {

                @Override
                public void checkCanSetUser(Optional<Principal> principal, String userName) {
                }

                @Override
                public void checkCanImpersonateUser(Identity identity, String propertyName) {
                }

                @Override
                public void checkCanSetSystemSessionProperty(Identity identity, String propertyName) {
                }

                @Override
                public Optional<ViewExpression> getColumnMask(Identity identity, CatalogSchemaTableName tableName, String columnName, Type type) {
                    return Optional.of(new ViewExpression("user", Optional.empty(), Optional.empty(), "system mask"));
                }
            };
        }
    });
    accessControlManager.setSystemAccessControl("test", ImmutableMap.of());
    CatalogName catalogName = registerBogusConnector(catalogManager, transactionManager, accessControlManager, "catalog");
    accessControlManager.addCatalogAccessControl(catalogName, new ConnectorAccessControl() {

        @Override
        public Optional<ViewExpression> getColumnMask(ConnectorTransactionHandle transactionHandle, Identity identity, SchemaTableName tableName, String columnName, Type type) {
            return Optional.of(new ViewExpression("user", Optional.empty(), Optional.empty(), "connector mask"));
        }
    });
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        List<ViewExpression> masks = accessControlManager.getColumnMasks(transactionId, new Identity(USER_NAME, Optional.of(PRINCIPAL)), new QualifiedObjectName("catalog", "schema", "table"), "column", BIGINT);
        assertEquals(masks.get(0).getExpression(), "connector mask");
        assertEquals(masks.get(1).getExpression(), "system mask");
    });
}