org.apache.hadoop.hive.ql.metadata.Table Java Examples

@Test
public void testImportContinuesWhenTableRegistrationFails() throws Exception {
    setupDB(hiveClient, TEST_DB_NAME);
    final String table2Name = TEST_TABLE_NAME + "_1";
    List<Table> hiveTables = setupTables(hiveClient, TEST_DB_NAME, TEST_TABLE_NAME, table2Name);

    returnExistingDatabase(TEST_DB_NAME, atlasClient, CLUSTER_NAME);
    when(hiveClient.getTable(TEST_DB_NAME, TEST_TABLE_NAME)).thenThrow(new RuntimeException("Timeout while reading data from hive metastore"));

    when(atlasClient.getEntity(HiveDataTypes.HIVE_TABLE.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, TEST_DB_NAME,
        table2Name))).thenReturn(
        getEntityReference(HiveDataTypes.HIVE_TABLE.getName(), "82e06b34-9151-4023-aa9d-b82103a50e77"));
    when(atlasClient.getEntity("82e06b34-9151-4023-aa9d-b82103a50e77")).thenReturn(createTableReference());
    String processQualifiedName = HiveMetaStoreBridge.getTableProcessQualifiedName(CLUSTER_NAME, hiveTables.get(1));
    when(atlasClient.getEntity(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
        processQualifiedName)).thenReturn(getEntityReference(HiveDataTypes.HIVE_PROCESS.getName(), "82e06b34-9151-4023-aa9d-b82103a50e77"));

    HiveMetaStoreBridge bridge = new HiveMetaStoreBridge(CLUSTER_NAME, hiveClient, atlasClient);
    try {
        bridge.importHiveMetadata(false);
    } catch (Exception e) {
        Assert.fail("Table registration failed with exception", e);
    }
}
 

private String getCreateTableString(Table table, String location){
    String            colString = "";
    List<FieldSchema> colList   = table.getAllCols();

    if (colList != null) {
        for (FieldSchema col : colList) {
            colString += col.getName() + " " + col.getType() + ",";
        }

        if (colList.size() > 0) {
            colString = colString.substring(0, colString.length() - 1);
            colString = "(" + colString + ")";
        }
    }

    String query = "create external table " + table.getTableName() +  colString + " location '" + location + "'";

    return query;
}
 

@Override
public void onAlterTable(AlterTableEvent tableEvent) {
    HiveOperationContext context = new HiveOperationContext(tableEvent);
    Table                oldTable = toTable(tableEvent.getOldTable());
    Table                newTable = toTable(tableEvent.getNewTable());

    if (isTableRename(oldTable, newTable)) {
        context.setOperation(ALTERTABLE_RENAME);
    } else if (isColumnRename(oldTable, newTable, context)) {
        context.setOperation(ALTERTABLE_RENAMECOL);
    } else if(isAlterTableProperty(tableEvent, "last_modified_time") ||
            isAlterTableProperty(tableEvent, "transient_lastDdlTime")) {
        context.setOperation(ALTERTABLE_PROPERTIES); // map other alter table operations to ALTERTABLE_PROPERTIES
    }

    hook.handleEvent(context);
}
 

@Test
public void testGetSchemaFromUrlUsingHiveSchema() throws IOException, HiveException {
  FileSystem fs = FileSystem.getLocal(new Configuration());

  String jobId = "123";
  State state = new State();
  state.setProp(ConfigurationKeys.JOB_ID_KEY, jobId);

  AvroSchemaManager asm = new AvroSchemaManager(fs, state);
  Partition partition = getTestPartition(new Table("testDb", "testTable"));
  Path schemaPath = asm.getSchemaUrl(partition);

  Schema actualSchema = AvroUtils.parseSchemaFromFile(schemaPath, fs);
  String expectedSchema = new String(Files.readAllBytes(
      Paths.get(getClass().getClassLoader().getResource("avroSchemaManagerTest/expectedSchema.avsc").getFile())));
  Assert.assertEquals(actualSchema.toString(), expectedSchema);
}
 

