org.apache.hadoop.hive.ql.hooks.ReadEntity Java Examples

private Set<ReadEntity> getInputs(String inputName, Entity.Type entityType) throws HiveException {
    final ReadEntity entity = new ReadEntity();

    if ( Entity.Type.DFS_DIR.equals(entityType)) {
        entity.setName(lower(new Path(inputName).toString()));
        entity.setTyp(Entity.Type.DFS_DIR);
    } else {
        entity.setName(getQualifiedTblName(inputName));
        entity.setTyp(entityType);
    }

    if (entityType == Entity.Type.TABLE) {
        entity.setT(hiveMetaStoreBridge.hiveClient.getTable(DEFAULT_DB, inputName));
    }

    return new LinkedHashSet<ReadEntity>() {{ add(entity); }};
}
 

/**
 * Add column level hierarchy to inputHierarchy
 *
 * @param inputHierarchy
 * @param entity
 * @param sentryContext
 */
private void addColumnHierarchy(List<List<DBModelAuthorizable>> inputHierarchy,
    ReadEntity entity) {
  List<DBModelAuthorizable> entityHierarchy = new ArrayList<DBModelAuthorizable>();
  entityHierarchy.add(hiveAuthzBinding.getAuthServer());
  entityHierarchy.addAll(getAuthzHierarchyFromEntity(entity));

  switch (entity.getType()) {
  case TABLE:
  case PARTITION:
    List<String> cols = entity.getAccessedColumns();
    for (String col : cols) {
      List<DBModelAuthorizable> colHierarchy = new ArrayList<DBModelAuthorizable>(entityHierarchy);
      colHierarchy.add(new Column(col));
      inputHierarchy.add(colHierarchy);
    }
    break;
  default:
    inputHierarchy.add(entityHierarchy);
  }
}
 

@Test
public void testCTAS() throws Exception {
    String tableName     = createTable();
    String ctasTableName = "table" + random();
    String query         = "create table " + ctasTableName + " as select * from " + tableName;

    runCommand(query);

    final Set<ReadEntity> readEntities = getInputs(tableName, Entity.Type.TABLE);
    final Set<WriteEntity> writeEntities = getOutputs(ctasTableName, Entity.Type.TABLE);

    HiveEventContext hiveEventContext = constructEvent(query, HiveOperation.CREATETABLE_AS_SELECT, readEntities,
            writeEntities);
    AtlasEntity processEntity1 = validateProcess(hiveEventContext);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, hiveEventContext);
    AtlasObjectId process = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
            BaseHiveEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process.getGuid(), processEntity1.getGuid());

    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);
    assertTableIsRegistered(DEFAULT_DB, ctasTableName);
}
 

@Override
public Task<? extends Serializable> createShowRolePrincipalsTask(ASTNode ast, Path resFile,
    HashSet<ReadEntity> inputs, HashSet<WriteEntity> outputs) throws SemanticException {
  String roleName;

  if (ast.getChildCount() == 1) {
    roleName = ast.getChild(0).getText();
  } else {
    // the parser should not allow this
    throw new AssertionError("Unexpected Tokens in SHOW ROLE PRINCIPALS");
  }

  RoleDDLDesc roleDDLDesc = new RoleDDLDesc(roleName, PrincipalType.ROLE,
   RoleDDLDesc.RoleOperation.SHOW_ROLE_PRINCIPALS, null);
  roleDDLDesc.setResFile(resFile.toString());
  return createTask(new DDLWork(inputs, outputs, roleDDLDesc));
  //return TaskFactory.get(new DDLWork(inputs, outputs, roleDDLDesc), conf);
}
 

/**
 * Get Authorizable from inputs and put into inputHierarchy
 *
 * @param inputHierarchy
 * @param entity
 * @param sentryContext
 */
