content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   * <p>
   * http://www.apache.org/licenses/LICENSE-2.0
   * <p>
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.atlas.examples.sampleapp;
  
  import com.sun.jersey.core.util.MultivaluedMapImpl;
  import org.apache.atlas.AtlasClientV2;
  import org.apache.atlas.AtlasServiceException;
  import org.apache.atlas.model.SearchFilter;
  import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
  import org.apache.atlas.model.typedef.AtlasBusinessMetadataDef;
  import org.apache.atlas.model.typedef.AtlasClassificationDef;
  import org.apache.atlas.model.typedef.AtlasEntityDef;
  import org.apache.atlas.model.typedef.AtlasEnumDef;
  import org.apache.atlas.model.typedef.AtlasRelationshipDef;
  import org.apache.atlas.model.typedef.AtlasStructDef;
  import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
  import org.apache.atlas.model.typedef.AtlasTypesDef;
  import org.apache.atlas.type.AtlasTypeUtil;
  
  import javax.ws.rs.core.MultivaluedMap;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_ATTR1;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_ATTR11;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_ATTR18;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_ATTR2;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_ATTR8;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_COMMENT;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_CREATE_TIME;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_DATA_TYPE;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_END_TIME;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_LAST_ACCESS_TIME;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_GRAPH;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_ID;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_PLAN;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_QUERY_TEXT;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_RANDOM_TABLE;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE1;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_SERDE2;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_START_TIME;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_TABLE_TYPE;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_TEMPORARY;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ATTR_USERNAME;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ENTITY_TYPE_DATASET;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ENTITY_TYPE_PROCESS;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.ENUM_TABLE_TYPE;
  import static org.apache.atlas.examples.sampleapp.SampleAppConstants.STRUCT_TYPE_SERDE;
  import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BOOLEAN;
  import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BYTE;
  import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_STRING;
  import static org.apache.atlas.model.typedef.AtlasRelationshipDef.RelationshipCategory.AGGREGATION;
  import static org.apache.atlas.model.typedef.AtlasRelationshipDef.RelationshipCategory.COMPOSITION;
  import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality.SET;
  import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality.SINGLE;
  import static org.apache.atlas.type.AtlasTypeUtil.createBusinessMetadataDef;
  import static org.apache.atlas.type.AtlasTypeUtil.createOptionalAttrDef;
  import static org.apache.atlas.type.AtlasTypeUtil.createRelationshipEndDef;
  import static org.apache.atlas.type.AtlasTypeUtil.createRelationshipTypeDef;
  import static org.apache.atlas.type.AtlasTypeUtil.createTraitTypeDef;
  
  public class TypeDefExample {
      private static final String[] SAMPLE_APP_TYPES = {
              SampleAppConstants.DATABASE_TYPE,
              SampleAppConstants.TABLE_TYPE,
              SampleAppConstants.COLUMN_TYPE,
              SampleAppConstants.PROCESS_TYPE,
              SampleAppConstants.PII_TAG,
              SampleAppConstants.CLASSIFIED_TAG,
              SampleAppConstants.FINANCE_TAG,
              SampleAppConstants.METRIC_TAG
      };
  
      private final AtlasClientV2 client;
      private       AtlasTypesDef typesDef;
  
      TypeDefExample(AtlasClientV2 client) {
          this.client = client;
      }
  
      public void createTypeDefinitions() throws Exception {
          AtlasEntityDef databaseDef  = createDatabaseDef();
          AtlasEntityDef tableDef     = createTableDef();
          AtlasEntityDef columnDef    = createColumnDef();
          AtlasEntityDef processDef   = createProcessDef();
          AtlasStructDef serDeDef     = createSerDeDef();
          AtlasEnumDef   tableTypeDef = createTableTypeEnumDef();
  
          List<AtlasClassificationDef>   classificationDefs  = createClassificationDefs();
          List<AtlasBusinessMetadataDef> businessMetadataDef = createBusinessMetadataDefs();
          List<AtlasRelationshipDef>     relationshipDefs    = createAtlasRelationshipDef();
  
          AtlasTypesDef typesDef = new AtlasTypesDef(Collections.singletonList(tableTypeDef),
                  Collections.singletonList(serDeDef),
                  classificationDefs,
                  Arrays.asList(databaseDef, tableDef, columnDef, processDef),
                  relationshipDefs,
                  businessMetadataDef);
  
          this.typesDef = batchCreateTypes(typesDef);
      }
  
      public void printTypeDefinitions() throws AtlasServiceException {
          for (String typeName : SAMPLE_APP_TYPES) {
              MultivaluedMap<String, String> searchParams = new MultivaluedMapImpl();
  
              searchParams.add(SearchFilter.PARAM_NAME, typeName);
  
              SearchFilter searchFilter = new SearchFilter(searchParams);
  
              AtlasTypesDef typesDef = client.getAllTypeDefs(searchFilter);
  
              if (typesDef.isEmpty()) {
                  throw new IllegalStateException("Type definition is empty for type: " + typeName);
              }
  
              SampleApp.log("Created type: " + typeName);
          }
      }
  
      public void removeTypeDefinitions() throws AtlasServiceException {
          if (typesDef != null) {
              client.deleteAtlasTypeDefs(typesDef);
  
              typesDef = null;
  
              SampleApp.log("Deleted TypesDef successfully!");
          }
      }
  
      private AtlasEntityDef createDatabaseDef() {
          return AtlasTypeUtil.createClassTypeDef(SampleAppConstants.DATABASE_TYPE,
                  Collections.singleton(ENTITY_TYPE_DATASET),
                  AtlasTypeUtil.createOptionalAttrDef("locationUri", "string"),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_CREATE_TIME, "long"),
                  new AtlasAttributeDef(ATTR_RANDOM_TABLE,
                          AtlasBaseTypeDef.getArrayTypeName(SampleAppConstants.TABLE_TYPE),
                          true, AtlasAttributeDef.Cardinality.SET));
      }
  
      private AtlasEntityDef createTableDef() {
          return AtlasTypeUtil.createClassTypeDef(SampleAppConstants.TABLE_TYPE,
                  Collections.singleton(ENTITY_TYPE_DATASET),
                  createOptionalAttrDef(ATTR_CREATE_TIME, "long"),
                  createOptionalAttrDef(ATTR_LAST_ACCESS_TIME, "date"),
                  createOptionalAttrDef(ATTR_TEMPORARY, "boolean"),
                  createOptionalAttrDef(ATTR_TABLE_TYPE, ENUM_TABLE_TYPE),
                  createOptionalAttrDef(ATTR_SERDE1, STRUCT_TYPE_SERDE),
                  createOptionalAttrDef(ATTR_SERDE2, STRUCT_TYPE_SERDE));
      }
  
      private AtlasEntityDef createColumnDef() {
          return AtlasTypeUtil.createClassTypeDef(SampleAppConstants.COLUMN_TYPE,
                  Collections.singleton(ENTITY_TYPE_DATASET),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_DATA_TYPE, "string"),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_COMMENT, "string"));
      }
  
      private AtlasEntityDef createProcessDef() {
          return AtlasTypeUtil.createClassTypeDef(SampleAppConstants.PROCESS_TYPE,
                  Collections.singleton(ENTITY_TYPE_PROCESS),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_USERNAME, "string"),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_START_TIME, "long"),
                  AtlasTypeUtil.createOptionalAttrDef(ATTR_END_TIME, "long"),
                  AtlasTypeUtil.createRequiredAttrDef(ATTR_QUERY_TEXT, "string"),
                  AtlasTypeUtil.createRequiredAttrDef(ATTR_QUERY_PLAN, "string"),
                  AtlasTypeUtil.createRequiredAttrDef(ATTR_QUERY_ID, "string"),
                  AtlasTypeUtil.createRequiredAttrDef(ATTR_QUERY_GRAPH, "string"));
      }
  
      private AtlasStructDef createSerDeDef() {
          return AtlasTypeUtil.createStructTypeDef(SampleAppConstants.STRUCT_TYPE_SERDE,
                  AtlasTypeUtil.createRequiredAttrDef(SampleAppConstants.ATTR_NAME, "string"),
                  AtlasTypeUtil.createRequiredAttrDef(ATTR_SERDE, "string"));
      }
  
      private AtlasEnumDef createTableTypeEnumDef() {
          return new AtlasEnumDef(SampleAppConstants.ENUM_TABLE_TYPE,
                  SampleAppConstants.ATTR_DESCRIPTION,
                  Arrays.asList(new AtlasEnumDef.AtlasEnumElementDef("MANAGED", null, 1),
                          new AtlasEnumDef.AtlasEnumElementDef("EXTERNAL", null, 2)));
      }
  
      private List<AtlasClassificationDef> createClassificationDefs() {
          AtlasClassificationDef classification = createTraitTypeDef(SampleAppConstants.CLASSIFIED_TAG, Collections.<String>emptySet(), AtlasTypeUtil.createRequiredAttrDef("tag", "string"));
          AtlasClassificationDef pii            = createTraitTypeDef(SampleAppConstants.PII_TAG, Collections.<String>emptySet());
          AtlasClassificationDef finance        = createTraitTypeDef(SampleAppConstants.FINANCE_TAG, Collections.<String>emptySet());
          AtlasClassificationDef metric         = createTraitTypeDef(SampleAppConstants.METRIC_TAG, Collections.emptySet());
  
          return Arrays.asList(classification, pii, finance, metric);
      }
  
      private List<AtlasBusinessMetadataDef> createBusinessMetadataDefs() {
          String description = "description";
  
          Map<String, String> options = new HashMap<>();
  
          options.put("maxStrLength", "20");
          options.put("applicableEntityTypes", "[\"" + SampleAppConstants.DATABASE_TYPE + "\",\"" + SampleAppConstants.TABLE_TYPE + "\"]");
  
          AtlasBusinessMetadataDef bmWithAllTypes = createBusinessMetadataDef(SampleAppConstants.BUSINESS_METADATA_TYPE,
                  description,
                  "1.0",
                  createOptionalAttrDef(ATTR_ATTR1, ATLAS_TYPE_BOOLEAN, options, description),
                  createOptionalAttrDef(ATTR_ATTR2, ATLAS_TYPE_BYTE, options, description),
                  createOptionalAttrDef(ATTR_ATTR8, ATLAS_TYPE_STRING, options, description));
  
          AtlasBusinessMetadataDef bmWithAllTypesMV = createBusinessMetadataDef(SampleAppConstants.BUSINESS_METADATA_TYPE_MV,
                  description,
                  "1.0",
                  createOptionalAttrDef(ATTR_ATTR11, "array<boolean>", options, description),
                  createOptionalAttrDef(ATTR_ATTR18, "array<string>", options, description));
  
          return Arrays.asList(bmWithAllTypes, bmWithAllTypesMV);
      }
  
      private List<AtlasRelationshipDef> createAtlasRelationshipDef() {
          AtlasRelationshipDef dbTablesDef = createRelationshipTypeDef(SampleAppConstants.TABLE_DATABASE_TYPE, SampleAppConstants.TABLE_DATABASE_TYPE,
                  "1.0", AGGREGATION, AtlasRelationshipDef.PropagateTags.NONE,
                  createRelationshipEndDef(SampleAppConstants.TABLE_TYPE, "db", SINGLE, false),
                  createRelationshipEndDef(SampleAppConstants.DATABASE_TYPE, "tables", SET, true));
  
          AtlasRelationshipDef tableColumnsDef = createRelationshipTypeDef(SampleAppConstants.TABLE_COLUMNS_TYPE, SampleAppConstants.TABLE_COLUMNS_TYPE,
                  "1.0", COMPOSITION, AtlasRelationshipDef.PropagateTags.NONE,
                  createRelationshipEndDef(SampleAppConstants.TABLE_TYPE, "columns", SET, true),
                  createRelationshipEndDef(SampleAppConstants.COLUMN_TYPE, "table", SINGLE, false));
  
          return Arrays.asList(dbTablesDef, tableColumnsDef);
      }
  
      private AtlasTypesDef batchCreateTypes(AtlasTypesDef typesDef) throws AtlasServiceException {
          AtlasTypesDef typesToCreate = new AtlasTypesDef();
  
          for (AtlasEnumDef enumDef : typesDef.getEnumDefs()) {
              if (client.typeWithNameExists(enumDef.getName())) {
                  SampleApp.log(enumDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getEnumDefs().add(enumDef);
              }
          }
  
          for (AtlasStructDef structDef : typesDef.getStructDefs()) {
              if (client.typeWithNameExists(structDef.getName())) {
                  SampleApp.log(structDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getStructDefs().add(structDef);
              }
          }
  
          for (AtlasEntityDef entityDef : typesDef.getEntityDefs()) {
              if (client.typeWithNameExists(entityDef.getName())) {
                  SampleApp.log(entityDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getEntityDefs().add(entityDef);
              }
          }
  
          for (AtlasClassificationDef classificationDef : typesDef.getClassificationDefs()) {
              if (client.typeWithNameExists(classificationDef.getName())) {
                  SampleApp.log(classificationDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getClassificationDefs().add(classificationDef);
              }
          }
  
          for (AtlasRelationshipDef relationshipDef : typesDef.getRelationshipDefs()) {
              if (client.typeWithNameExists(relationshipDef.getName())) {
                  SampleApp.log(relationshipDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getRelationshipDefs().add(relationshipDef);
              }
          }
  
          for (AtlasBusinessMetadataDef businessMetadataDef : typesDef.getBusinessMetadataDefs()) {
              if (client.typeWithNameExists(businessMetadataDef.getName())) {
                  SampleApp.log(businessMetadataDef.getName() + ": type already exists. Skipping");
              } else {
                  typesToCreate.getBusinessMetadataDefs().add(businessMetadataDef);
              }
          }
  
          return client.createAtlasTypeDefs(typesToCreate);
      }
  }
  