@Test
public void testStateStoreReadWrite() throws Exception {

  String dbName = "testStateStoreReadWrite";
  LocalHiveMetastoreTestUtils.getInstance().dropDatabaseIfExists(dbName);

  PartitionLevelWatermarker watermarker0 = new PartitionLevelWatermarker(new SourceState());
  Table mockTable = localTestTable(dbName, "table1", true);

  watermarker0.onTableProcessBegin(mockTable, 0l);
  long now = new DateTime().getMillis();
  watermarker0.onPartitionProcessBegin(localTestPartition(mockTable, ImmutableList.of("2016")), 0, now);
  List<WorkUnit> workunits = Lists.newArrayList();
  watermarker0.onGetWorkunitsEnd(workunits);

  @SuppressWarnings("deprecation")
  WorkUnitState previousWus = new WorkUnitState(workunits.get(0));
  watermarker0.setActualHighWatermark(previousWus);

  SourceState state = new SourceState(new State(), Lists.newArrayList(previousWus));
  PartitionLevelWatermarker watermarker = new PartitionLevelWatermarker(state);

  Assert.assertEquals(watermarker.getPreviousWatermarks().size(), 1);
  Assert.assertEquals(watermarker.getPreviousWatermarks().get(dbName + "@table1"), ImmutableMap.of("2016", now));

}
 

private AtlasEntityWithExtInfo registerTable(AtlasEntity dbEntity, Table table) throws AtlasHookException {
    try {
        AtlasEntityWithExtInfo ret;
        AtlasEntityWithExtInfo tableEntity = findTableEntity(table);

        if (tableEntity == null) {
            tableEntity = toTableEntity(dbEntity, table);

            ret = registerInstance(tableEntity);
        } else {
            LOG.info("Table {}.{} is already registered with id {}. Updating entity.", table.getDbName(), table.getTableName(), tableEntity.getEntity().getGuid());

            ret = toTableEntity(dbEntity, table, tableEntity);

            updateInstance(ret);
        }

        return ret;
    } catch (Exception e) {
        throw new AtlasHookException("HiveMetaStoreBridge.registerTable() failed.", e);
    }
}
 

@Before
public void setup() throws Exception {
  conf = new HiveConf();
  baseDir = Files.createTempDir();
  baseDir.setWritable(true, false);
  conf.setVar(HiveConf.ConfVars.SCRATCHDIR, baseDir.getAbsolutePath());
  SessionState.start(conf);
  conf.setVar(ConfVars.HIVE_AUTHORIZATION_TASK_FACTORY,
      SentryHiveAuthorizationTaskFactoryImpl.class.getName());

  db = Mockito.mock(Hive.class);
  table = new Table(DB, TABLE);
  partition = new Partition(table);
  context = new Context(conf);
  parseDriver = new ParseDriver();
  analyzer = new DDLSemanticAnalyzer(conf, db);
  SessionState.start(conf);
  Mockito.when(db.getTable(TABLE, false)).thenReturn(table);
  Mockito.when(db.getPartition(table, new HashMap<String, String>(), false))
  .thenReturn(partition);

  HadoopDefaultAuthenticator auth = new HadoopDefaultAuthenticator();
  auth.setConf(conf);
  currentUser = auth.getUserName();

}
 

HBaseTableInfo(Table table) {
    Map<String, String> parameters = table.getParameters();

    if (MapUtils.isNotEmpty(parameters)) {
        hbaseNameSpace = HBASE_DEFAULT_NAMESPACE;
        hbaseTableName = parameters.get(HBASE_PARAM_TABLE_NAME);

        if (hbaseTableName != null) {
            if (hbaseTableName.contains(HBASE_NAMESPACE_TABLE_DELIMITER)) {
                String[] hbaseTableInfo = hbaseTableName.split(HBASE_NAMESPACE_TABLE_DELIMITER);

                if (hbaseTableInfo.length > 1) {
                    hbaseNameSpace = hbaseTableInfo[0];
                    hbaseTableName = hbaseTableInfo[1];
                }
            }
        }
    }
}
 

@VisibleForTesting
static void addMetadataToTargetTable(Table targetTable, Path targetLocation, String targetDatabase, long startTime)
    throws IOException {
  targetTable.setDbName(targetDatabase);
  targetTable.setDataLocation(targetLocation);
  /*
   * Need to set the table owner as the flow executor
   */
  targetTable.setOwner(UserGroupInformation.getCurrentUser().getShortUserName());
  targetTable.getTTable().putToParameters(HiveDataset.REGISTERER, GOBBLIN_DISTCP);
  targetTable.getTTable().putToParameters(HiveDataset.REGISTRATION_GENERATION_TIME_MILLIS,
      Long.toString(startTime));

  /**
   * Only set the this constants when source table has it.
   */
  targetTable.getTTable().getSd().getSerdeInfo().getParameters()
      .computeIfPresent(HiveConstants.PATH, (k,v) -> targetLocation.toString());
  targetTable.getTTable().unsetCreateTime();
}
 