private void getInputHierarchyFromInputs(List<List<DBModelAuthorizable>> inputHierarchy,
    Set<ReadEntity> inputs) {
  for (ReadEntity readEntity: inputs) {
    // skip the tables/view that are part of expanded view definition
    // skip the Hive generated dummy entities created for queries like 'select <expr>'
    if (isChildTabForView(readEntity) || isDummyEntity(readEntity)) {
      continue;
    }
    if (readEntity.getAccessedColumns() != null && !readEntity.getAccessedColumns().isEmpty()) {
      addColumnHierarchy(inputHierarchy, readEntity);
    } else {
      List<DBModelAuthorizable> entityHierarchy = new ArrayList<DBModelAuthorizable>();
      entityHierarchy.add(hiveAuthzBinding.getAuthServer());
      entityHierarchy.addAll(getAuthzHierarchyFromEntity(readEntity));
      inputHierarchy.add(entityHierarchy);
    }
  }
}
 

private Task<? extends Serializable> analyzeGrantRevokeRole(boolean isGrant, ASTNode ast,
    HashSet<ReadEntity> inputs, HashSet<WriteEntity> outputs) throws SemanticException {
  List<PrincipalDesc> principalDesc = analyzePrincipalListDef(
      (ASTNode) ast.getChild(0));

  List<String> roles = new ArrayList<String>();
  for (int i = 1; i < ast.getChildCount(); i++) {
    roles.add(BaseSemanticAnalyzer.unescapeIdentifier(ast.getChild(i).getText()));
  }
  String roleOwnerName = "";
  if (SessionState.get() != null
      && SessionState.get().getAuthenticator() != null) {
    roleOwnerName = SessionState.get().getAuthenticator().getUserName();
  }
  for (PrincipalDesc princ : principalDesc) {
    if (princ.getType() != PrincipalType.GROUP) {
      String msg = SentryHiveConstants.GRANT_REVOKE_NOT_SUPPORTED_ON_OBJECT + princ.getType();
      throw new SemanticException(msg);
    }
  }
  GrantRevokeRoleDDL grantRevokeRoleDDL = new GrantRevokeRoleDDL(isGrant,
      roles, principalDesc, roleOwnerName, PrincipalType.USER, false);
  return createTask(new DDLWork(inputs, outputs, grantRevokeRoleDDL));
}
 

/**
 * Check if the given read entity is a table that has parents of type Table
 * Hive compiler performs a query rewrite by replacing view with its definition. In the process, tt captures both
 * the original view and the tables/view that it selects from .
 * The access authorization is only interested in the top level views and not the underlying tables.
 * @param readEntity
 * @return
 */
private boolean isChildTabForView(ReadEntity readEntity) {
  // If this is a table added for view, then we need to skip that
  if (!readEntity.getType().equals(Type.TABLE) && !readEntity.getType().equals(Type.PARTITION)) {
    return false;
  }
  if (readEntity.getParents() != null && readEntity.getParents().size() > 0) {
    for (ReadEntity parentEntity : readEntity.getParents()) {
      if (!parentEntity.getType().equals(Type.TABLE)) {
        return false;
      }
    }
    return true;
  } else {
    return false;
  }
}
 

public SentryOnFailureHookContextImpl(String command,
    Set<ReadEntity> inputs, Set<WriteEntity> outputs, HiveOperation hiveOp,
    Database db, Table tab, AccessURI udfURI, AccessURI partitionURI,
    String userName, String ipAddress, AuthorizationException e,
    Configuration conf) {
  this.command = command;
  this.inputs = inputs;
  this.outputs = outputs;
  this.hiveOp = hiveOp;
  this.userName = userName;
  this.ipAddress = ipAddress;
  this.database = db;
  this.table = tab;
  this.udfURI = udfURI;
  this.partitionURI = partitionURI;
  this.authException = e;
  this.conf = conf;
}
 

