ATLAS-4922: Atlas Async Import using Kafka [4] - Addressing few of the PR comments

Author: Chandra Kanth Peravelli

intg/src/main/java/org/apache/atlas/AtlasErrorCode.java

@@ -205,7 +205,7 @@ public enum AtlasErrorCode {
     FILE_NAME_NOT_FOUND(404, "ATLAS-404-00-014", "File name should not be blank"),
     NO_TYPE_NAME_ON_VERTEX(404, "ATLAS-404-00-015", "No typename found for given entity with guid: {0}"),
     NO_LINEAGE_CONSTRAINTS_FOR_GUID(404, "ATLAS-404-00-016", "No lineage constraints found for requested entity with guid : {0}"),
-    IMPORT_NOT_FOUND(404, "ATLAS-404-00-017", "Give import request {0} is not found"),
+    IMPORT_NOT_FOUND(404, "ATLAS-404-00-017", "Import {0} is not found"),
 
     METHOD_NOT_ALLOWED(405, "ATLAS-405-00-001", "Error 405 - The request method {0} is inappropriate for the URL: {1}"),
 
@@ -224,7 +224,7 @@ public enum AtlasErrorCode {
     GLOSSARY_IMPORT_FAILED(409, "ATLAS-409-00-011", "Glossary import failed"),
     METRICSSTAT_ALREADY_EXISTS(409, "ATLAS-409-00-012", "Metric Statistics already collected at {0}"),
     PENDING_TASKS_ALREADY_IN_PROGRESS(409, "ATLAS-409-00-013", "There are already {0} pending tasks in queue"),
-    IMPORT_ALREADY_IN_PROGRESS(409, "ATLAS-409-00-016", "Given import request {0} is already in progress or completed"),
+    ABORT_IMPORT(409, "ATLAS-409-00-016", "Import id={0} is currently in state {1}, cannot be aborted"),
 
     // All internal errors go here
     INTERNAL_ERROR(500, "ATLAS-500-00-001", "Internal server error {0}"),
@@ -252,9 +252,9 @@ public enum AtlasErrorCode {
     NOTIFICATION_EXCEPTION(500, "ATLAS-500-00-018", "{0}"),
     IMPORT_UPDATE_FAILED(500, "ATLAS-500-00-019", "Failed to update import: {0}"),
     IMPORT_REGISTRATION_FAILED(500, "ATLAS-500-00-020", "Failed to register import request"),
-    IMPORT_FAILED(500, "ATLAS-500-00-021", "Given import {0} failed"),
-    ABORT_IMPORT_FAILED(500, "ATLAS-500-00-022", "Failed to abort given import {0}"),
-    IMPORT_QUEUEING_FAILED(500, "ATLAS-500-00-023", "Failed to add given import {0} to request queue, please try again later");
+    IMPORT_FAILED(500, "ATLAS-500-00-021", "Import with id={0} failed"),
+    ABORT_IMPORT_FAILED(500, "ATLAS-500-00-022", "Failed to abort import id={0}"),
+    IMPORT_QUEUEING_FAILED(500, "ATLAS-500-00-023", "Failed to add import id={0} to request queue, please try again later");
 
     private static final Logger LOG = LoggerFactory.getLogger(AtlasErrorCode.class);
     private final        String errorCode;

intg/src/main/java/org/apache/atlas/model/impexp/AtlasAsyncImportRequest.java

@@ -21,8 +21,6 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonInclude;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import org.apache.atlas.model.AtlasBaseModelObject;
 import org.apache.atlas.utils.AtlasEntityUtil;
 
@@ -42,7 +40,7 @@
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_ONLY;
 
 @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
-@JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
+@JsonInclude(JsonInclude.Include.NON_NULL)
 @JsonIgnoreProperties(ignoreUnknown = true)
 public class AtlasAsyncImportRequest extends AtlasBaseModelObject implements Serializable {
     private static final long serialVersionUID = 1L;
@@ -77,35 +75,20 @@ public String toString() {
 
     @JsonIgnore
     private String requestId;
-
-    @JsonProperty("importId")
-    private String importId;
-
-    @JsonProperty("status")
-    private ImportStatus status;
-
     @JsonIgnore
     private int skipTo;
 
-    @JsonInclude(JsonInclude.Include.NON_NULL)
-    @JsonProperty("atlasImportResult")
-    private AtlasImportResult atlasImportResult;
-
-    @JsonProperty("importDetails")
+    private String importId;
+    private ImportStatus status;
     private ImportDetails importDetails = new ImportDetails();
-
-    @JsonProperty("receivedAt")
     private long receivedAt;
-
-    @JsonProperty("stagedAt")
     private long stagedAt;
-
-    @JsonProperty("startedProcessingAt")
     private long startedProcessingAt;
-
-    @JsonProperty("completedAt")
     private long completedAt;
 
+    @JsonInclude(JsonInclude.Include.NON_NULL)
+    private AtlasImportResult atlasImportResult;
+
     public AtlasAsyncImportRequest() {}
 
     public AtlasAsyncImportRequest(String guid) {

intg/src/main/java/org/apache/atlas/model/notification/ImportNotification.java

@@ -19,8 +19,8 @@
 
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
-import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 
@@ -37,7 +37,7 @@
  * Class representing atlas import notification, extending HookNotification.
  */
 @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
-@JsonSerialize(include = JsonSerialize.Inclusion.ALWAYS)
+@JsonInclude
 @JsonIgnoreProperties(ignoreUnknown = true)
 @XmlRootElement
 @XmlAccessorType(XmlAccessType.PROPERTY)
@@ -73,7 +73,7 @@ public StringBuilder toString(StringBuilder sb) {
      * Notification for type definitions import
      */
     @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
-    @JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
+    @JsonInclude(JsonInclude.Include.NON_NULL)
     @JsonIgnoreProperties(ignoreUnknown = true)
     @XmlRootElement
     @XmlAccessorType(XmlAccessType.PROPERTY)
@@ -113,7 +113,7 @@ public String toString() {
      * Notification for entities import
      */
     @JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
-    @JsonSerialize(include = JsonSerialize.Inclusion.ALWAYS)
+    @JsonInclude
     @JsonIgnoreProperties(ignoreUnknown = true)
     @XmlRootElement
     @XmlAccessorType(XmlAccessType.PROPERTY)

notification/src/main/java/org/apache/atlas/kafka/AtlasKafkaConsumer.java

@@ -61,20 +61,12 @@ public AtlasKafkaConsumer(AtlasNotificationMessageDeserializer<T> deserializer,
 
     @Override
     public Set<TopicPartition> getTopicPartition() {
-        Set<TopicPartition> ret = null;
-        if (kafkaConsumer != null) {
-            ret = kafkaConsumer.assignment();
-        }
-        return ret;
+        return kafkaConsumer != null ? kafkaConsumer.assignment() : null;
     }
 
     @Override
     public Set<String> subscription() {
-        Set<String> ret = null;
-        if (kafkaConsumer != null) {
-            ret = kafkaConsumer.subscription();
-        }
-        return ret;
+        return kafkaConsumer != null ? kafkaConsumer.subscription() : null;
     }
 
     @Override

notification/src/main/java/org/apache/atlas/kafka/KafkaNotification.java

@@ -45,14 +45,12 @@
 import javax.inject.Inject;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Future;
-import java.util.stream.Stream;
 
 import static org.apache.atlas.security.SecurityProperties.TLS_ENABLED;
 import static org.apache.atlas.security.SecurityProperties.TRUSTSTORE_PASSWORD_KEY;
@@ -83,7 +81,7 @@ public class KafkaNotification extends AbstractNotification implements Service {
     private static final boolean SORT_NOT_NEEDED = false;
 
     private static final Map<NotificationType, String> PRODUCER_TOPIC_MAP    = new HashMap<>();
-    private static Map<NotificationType, String[]> CONSUMER_TOPICS_MAP = new HashMap<>();
+    private static final Map<NotificationType, List<String>> CONSUMER_TOPICS_MAP = new HashMap<>();
 
     private final Properties                                 properties;
     private final Long                                       pollTimeOutMs;
@@ -154,7 +152,7 @@ protected KafkaNotification(Properties properties) {
         LOG.info("<== KafkaNotification()");
     }
 
-    public static String[] trimAndPurge(String[] strings) {
+    public static List<String> trimAndPurge(String[] strings) {
         List<String> ret = new ArrayList<>();
 
         if (strings != null) {
@@ -167,7 +165,7 @@ public static String[] trimAndPurge(String[] strings) {
             }
         }
 
-        return ret.toArray(new String[ret.size()]);
+        return ret;
     }
 
     @Override
@@ -188,10 +186,7 @@ public void stop() {
 
     @Override
     public <T> List<NotificationConsumer<T>> createConsumers(NotificationType notificationType, int numConsumers) {
-        boolean enableAutoCommit = Boolean.parseBoolean(properties.getProperty("enable.auto.commit", properties.getProperty("auto.commit.enable", "false")));
-        if (notificationType.equals(NotificationType.ASYNC_IMPORT)) {
-            enableAutoCommit = true;
-        }
+        boolean enableAutoCommit = notificationType.equals(NotificationType.ASYNC_IMPORT) || Boolean.parseBoolean(properties.getProperty("enable.auto.commit", properties.getProperty("auto.commit.enable", "false")));
         return createConsumers(notificationType, numConsumers, enableAutoCommit);
     }
 
@@ -215,13 +210,18 @@ public void close() {
     }
 
     @Override
-    public void closeConsumer(NotificationType notificationType) {
-        List<KafkaConsumer> notificationConsumers = this.consumers.get(notificationType);
-        for (final KafkaConsumer consumer : notificationConsumers) {
-            consumer.unsubscribe();
-            consumer.close();
-        }
-        this.consumers.remove(notificationType);
+    public void closeConsumer(NotificationType notificationTypeToClose, String topic) {
+        this.consumers.computeIfPresent(notificationTypeToClose, (notificationType, notificationConsumers) -> {
+            notificationConsumers.removeIf(consumer -> {
+                if (consumer.subscription().contains(topic)) {
+                    consumer.unsubscribe();
+                    consumer.close();
+                    return true;
+                }
+                return false;
+            });
+            return notificationConsumers.isEmpty() ? null : notificationConsumers;
+        });
     }
 
     // ----- NotificationInterface -------------------------------------------
@@ -243,16 +243,16 @@ public boolean isReady(NotificationType notificationType) {
     public <T> List<NotificationConsumer<T>> createConsumers(NotificationType notificationType, int numConsumers, boolean autoCommitEnabled) {
         LOG.info("==> KafkaNotification.createConsumers(notificationType={}, numConsumers={}, autoCommitEnabled={})", notificationType, numConsumers, autoCommitEnabled);
 
-        String[] topics = CONSUMER_TOPICS_MAP.get(notificationType);
+        List<String> topics = CONSUMER_TOPICS_MAP.get(notificationType);
 
-        if (numConsumers < topics.length) {
-            LOG.warn("consumers count {} is fewer than number of topics {}. Creating {} consumers, so that consumer count is equal to number of topics.", numConsumers, topics.length, topics.length);
+        if (numConsumers < topics.size()) {
+            LOG.warn("consumers count {} is fewer than number of topics {}. Creating {} consumers, so that consumer count is equal to number of topics.", numConsumers, topics.size(), topics.size());
 
-            numConsumers = topics.length;
-        } else if (numConsumers > topics.length) {
-            LOG.warn("consumers count {} is higher than number of topics {}. Creating {} consumers, so that consumer count is equal to number of topics", numConsumers, topics.length, topics.length);
+            numConsumers = topics.size();
+        } else if (numConsumers > topics.size()) {
+            LOG.warn("consumers count {} is higher than number of topics {}. Creating {} consumers, so that consumer count is equal to number of topics", numConsumers, topics.size(), topics.size());
 
-            numConsumers = topics.length;
+            numConsumers = topics.size();
         }
 
         List<KafkaConsumer> notificationConsumers = this.consumers.get(notificationType);
@@ -318,11 +318,7 @@ public Properties getConsumerProperties(NotificationType notificationType) {
         String groupId = properties.getProperty(notificationType.toString().toLowerCase() + "." + CONSUMER_GROUP_ID_PROPERTY);
 
         if (StringUtils.isEmpty(groupId)) {
-            if (!notificationType.equals(NotificationType.ASYNC_IMPORT)) {
-                groupId = "atlas";
-            } else {
-                groupId = "atlas-import";
-            }
+            groupId = notificationType.equals(NotificationType.ASYNC_IMPORT) ? "atlas-import" : "atlas";
         }
 
         if (StringUtils.isEmpty(groupId)) {
@@ -343,8 +339,8 @@ public KafkaConsumer getOrCreateKafkaConsumer(KafkaConsumer existingConsumer, Pr
 
         try {
             if (ret == null || !isKafkaConsumerOpen(ret)) {
-                String[] topics = CONSUMER_TOPICS_MAP.get(notificationType);
-                String   topic  = topics[idxConsumer % topics.length];
+                List<String> topics = CONSUMER_TOPICS_MAP.get(notificationType);
+                String   topic  = topics.get(idxConsumer % topics.size());
 
                 LOG.debug("Creating new KafkaConsumer for topic : {}, index : {}", topic, idxConsumer);
 
@@ -452,39 +448,29 @@ private KafkaProducer getOrCreateProducerByCriteria(Object producerCriteria, Map
 
     @Override
     public void addTopicToNotificationType(NotificationType notificationType, String topic) {
-        String[] topics = CONSUMER_TOPICS_MAP.get(notificationType);
-        String[] updatedTopics;
-        if (topics == null) {
-            updatedTopics = new String[] {topic};
-        } else {
-            updatedTopics = Stream.concat(Arrays.stream(topics), Stream.of(topic)).toArray(String[]::new);
-        }
-        CONSUMER_TOPICS_MAP.put(notificationType, updatedTopics);
+        CONSUMER_TOPICS_MAP.computeIfAbsent(notificationType, k -> new ArrayList<>()).add(topic);
     }
 
     @Override
     public void closeProducer(NotificationType notificationType, String topic) {
-        KafkaProducer producerToClose = producersByTopic.get(topic);
-        if (producerToClose != null) {
-            producersByTopic.remove(topic);
-            producerToClose.close();
-        }
+        producersByTopic.computeIfPresent(topic, (key, producer) -> {
+            // Close the KafkaProducer before removal
+            producer.close();
+            // Returning null removes the key from the map
+            return null;
+        });
         PRODUCER_TOPIC_MAP.remove(notificationType, topic);
     }
 
     @Override
-    public void deleteTopics(NotificationType notificationType, String topicName) {
+    public void deleteTopic(NotificationType notificationType, String topicName) {
         try (AdminClient adminClient = AdminClient.create(this.properties)) {
             adminClient.deleteTopics(Collections.singleton(topicName));
         }
-        String[] topics = CONSUMER_TOPICS_MAP.get(notificationType);
-        String[] updatedTopics;
-        if (topics == null) {
-            updatedTopics = new String[] {};
-        } else {
-            updatedTopics = Arrays.stream(topics).filter(topic -> !topic.equals(topicName)).toArray(String[]::new);
-        }
-        CONSUMER_TOPICS_MAP.put(notificationType, updatedTopics);
+        CONSUMER_TOPICS_MAP.computeIfPresent(notificationType, (key, topics) -> {
+            topics.remove(topicName);
+            return topics.isEmpty() ? null : topics;
+        });
     }
 
     // kafka-client doesn't have method to check if consumer is open, hence checking list topics and catching exception

notification/src/main/java/org/apache/atlas/notification/AbstractNotification.java

@@ -237,8 +237,8 @@ public <T> void send(NotificationType type, List<T> messages, MessageSource sour
     @Override
     public <T> void send(String topic, List<T> messages, MessageSource source) throws NotificationException {
         List<String> strMessages = new ArrayList<>(messages.size());
-        for (int index = 0; index < messages.size(); index++) {
-            createNotificationMessages(messages.get(index), strMessages, source);
+        for (T message : messages) {
+            createNotificationMessages(message, strMessages, source);
         }
         sendInternal(topic, strMessages);
     }

notification/src/main/java/org/apache/atlas/notification/NotificationInterface.java

@@ -131,14 +131,16 @@ default void closeProducer(NotificationType notificationType, String topic) {}
      * @param notificationType The type of notification related to the topic.
      * @param topicName The name of the topic to be deleted.
      */
-    default void deleteTopics(NotificationType notificationType, String topicName) {}
+    default void deleteTopic(NotificationType notificationType, String topicName) {}
 
     /**
      * Closes the consumer associated with the specified notification type.
      *
      * @param notificationType The type of notification for which the consumer is to be closed.
+     * @param topic The consumer to close with assignment.
+     *
      */
-    default void closeConsumer(NotificationType notificationType) {}
+    default void closeConsumer(NotificationType notificationType, String topic) {}
 
     /**
      * Atlas notification types.
@@ -150,11 +152,11 @@ enum NotificationType {
         // Notifications from the Atlas integration hooks - unsorted.
         HOOK_UNSORTED(new HookMessageDeserializer()),
 
-        // Notifications from Atlas async importer
-        ASYNC_IMPORT(new HookMessageDeserializer()),
-
         // Notifications to entity change consumers.
-        ENTITIES(new EntityMessageDeserializer());
+        ENTITIES(new EntityMessageDeserializer()),
+
+        // Notifications from Atlas async importer
+        ASYNC_IMPORT(new HookMessageDeserializer());
 
         private final AtlasNotificationMessageDeserializer deserializer;
 

notification/src/test/java/org/apache/atlas/kafka/KafkaConsumerTest.java

@@ -58,7 +58,7 @@ public class KafkaConsumerTest {
     private static final String TRAIT_NAME = "MyTrait";
 
     private static final String   ATLAS_HOOK_TOPIC           = AtlasConfiguration.NOTIFICATION_HOOK_TOPIC_NAME.getString();
-    private static final String[] ATLAS_HOOK_CONSUMER_TOPICS = KafkaNotification.trimAndPurge(AtlasConfiguration.NOTIFICATION_HOOK_CONSUMER_TOPIC_NAMES.getStringArray(ATLAS_HOOK_TOPIC));
+    private static final List<String> ATLAS_HOOK_CONSUMER_TOPICS = KafkaNotification.trimAndPurge(AtlasConfiguration.NOTIFICATION_HOOK_CONSUMER_TOPIC_NAMES.getStringArray(ATLAS_HOOK_TOPIC));
 
     @Mock
     private KafkaConsumer<String, String> kafkaConsumer;