private HiveDataset getHiveDataset(String tableString, FileSystem fs, State state) throws IOException {
  try {
    HiveMetastoreClientPool pool = HiveMetastoreClientPool.get(state.getProperties(),
        Optional.fromNullable(state.getProp(HIVE_METASTORE_URI_KEY)));

    List<String> tokens = Splitter.on(".").splitToList(tableString);
    DbAndTable sourceDbAndTable = new DbAndTable(tokens.get(0), tokens.get(1));

    try (AutoReturnableObject<IMetaStoreClient> client = pool.getClient()) {
      Table sourceTable = new Table(client.get().getTable(sourceDbAndTable.getDb(), sourceDbAndTable.getTable()));
      return new HiveDataset(fs, pool, sourceTable, ConfigUtils.propertiesToConfig(state.getProperties()));
    }
  } catch (TException exc) {
    throw new RuntimeException(exc);
  }
}
 

private Table constructAvroTable(String database, String tableName, Schema schema, Partitioner partitioner)
    throws HiveMetaStoreException {
  Table table = new Table(database, tableName);
  table.setTableType(TableType.EXTERNAL_TABLE);
  table.getParameters().put("EXTERNAL", "TRUE");
  String tablePath = FileUtils.hiveDirectoryName(url, topicsDir, tableName);
  table.setDataLocation(new Path(tablePath));
  table.setSerializationLib(avroSerde);
  try {
    table.setInputFormatClass(avroInputFormat);
    table.setOutputFormatClass(avroOutputFormat);
  } catch (HiveException e) {
    throw new HiveMetaStoreException("Cannot find input/output format:", e);
  }
  List<FieldSchema> columns = HiveSchemaConverter.convertSchema(schema);
  table.setFields(columns);
  table.setPartCols(partitioner.partitionFields());
  table.getParameters().put(AVRO_SCHEMA_LITERAL, avroData.fromConnectSchema(schema).toString());
  return table;
}
 

private Table constructParquetTable(String database, String tableName, Schema schema, Partitioner partitioner) throws HiveMetaStoreException {
  Table table = new Table(database, tableName);
  table.setTableType(TableType.EXTERNAL_TABLE);
  table.getParameters().put("EXTERNAL", "TRUE");
  String tablePath = FileUtils.hiveDirectoryName(url, topicsDir, tableName);
  table.setDataLocation(new Path(tablePath));
  table.setSerializationLib(getHiveParquetSerde());
  try {
    table.setInputFormatClass(getHiveParquetInputFormat());
    table.setOutputFormatClass(getHiveParquetOutputFormat());
  } catch (HiveException e) {
    throw new HiveMetaStoreException("Cannot find input/output format:", e);
  }
  // convert copycat schema schema to Hive columns
  List<FieldSchema> columns = HiveSchemaConverter.convertSchema(schema);
  table.setFields(columns);
  table.setPartCols(partitioner.partitionFields());
  return table;
}
 

@Test
public void testDroppedPartitions() throws Exception {
  WorkUnitState previousWus = new WorkUnitState();
  previousWus.setProp(ConfigurationKeys.DATASET_URN_KEY, "db@test_dataset_urn");
  previousWus.setProp(PartitionLevelWatermarker.IS_WATERMARK_WORKUNIT_KEY, true);
  previousWus
      .setActualHighWatermark(new MultiKeyValueLongWatermark(ImmutableMap.of("2015-01", 100l, "2015-02", 101l)));

  SourceState state = new SourceState(new State(), Lists.newArrayList(previousWus));
  PartitionLevelWatermarker watermarker = new PartitionLevelWatermarker(state);

  Table table = mockTable("test_dataset_urn");
  Mockito.when(table.getPartitionKeys()).thenReturn(ImmutableList.of(new FieldSchema("year", "string", "")));

  Partition partition2015 = mockPartition(table, ImmutableList.of("2015"));

  // partition 2015 replaces 2015-01 and 2015-02
  Mockito.when(partition2015.getParameters()).thenReturn(
      ImmutableMap.of(AbstractAvroToOrcConverter.REPLACED_PARTITIONS_HIVE_METASTORE_KEY, "2015-01|2015-02"));
  watermarker.onPartitionProcessBegin(partition2015, 0l, 0l);

  Assert.assertEquals(watermarker.getExpectedHighWatermarks().get("db@test_dataset_urn"), ImmutableMap.of("2015", 0l));
}
 