@Override
public Task<? extends Serializable> createRevokeTask(ASTNode ast, HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs) throws SemanticException {
  List<PrivilegeDesc> privilegeDesc = analyzePrivilegeListDef((ASTNode) ast.getChild(0));
  List<PrincipalDesc> principalDesc = analyzePrincipalListDef((ASTNode) ast.getChild(1));
  PrivilegeObjectDesc privilegeObj = null;
  if (ast.getChildCount() > 2) {
    ASTNode astChild = (ASTNode) ast.getChild(2);
    privilegeObj = analyzePrivilegeObject(astChild);
  }
  if (privilegeObj != null && privilegeObj.getPartSpec() != null) {
    throw new SemanticException(SentryHiveConstants.PARTITION_PRIVS_NOT_SUPPORTED);
  }
  for (PrincipalDesc princ : principalDesc) {
    if (princ.getType() != PrincipalType.ROLE) {
      String msg = SentryHiveConstants.GRANT_REVOKE_NOT_SUPPORTED_FOR_PRINCIPAL + princ.getType();
      throw new SemanticException(msg);
    }
  }
  RevokeDesc revokeDesc = new RevokeDesc(privilegeDesc, principalDesc, privilegeObj);
  return createTask(new DDLWork(inputs, outputs, revokeDesc));
}
 

@Override
public Task<? extends Serializable> createShowRoleGrantTask(ASTNode ast, Path resultFile,
    HashSet<ReadEntity> inputs, HashSet<WriteEntity> outputs) throws SemanticException {
  ASTNode child = (ASTNode) ast.getChild(0);
  PrincipalType principalType = PrincipalType.USER;
  switch (child.getType()) {
  case HiveParser.TOK_USER:
    principalType = PrincipalType.USER;
    break;
  case HiveParser.TOK_GROUP:
    principalType = PrincipalType.GROUP;
    break;
  case HiveParser.TOK_ROLE:
    principalType = PrincipalType.ROLE;
    break;
  }
  if (principalType != PrincipalType.GROUP) {
    String msg = SentryHiveConstants.GRANT_REVOKE_NOT_SUPPORTED_FOR_PRINCIPAL + principalType;
    throw new SemanticException(msg);
  }
  String principalName = BaseSemanticAnalyzer.unescapeIdentifier(child.getChild(0).getText());
  RoleDDLDesc roleDesc = new RoleDDLDesc(principalName, principalType,
      RoleDDLDesc.RoleOperation.SHOW_ROLE_GRANT, null);
  roleDesc.setResFile(resultFile.toString());
  return createTask(new DDLWork(inputs, outputs,  roleDesc));
}
 

private void assertProcessIsNotRegistered(HiveHook.HiveEventContext event) throws Exception {
    try {
        SortedSet<ReadEntity> sortedHiveInputs = event.getInputs() == null ? null : new TreeSet<ReadEntity>(entityComparator);
        SortedSet<WriteEntity> sortedHiveOutputs = event.getOutputs() == null ? null : new TreeSet<WriteEntity>(entityComparator);
        if ( event.getInputs() != null) {
            sortedHiveInputs.addAll(event.getInputs());
        }
        if ( event.getOutputs() != null) {
            sortedHiveOutputs.addAll(event.getOutputs());
        }
        String processQFName = getProcessQualifiedName(hiveMetaStoreBridge, event, sortedHiveInputs, sortedHiveOutputs, getSortedProcessDataSets(event.getInputs()), getSortedProcessDataSets(event.getOutputs()));
        LOG.debug("Searching for process with query {}", processQFName);
        assertEntityIsNotRegistered(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, processQFName);
    } catch( Exception e) {
        LOG.error("Exception : ", e);
    }
}
 

private String assertProcessIsRegistered(final HiveHook.HiveEventContext event, final Set<ReadEntity> inputTbls, final Set<WriteEntity> outputTbls) throws Exception {
    try {
        SortedSet<ReadEntity> sortedHiveInputs = event.getInputs() == null ? null : new TreeSet<ReadEntity>(entityComparator);
        SortedSet<WriteEntity> sortedHiveOutputs = event.getOutputs() == null ? null : new TreeSet<WriteEntity>(entityComparator);
        if ( event.getInputs() != null) {
            sortedHiveInputs.addAll(event.getInputs());
        }
        if ( event.getOutputs() != null) {
            sortedHiveOutputs.addAll(event.getOutputs());
        }
        String processQFName = getProcessQualifiedName(hiveMetaStoreBridge, event, sortedHiveInputs, sortedHiveOutputs, getSortedProcessDataSets(inputTbls), getSortedProcessDataSets(outputTbls));
        LOG.debug("Searching for process with query {}", processQFName);
        return assertEntityIsRegistered(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, processQFName, new AssertPredicate() {
            @Override
            public void assertOnEntity(final Referenceable entity) throws Exception {
                List<String> recentQueries = (List<String>) entity.get("recentQueries");
                Assert.assertEquals(recentQueries.get(0), lower(event.getQueryStr()));
            }
        });
    } catch(Exception e) {
        LOG.error("Exception : ", e);
        throw e;
    }
}
 

private void assertProcessIsNotRegistered(HiveEventContext event) throws Exception {
    try {
        SortedSet<ReadEntity>  sortedHiveInputs  = event.getInputs() == null ? null : new TreeSet<ReadEntity>(entityComparator);
        SortedSet<WriteEntity> sortedHiveOutputs = event.getOutputs() == null ? null : new TreeSet<WriteEntity>(entityComparator);

        if (event.getInputs() != null) {
            sortedHiveInputs.addAll(event.getInputs());
        }

        if (event.getOutputs() != null) {
            sortedHiveOutputs.addAll(event.getOutputs());
        }

        String processQFName = getProcessQualifiedName(hiveMetaStoreBridge, event, sortedHiveInputs, sortedHiveOutputs, getSortedProcessDataSets(event.getInputs()), getSortedProcessDataSets(event.getOutputs()));

        LOG.debug("Searching for process with query {}", processQFName);

        assertEntityIsNotRegistered(HiveDataTypes.HIVE_PROCESS.getName(), ATTRIBUTE_QUALIFIED_NAME, processQFName);
    } catch(Exception e) {
        LOG.error("Exception : ", e);
    }
}
 

@Test(enabled = false)
public void testInsertIntoTempTable() throws Exception {
    String tableName = createTable();
    String insertTableName = createTable(false, false, true);
    assertTableIsRegistered(DEFAULT_DB, tableName);
    assertTableIsNotRegistered(DEFAULT_DB, insertTableName, true);

    String query =
        "insert into " + insertTableName + " select id, name from " + tableName;

    runCommand(query);

    Set<ReadEntity> inputs = getInputs(tableName, Entity.Type.TABLE);
    Set<WriteEntity> outputs = getOutputs(insertTableName, Entity.Type.TABLE);
    outputs.iterator().next().setName(getQualifiedTblName(insertTableName + HiveMetaStoreBridge.TEMP_TABLE_PREFIX + SessionState.get().getSessionId()));
    outputs.iterator().next().setWriteType(WriteEntity.WriteType.INSERT);

    validateProcess(constructEvent(query,  HiveOperation.QUERY, inputs, outputs));

    assertTableIsRegistered(DEFAULT_DB, tableName);
    assertTableIsRegistered(DEFAULT_DB, insertTableName, null, true);
}
 

private Referenceable validateProcess(HiveHook.HiveEventContext event, Set<ReadEntity> inputTables, Set<WriteEntity> outputTables) throws Exception {
    String processId = assertProcessIsRegistered(event, inputTables, outputTables);
    Referenceable process = atlasClient.getEntity(processId);
    if (inputTables == null) {
        Assert.assertNull(process.get(INPUTS));
    } else {
        Assert.assertEquals(((List<Referenceable>) process.get(INPUTS)).size(), inputTables.size());
        validateInputTables(process, inputTables);
    }

    if (outputTables == null) {
        Assert.assertNull(process.get(OUTPUTS));
    } else {
        Assert.assertEquals(((List<Id>) process.get(OUTPUTS)).size(), outputTables.size());
        validateOutputTables(process, outputTables);
    }

    return process;
}
 