private static long getFileSizeForPartition(Read readRequest, Partition partitionToRead)
    throws Exception {
  IMetaStoreClient client = null;
  try {
    HiveConf hiveConf = HCatalogUtils.createHiveConf(readRequest);
    client = HCatalogUtils.createMetaStoreClient(hiveConf);
    List<org.apache.hadoop.hive.ql.metadata.Partition> p = new ArrayList<>();
    Table table = HCatUtil.getTable(client, readRequest.getDatabase(), readRequest.getTable());
    final org.apache.hadoop.hive.ql.metadata.Partition partition =
        new org.apache.hadoop.hive.ql.metadata.Partition(table, partitionToRead);
    p.add(partition);
    final List<Long> fileSizeForPartitions = StatsUtils.getFileSizeForPartitions(hiveConf, p);
    return fileSizeForPartitions.get(0);
  } finally {
    // IMetaStoreClient is not AutoCloseable, closing it manually
    if (client != null) {
      client.close();
    }
  }
}
 

/**
 * Returns the size of the table in bytes, does not take into consideration filter/partition
 * details passed, if any.
 */
@Override
public long getEstimatedSizeBytes(PipelineOptions pipelineOptions) throws Exception {
  IMetaStoreClient client = null;
  try {
    HiveConf hiveConf = HCatalogUtils.createHiveConf(spec);
    client = HCatalogUtils.createMetaStoreClient(hiveConf);
    Table table = HCatUtil.getTable(client, spec.getDatabase(), spec.getTable());
    return StatsUtils.getFileSizeForTable(hiveConf, table);
  } finally {
    // IMetaStoreClient is not AutoCloseable, closing it manually
    if (client != null) {
      client.close();
    }
  }
}
 

public HiveDataset(FileSystem fs, HiveMetastoreClientPool clientPool, Table table, Properties properties, Config datasetConfig) {
  this.fs = fs;
  this.clientPool = clientPool;
  this.table = table;
  this.properties = properties;

  this.tableRootPath = PathUtils.isGlob(this.table.getDataLocation()) ? Optional.<Path> absent() :
      Optional.fromNullable(this.table.getDataLocation());

  this.tableIdentifier = this.table.getDbName() + "." + this.table.getTableName();

  this.datasetNamePattern = Optional.fromNullable(ConfigUtils.getString(datasetConfig, DATASET_NAME_PATTERN_KEY, null));
  this.dbAndTable = new DbAndTable(table.getDbName(), table.getTableName());
  if (this.datasetNamePattern.isPresent()) {
    this.logicalDbAndTable = parseLogicalDbAndTable(this.datasetNamePattern.get(), this.dbAndTable, LOGICAL_DB_TOKEN, LOGICAL_TABLE_TOKEN);
  } else {
    this.logicalDbAndTable = this.dbAndTable;
  }
  this.datasetConfig = resolveConfig(datasetConfig, dbAndTable, logicalDbAndTable);
  this.metricContext = Instrumented.getMetricContext(new State(properties), HiveDataset.class,
      Lists.<Tag<?>> newArrayList(new Tag<>(DATABASE, table.getDbName()), new Tag<>(TABLE, table.getTableName())));
}
 

private Referenceable registerTable(Referenceable dbReference, Table table) throws AtlasHookException {
    try {
        String dbName = table.getDbName();
        String tableName = table.getTableName();
        LOG.info("Attempting to register table [{}]", tableName);
        Referenceable tableReference = getTableReference(table);
        LOG.info("Found result {}", tableReference);
        if (tableReference == null) {
            tableReference = createTableInstance(dbReference, table);
            tableReference = registerInstance(tableReference);
        } else {
            LOG.info("Table {}.{} is already registered with id {}. Updating entity.", dbName, tableName,
                    tableReference.getId().id);
            tableReference = createOrUpdateTableInstance(dbReference, tableReference, table);
            updateInstance(tableReference);
        }
        return tableReference;
    } catch (Exception e) {
        throw new AtlasHookException("HiveMetaStoreBridge.getStorageDescQFName() failed.", e);
    }
}
 