private static void addInputs(HiveMetaStoreBridge hiveBridge, HiveOperation op, SortedSet<ReadEntity> sortedInputs, StringBuilder buffer, final Map<ReadEntity, Referenceable> refs, final boolean ignoreHDFSPathsInQFName) throws HiveException {
    if (refs != null) {
        if (sortedInputs != null) {
            Set<String> dataSetsProcessed = new LinkedHashSet<>();
            for (Entity input : sortedInputs) {

                if (!dataSetsProcessed.contains(input.getName().toLowerCase())) {
                    //HiveOperation.QUERY type encompasses INSERT, INSERT_OVERWRITE, UPDATE, DELETE, PATH_WRITE operations
                    if (ignoreHDFSPathsInQFName &&
                        (Type.DFS_DIR.equals(input.getType()) || Type.LOCAL_DIR.equals(input.getType()))) {
                        LOG.debug("Skipping dfs dir input addition to process qualified name {} ", input.getName());
                    } else if (refs.containsKey(input)) {
                        if ( input.getType() == Type.PARTITION || input.getType() == Type.TABLE) {
                            final Date createTime = HiveMetaStoreBridge.getTableCreatedTime(hiveBridge.hiveClient.getTable(input.getTable().getDbName(), input.getTable().getTableName()));
                            addDataset(buffer, refs.get(input), createTime.getTime());
                        } else {
                            addDataset(buffer, refs.get(input));
                        }
                    }
                    dataSetsProcessed.add(input.getName().toLowerCase());
                }
            }

        }
    }
}
 

@Test
public void testInsertIntoPartition() throws Exception {
    final boolean isPartitionedTable = true;
    String tableName = createTable(isPartitionedTable);
    String insertTableName = createTable(isPartitionedTable);
    String query =
        "insert into " + insertTableName + " partition(dt = '"+ PART_FILE + "') select id, name from " + tableName
            + " where dt = '"+ PART_FILE + "'";
    runCommand(query);

    final Set<ReadEntity> inputs = getInputs(tableName, Entity.Type.TABLE);
    final Set<WriteEntity> outputs = getOutputs(insertTableName, Entity.Type.TABLE);
    outputs.iterator().next().setWriteType(WriteEntity.WriteType.INSERT);

    final Set<ReadEntity> partitionIps = new LinkedHashSet<ReadEntity>() {
        {
            addAll(inputs);
            add(getPartitionInput());

        }
    };

    final Set<WriteEntity> partitionOps = new LinkedHashSet<WriteEntity>() {
        {
            addAll(outputs);
            add(getPartitionOutput());

        }
    };

    validateProcess(constructEvent(query,  HiveOperation.QUERY, partitionIps, partitionOps), inputs, outputs);

    assertTableIsRegistered(DEFAULT_DB, tableName);
    assertTableIsRegistered(DEFAULT_DB, insertTableName);

    //TODO -Add update test case
}
 

@Override
public Task<? extends Serializable> createShowGrantTask(ASTNode ast, Path resultFile, HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs) throws SemanticException {
  SentryHivePrivilegeObjectDesc privHiveObj = null;

  ASTNode principal = (ASTNode) ast.getChild(0);
  PrincipalType type = PrincipalType.USER;
  switch (principal.getType()) {
  case HiveParser.TOK_USER:
    type = PrincipalType.USER;
    break;
  case HiveParser.TOK_GROUP:
    type = PrincipalType.GROUP;
    break;
  case HiveParser.TOK_ROLE:
    type = PrincipalType.ROLE;
    break;
  }
  if (type != PrincipalType.ROLE) {
    String msg = SentryHiveConstants.GRANT_REVOKE_NOT_SUPPORTED_FOR_PRINCIPAL + type;
    throw new SemanticException(msg);
  }
  String principalName = BaseSemanticAnalyzer.unescapeIdentifier(principal.getChild(0).getText());
  PrincipalDesc principalDesc = new PrincipalDesc(principalName, type);

  // Partition privileges are not supported by Sentry
  if (ast.getChildCount() > 1) {
    ASTNode child = (ASTNode) ast.getChild(1);
    if (child.getToken().getType() == HiveParser.TOK_PRIV_OBJECT_COL) {
      privHiveObj = analyzePrivilegeObject(child);
    } else {
      throw new SemanticException("Unrecognized Token: " + child.getToken().getType());
    }
  }

  ShowGrantDesc showGrant = new ShowGrantDesc(resultFile.toString(),
      principalDesc, privHiveObj);
  return createTask(new DDLWork(inputs, outputs, showGrant));
}
 

private String sortEventsAndGetProcessQualifiedName(final HiveHook.HiveEventContext event) throws HiveException{
    SortedSet<ReadEntity> sortedHiveInputs = event.getInputs() == null ? null : new TreeSet<ReadEntity>(entityComparator);
    SortedSet<WriteEntity> sortedHiveOutputs = event.getOutputs() == null ? null : new TreeSet<WriteEntity>(entityComparator);

    if ( event.getInputs() != null) {
        sortedHiveInputs.addAll(event.getInputs());
    }
    if ( event.getOutputs() != null) {
        sortedHiveOutputs.addAll(event.getOutputs());
    }
    return getProcessQualifiedName(hiveMetaStoreBridge, event, sortedHiveInputs, sortedHiveOutputs, getSortedProcessDataSets(event.getInputs()), getSortedProcessDataSets(event.getOutputs()));
}
 

private HiveHook.HiveEventContext constructEvent(String query, HiveOperation op, Set<ReadEntity> inputs, Set<WriteEntity> outputs) {
    HiveHook.HiveEventContext event = new HiveHook.HiveEventContext();
    event.setQueryStr(query);
    event.setOperation(op);
    event.setInputs(inputs);
    event.setOutputs(outputs);
    return event;
}
 

@Override
public Task<? extends Serializable> createShowRolesTask(ASTNode ast, Path resFile,
    HashSet<ReadEntity> inputs, HashSet<WriteEntity> outputs) throws SemanticException {
  RoleDDLDesc showRolesDesc = new RoleDDLDesc(null, null, RoleDDLDesc.RoleOperation.SHOW_ROLES,
      null);
  showRolesDesc.setResFile(resFile.toString());
  return createTask(new DDLWork(inputs, outputs, showRolesDesc));
}
 

@Override
public Task<? extends Serializable> createCreateRoleTask(ASTNode ast, HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs) throws SemanticException {
  String roleName = BaseSemanticAnalyzer.unescapeIdentifier(ast.getChild(0).getText());
  if (AccessConstants.RESERVED_ROLE_NAMES.contains(roleName.toUpperCase())) {
    String msg = "Roles cannot be one of the reserved roles: " + AccessConstants.RESERVED_ROLE_NAMES;
    throw new SemanticException(msg);
  }
  RoleDDLDesc roleDesc = new RoleDDLDesc(roleName, RoleDDLDesc.RoleOperation.CREATE_ROLE);
  return createTask(new DDLWork(inputs, outputs, roleDesc));
}
 

@Override
public Task<? extends Serializable> createDropRoleTask(ASTNode ast, HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs) throws SemanticException {
  String roleName = BaseSemanticAnalyzer.unescapeIdentifier(ast.getChild(0).getText());
  if (AccessConstants.RESERVED_ROLE_NAMES.contains(roleName.toUpperCase())) {
    String msg = "Roles cannot be one of the reserved roles: " + AccessConstants.RESERVED_ROLE_NAMES;
    throw new SemanticException(msg);
  }
  RoleDDLDesc roleDesc = new RoleDDLDesc(roleName, RoleDDLDesc.RoleOperation.DROP_ROLE);
  return createTask(new DDLWork(inputs, outputs, roleDesc));
}
 