private Referenceable replaceTableQFName(HiveEventContext event, Table oldTable, Table newTable, final Referenceable tableEntity, final String oldTableQFName, final String newTableQFName) throws HiveException {
    tableEntity.set(AtlasClient.NAME,  oldTable.getTableName().toLowerCase());
    tableEntity.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, oldTableQFName);

    //Replace table entity with new name
    final Referenceable newEntity = new Referenceable(HiveDataTypes.HIVE_TABLE.getName());
    newEntity.set(AtlasClient.NAME, newTable.getTableName().toLowerCase());
    newEntity.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, newTableQFName);

    ArrayList<String> alias_list = new ArrayList<>();
    alias_list.add(oldTable.getTableName().toLowerCase());
    newEntity.set(HiveMetaStoreBridge.TABLE_ALIAS_LIST, alias_list);
    event.addMessage(new HookNotification.EntityPartialUpdateRequest(event.getUser(),
        HiveDataTypes.HIVE_TABLE.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
        oldTableQFName, newEntity));

    return newEntity;
}
 

@Test
public void testImportThatUpdatesRegisteredTable() throws Exception {
    setupDB(hiveClient, TEST_DB_NAME);

    List<Table> hiveTables = setupTables(hiveClient, TEST_DB_NAME, TEST_TABLE_NAME);

    returnExistingDatabase(TEST_DB_NAME, atlasClient, CLUSTER_NAME);

    // return existing table
    when(atlasClient.getEntity(HiveDataTypes.HIVE_TABLE.getName(),
        AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, TEST_DB_NAME, TEST_TABLE_NAME)))
        .thenReturn(getEntityReference(HiveDataTypes.HIVE_TABLE.getName(), "82e06b34-9151-4023-aa9d-b82103a50e77"));
    when(atlasClient.getEntity("82e06b34-9151-4023-aa9d-b82103a50e77")).thenReturn(createTableReference());
    String processQualifiedName = HiveMetaStoreBridge.getTableProcessQualifiedName(CLUSTER_NAME, hiveTables.get(0));
    when(atlasClient.getEntity(HiveDataTypes.HIVE_PROCESS.getName(),
        AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, processQualifiedName)).thenReturn(getEntityReference(HiveDataTypes.HIVE_PROCESS.getName(), "82e06b34-9151-4023-aa9d-b82103a50e77"));

    HiveMetaStoreBridge bridge = new HiveMetaStoreBridge(CLUSTER_NAME, hiveClient, atlasClient);
    bridge.importHiveMetadata(true);

    // verify update is called on table
    verify(atlasClient).updateEntity(eq("82e06b34-9151-4023-aa9d-b82103a50e77"),
            (Referenceable) argThat(new MatchesReferenceableProperty(HiveMetaStoreBridge.TABLE_TYPE_ATTR,
                    TableType.EXTERNAL_TABLE.name())));
}
 

private List<Table> setupTables(Hive hiveClient, String databaseName, String... tableNames) throws HiveException {
    List<Table> tables = new ArrayList<>();
    when(hiveClient.getAllTables(databaseName)).thenReturn(Arrays.asList(tableNames));
    for(String tableName : tableNames) {
        Table testTable = createTestTable(databaseName, tableName);
        when(hiveClient.getTable(databaseName, tableName)).thenReturn(testTable);
        tables.add(testTable);
    }
    return tables;
}
 

@Test(expectedExceptions = IllegalStateException.class)
public void testPartitionBeginBegoreTableBegin() throws Exception {
  SourceState state = new SourceState();
  PartitionLevelWatermarker watermarker = new PartitionLevelWatermarker(state);

  Table table = mockTable("test_dataset_urn");
  Partition partition = mockPartition(table, ImmutableList.of(""));

  watermarker.onPartitionProcessBegin(partition, 0l, 0l);
}
 

@Override
public void alterSchema(String database, String tableName, Schema schema) {
  Table table = hiveMetaStore.getTable(database, tableName);
  List<FieldSchema> columns = HiveSchemaConverter.convertSchema(schema);
  table.setFields(columns);
  hiveMetaStore.alterTable(table);
}
 

public String getQualifiedName(Table table) {
    String tableName = table.getTableName();

    if (table.isTemporary()) {
        if (SessionState.get() != null && SessionState.get().getSessionId() != null) {
            tableName = tableName + TEMP_TABLE_PREFIX + SessionState.get().getSessionId();
        } else {
            tableName = tableName + TEMP_TABLE_PREFIX + RandomStringUtils.random(10);
        }
    }

    return (table.getDbName() + QNAME_SEP_ENTITY_NAME + tableName + QNAME_SEP_METADATA_NAMESPACE).toLowerCase() + getMetadataNamespace();
}
 