@Override
public Task<? extends Serializable> createGrantTask(ASTNode ast, HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs) throws SemanticException {
  List<PrivilegeDesc> privilegeDesc = analyzePrivilegeListDef(
      (ASTNode) ast.getChild(0));
  List<PrincipalDesc> principalDesc = analyzePrincipalListDef(
      (ASTNode) ast.getChild(1));
  SentryHivePrivilegeObjectDesc privilegeObj = null;
  boolean grantOption = false;
  if (ast.getChildCount() > 2) {
    for (int i = 2; i < ast.getChildCount(); i++) {
      ASTNode astChild = (ASTNode) ast.getChild(i);
      if (astChild.getType() == HiveParser.TOK_GRANT_WITH_OPTION) {
        grantOption = true;
      } else if (astChild.getType() == HiveParser.TOK_PRIV_OBJECT) {
        privilegeObj = analyzePrivilegeObject(astChild);
      }
    }
  }
  String userName = null;
  if (SessionState.get() != null
      && SessionState.get().getAuthenticator() != null) {
    userName = SessionState.get().getAuthenticator().getUserName();
  }
  Preconditions.checkNotNull(privilegeObj, "privilegeObj is null for " + ast.dump());
  if (privilegeObj.getPartSpec() != null) {
    throw new SemanticException(SentryHiveConstants.PARTITION_PRIVS_NOT_SUPPORTED);
  }
  for (PrincipalDesc princ : principalDesc) {
    if (princ.getType() != PrincipalType.ROLE) {
      String msg = SentryHiveConstants.GRANT_REVOKE_NOT_SUPPORTED_FOR_PRINCIPAL + princ.getType();
      throw new SemanticException(msg);
    }
  }
  GrantDesc grantDesc = new GrantDesc(privilegeObj, privilegeDesc,
      principalDesc, userName, PrincipalType.USER, grantOption);
  return createTask(new DDLWork(inputs, outputs, grantDesc));
}
 

@Override
public Task<? extends Serializable> createShowCurrentRoleTask(HashSet<ReadEntity> inputs,
    HashSet<WriteEntity> outputs, Path resultFile) throws SemanticException {
  RoleDDLDesc ddlDesc = new RoleDDLDesc(null, RoleDDLDesc.RoleOperation.SHOW_CURRENT_ROLE);
  ddlDesc.setResFile(resultFile.toString());
  return createTask(new DDLWork(inputs, outputs, ddlDesc));
}
 

@Test
public void testLoadDFSPathPartitioned() throws Exception {
    String tableName = createTable(true, true, false);

    assertTableIsRegistered(DEFAULT_DB, tableName);

    final String loadFile = createTestDFSFile("loadDFSFile");
    String query = "load data inpath '" + loadFile + "' into table " + tableName + " partition(dt = '"+ PART_FILE + "')";
    runCommand(query);

    final Set<WriteEntity> outputs = getOutputs(tableName, Entity.Type.TABLE);
    final Set<ReadEntity> inputs = getInputs(loadFile, Entity.Type.DFS_DIR);

    final Set<WriteEntity> partitionOps = new LinkedHashSet<>(outputs);
    partitionOps.addAll(getOutputs(DEFAULT_DB + "@" + tableName + "@dt=" + PART_FILE, Entity.Type.PARTITION));

    Referenceable processReference = validateProcess(constructEvent(query, HiveOperation.LOAD, inputs, partitionOps), inputs, outputs);
    validateHDFSPaths(processReference, INPUTS, loadFile);
    validateOutputTables(processReference, outputs);

    final String loadFile2 = createTestDFSFile("loadDFSFile1");
    query = "load data inpath '" + loadFile2 + "' into table " + tableName + " partition(dt = '"+ PART_FILE + "')";
    runCommand(query);

    Set<ReadEntity> process2Inputs = getInputs(loadFile2, Entity.Type.DFS_DIR);
    Set<ReadEntity> expectedInputs = new LinkedHashSet<>();
    expectedInputs.addAll(process2Inputs);
    expectedInputs.addAll(inputs);

    validateProcess(constructEvent(query, HiveOperation.LOAD, expectedInputs, partitionOps), expectedInputs, outputs);

}
 

@Test
public void testDropAndRecreateCTASOutput() throws Exception {
    String tableName = createTable();
    String ctasTableName = "table" + random();
    String query = "create table " + ctasTableName + " as select * from " + tableName;
    runCommand(query);

    assertTableIsRegistered(DEFAULT_DB, ctasTableName);

    Set<ReadEntity> inputs = getInputs(tableName, Entity.Type.TABLE);
    Set<WriteEntity> outputs =  getOutputs(ctasTableName, Entity.Type.TABLE);

    final HiveHook.HiveEventContext hiveEventContext = constructEvent(query, HiveOperation.CREATETABLE_AS_SELECT, inputs, outputs);
    String processId = assertProcessIsRegistered(hiveEventContext);

    final String drpquery = String.format("drop table %s ", ctasTableName);
    runCommandWithDelay(drpquery, 100);
    assertTableIsNotRegistered(DEFAULT_DB, ctasTableName);

    runCommand(query);
    assertTableIsRegistered(DEFAULT_DB, ctasTableName);
    outputs =  getOutputs(ctasTableName, Entity.Type.TABLE);
    String process2Id = assertProcessIsRegistered(hiveEventContext, inputs, outputs);

    assertNotEquals(process2Id, processId);

    Referenceable processRef = atlasClient.getEntity(processId);
    validateOutputTables(processRef, outputs);
}
 

private boolean skipProcess() {
    Set<ReadEntity>  inputs  = getInputs();
    Set<WriteEntity> outputs = getOutputs();

    boolean ret = CollectionUtils.isEmpty(inputs) && CollectionUtils.isEmpty(outputs);

    if (!ret) {
        if (getContext().getHiveOperation() == HiveOperation.QUERY) {
            // Select query has only one output
            if (outputs.size() == 1) {
                WriteEntity output = outputs.iterator().next();

                if (output.getType() == Entity.Type.DFS_DIR || output.getType() == Entity.Type.LOCAL_DIR) {
                    if (output.getWriteType() == WriteEntity.WriteType.PATH_WRITE && output.isTempURI()) {
                        ret = true;
                    }
                }
                // DELETE and UPDATE initially have one input and one output.
                // Since they do not support sub-query, they won't create a lineage that have one input and one output. (One input only)
                // It's safe to filter them out here.
                if (output.getWriteType() == WriteEntity.WriteType.DELETE || output.getWriteType() == WriteEntity.WriteType.UPDATE) {
                    ret = true;
                }
            }
        }
    }

    return ret;
}
 

@Test
public void testCTAS() throws Exception {
    String tableName = createTable();
    String ctasTableName = "table" + random();
    String query = "create table " + ctasTableName + " as select * from " + tableName;
    runCommand(query);

    final Set<ReadEntity> readEntities = getInputs(tableName, Entity.Type.TABLE);
    final Set<WriteEntity> writeEntities = getOutputs(ctasTableName, Entity.Type.TABLE);

    assertProcessIsRegistered(constructEvent(query, HiveOperation.CREATETABLE_AS_SELECT, readEntities, writeEntities));
    assertTableIsRegistered(DEFAULT_DB, ctasTableName);

}
 

@Test
public void testIgnoreSelect() throws Exception {
    String tableName = createTable();
    String query = "select * from " + tableName;
    runCommand(query);
    Set<ReadEntity> inputs = getInputs(tableName, Entity.Type.TABLE);
    HiveHook.HiveEventContext hiveEventContext = constructEvent(query, HiveOperation.QUERY, inputs, null);
    assertProcessIsNotRegistered(hiveEventContext);

    //check with uppercase table name
    query = "SELECT * from " + tableName.toUpperCase();
    runCommand(query);
    assertProcessIsNotRegistered(hiveEventContext);
}