private Partition getTestPartition(Table table) throws HiveException {
  Partition partition = new Partition(table, ImmutableMap.of("partition_key", "1"), null);
  StorageDescriptor sd = new StorageDescriptor();
  sd.setSerdeInfo(new SerDeInfo("avro", AvroSerDe.class.getName(), null));
  sd.setCols(Lists.newArrayList(new FieldSchema("foo", "int", null)));
  partition.getTPartition().setSd(sd);
  return partition;
}
 

protected static Table refreshTable(HiveMetaStoreBridge dgiBridge, String dbName, String tableName) {
    try {
        return dgiBridge.getHiveClient().getTable(dbName, tableName);
    } catch (HiveException excp) { // this might be the case for temp tables
        LOG.warn("failed to get details for table {}.{}. Ignoring. {}: {}", dbName, tableName, excp.getClass().getCanonicalName(), excp.getMessage());
    }

    return null;
}
 

protected static void addInputs(HiveMetaStoreBridge hiveBridge, HiveOperation op, SortedSet<ReadEntity> sortedInputs, StringBuilder buffer, final Map<ReadEntity, AtlasEntity> 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 &&
                            (Entity.Type.DFS_DIR.equals(input.getType()) || Entity.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() == Entity.Type.PARTITION || input.getType() == Entity.Type.TABLE) {
                            Table inputTable = refreshTable(hiveBridge, input.getTable().getDbName(), input.getTable().getTableName());

                            if (inputTable != null) {
                                addDataset(buffer, refs.get(input), HiveMetaStoreBridge.getTableCreatedTime(inputTable));
                            }
                        } else {
                            addDataset(buffer, refs.get(input));
                        }
                    }

                    dataSetsProcessed.add(input.getName().toLowerCase());
                }
            }

        }
    }
}
 

private Table createTestTable(String databaseName, String tableName) throws HiveException {
    Table table = new Table(databaseName, tableName);
    table.setInputFormatClass(TextInputFormat.class);
    table.setFields(new ArrayList<FieldSchema>() {{
        add(new FieldSchema("col1", "string", "comment1"));
    }
    });
    table.setTableType(TableType.EXTERNAL_TABLE);
    table.setDataLocation(new Path("somehdfspath"));
    return table;
}
 

@Test
public void testImportFailsWhenTableRegistrationFails() throws Exception {
    setupDB(hiveClient, TEST_DB_NAME);
    final String table2Name = TEST_TABLE_NAME + "_1";
    List<Table> hiveTables = setupTables(hiveClient, TEST_DB_NAME, TEST_TABLE_NAME, table2Name);

    returnExistingDatabase(TEST_DB_NAME, atlasClientV2, METADATA_NAMESPACE);
    when(hiveClient.getTable(TEST_DB_NAME, TEST_TABLE_NAME)).thenThrow(new RuntimeException("Timeout while reading data from hive metastore"));


    when(atlasClientV2.getEntityByAttribute(HiveDataTypes.HIVE_TABLE.getName(),
            Collections.singletonMap(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
            HiveMetaStoreBridge.getTableQualifiedName(METADATA_NAMESPACE, TEST_DB_NAME, TEST_TABLE_NAME))))
    .thenReturn(new AtlasEntity.AtlasEntityWithExtInfo(
            getEntity(HiveDataTypes.HIVE_TABLE.getName(), AtlasClient.GUID, "82e06b34-9151-4023-aa9d-b82103a50e77")));


    when(atlasEntityWithExtInfo.getEntity("82e06b34-9151-4023-aa9d-b82103a50e77"))
            .thenReturn(createTableReference());

    Table testTable  = hiveTables.get(1);
    String processQualifiedName = HiveMetaStoreBridge.getTableProcessQualifiedName(METADATA_NAMESPACE, testTable);

    when(atlasClientV2.getEntityByAttribute(HiveDataTypes.HIVE_PROCESS.getName(),
            Collections.singletonMap(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
                    processQualifiedName)))
            .thenReturn(new AtlasEntity.AtlasEntityWithExtInfo(
                    getEntity(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.GUID, "82e06b34-9151-4023-aa9d-b82103a50e77")));

    HiveMetaStoreBridge bridge = new HiveMetaStoreBridge(METADATA_NAMESPACE, hiveClient, atlasClientV2);
    try {
        bridge.importHiveMetadata(null, null, true);
        Assert.fail("Table registration is supposed to fail");
    } catch (Exception e) {
        //Expected
    }
}
 

protected static void addOutputs(HiveMetaStoreBridge hiveBridge, HiveOperation op, SortedSet<WriteEntity> sortedOutputs, StringBuilder buffer, final Map<WriteEntity, AtlasEntity> refs, final boolean ignoreHDFSPathsInQFName) throws HiveException {
    if (refs != null) {
        Set<String> dataSetsProcessed = new LinkedHashSet<>();
        if (sortedOutputs != null) {
            for (WriteEntity output : sortedOutputs) {
                final Entity entity = output;
                if (!dataSetsProcessed.contains(output.getName().toLowerCase())) {
                    if (ignoreHDFSPathsInQFName &&
                            (Entity.Type.DFS_DIR.equals(output.getType()) || Entity.Type.LOCAL_DIR.equals(output.getType()))) {
                        LOG.debug("Skipping dfs dir output addition to process qualified name {} ", output.getName());
                    } else if (refs.containsKey(output)) {
                        //HiveOperation.QUERY type encompasses INSERT, INSERT_OVERWRITE, UPDATE, DELETE, PATH_WRITE operations
                        if (addQueryType(op, (WriteEntity) entity)) {
                            buffer.append(SEP);
                            buffer.append(((WriteEntity) entity).getWriteType().name());
                        }

                        if ( output.getType() == Entity.Type.PARTITION || output.getType() == Entity.Type.TABLE) {
                            Table outputTable = refreshTable(hiveBridge, output.getTable().getDbName(), output.getTable().getTableName());

                            if (outputTable != null) {
                                addDataset(buffer, refs.get(output), HiveMetaStoreBridge.getTableCreatedTime(outputTable));
                            }
                        } else {
                            addDataset(buffer, refs.get(output));
                        }
                    }

                    dataSetsProcessed.add(output.getName().toLowerCase());
                }
            }
        }
    }
}
 

@Test
public void testImportWhenPartitionKeysAreNull() throws Exception {
    setupDB(hiveClient, TEST_DB_NAME);
    List<Table> hiveTables = setupTables(hiveClient, TEST_DB_NAME, TEST_TABLE_NAME);
    Table hiveTable = hiveTables.get(0);

    returnExistingDatabase(TEST_DB_NAME, atlasClientV2, METADATA_NAMESPACE);


    when(atlasClientV2.getEntityByAttribute(HiveDataTypes.HIVE_TABLE.getName(),
            Collections.singletonMap(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
                    HiveMetaStoreBridge.getTableQualifiedName(METADATA_NAMESPACE, TEST_DB_NAME, TEST_TABLE_NAME))))
    .thenReturn(new AtlasEntity.AtlasEntityWithExtInfo(
                    getEntity(HiveDataTypes.HIVE_TABLE.getName(), AtlasClient.GUID, "82e06b34-9151-4023-aa9d-b82103a50e77")));

    String processQualifiedName = HiveMetaStoreBridge.getTableProcessQualifiedName(METADATA_NAMESPACE, hiveTable);

    when(atlasClientV2.getEntityByAttribute(HiveDataTypes.HIVE_PROCESS.getName(),
            Collections.singletonMap(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
                    processQualifiedName)))
    .thenReturn(new AtlasEntity.AtlasEntityWithExtInfo(
                    getEntity(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.GUID, "82e06b34-9151-4023-aa9d-b82103a50e77")));

    when(atlasEntityWithExtInfo.getEntity("82e06b34-9151-4023-aa9d-b82103a50e77"))
    .thenReturn(createTableReference());

    Partition partition = mock(Partition.class);
    when(partition.getTable()).thenReturn(hiveTable);
    List partitionValues = Arrays.asList(new String[]{});
    when(partition.getValues()).thenReturn(partitionValues);

    when(hiveClient.getPartitions(hiveTable)).thenReturn(Arrays.asList(new Partition[]{partition}));

    HiveMetaStoreBridge bridge = new HiveMetaStoreBridge(METADATA_NAMESPACE, hiveClient, atlasClientV2);
    try {
        bridge.importHiveMetadata(null, null, true);
    } catch (Exception e) {
        Assert.fail("Partition with null key caused import to fail with exception ", e);
    }
}