Merge pull request #2894 from atlanhq/enable-maintenance

feat: add feature to enable maintenance mode from Metastore from configuration
atlanhq · Mar 29, 2024 · fcaac1e · fcaac1e
2 parents 568a750 + 8851e73
commit fcaac1e
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 1 deletion.
diff --git a/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java b/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java
@@ -111,7 +111,9 @@ public enum AtlasConfiguration {
     HERACLES_CLIENT_PAGINATION_SIZE("atlas.heracles.admin.resource-pagination-size", 100),
     HERACLES_API_SERVER_URL("atlas.heracles.api.service.url", "http://heracles-service.heracles.svc.cluster.local"),
 
-    INDEXSEARCH_ASYNC_SEARCH_KEEP_ALIVE_TIME_IN_SECONDS("atlas.indexsearch.async.search.keep.alive.time.in.seconds", 300);
+    INDEXSEARCH_ASYNC_SEARCH_KEEP_ALIVE_TIME_IN_SECONDS("atlas.indexsearch.async.search.keep.alive.time.in.seconds", 300),
+
+    ATLAS_MAINTENANCE_MODE("atlas.maintenance.mode", false);
 
 
     private static final Configuration APPLICATION_PROPERTIES;

diff --git a/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java b/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
@@ -274,6 +274,8 @@ public enum AtlasErrorCode {
     RUNTIME_EXCEPTION(500, "ATLAS-500-00-020", "Runtime exception {0}"),
     KEYCLOAK_INIT_FAILED(500, "ATLAS-500-00-022", "Failed to initialize keycloak client: {0}"),
 
+    MAINTENANCE_MODE_ENABLED(503, "ATLAS-503-00-001", "Atlas is in maintenance mode for this specific operation. Please try again later."),
+
     BATCH_SIZE_TOO_LARGE(406, "ATLAS-406-00-001", "Batch size is too large, please use a smaller batch size"),
 
 

diff --git a/repository/src/main/java/org/apache/atlas/tasks/TaskQueueWatcher.java b/repository/src/main/java/org/apache/atlas/tasks/TaskQueueWatcher.java
@@ -76,6 +76,11 @@ public void shutdown() {
 
     @Override
     public void run() {
+        boolean isMaintenanceMode = AtlasConfiguration.ATLAS_MAINTENANCE_MODE.getBoolean();
+        if (isMaintenanceMode) {
+            LOG.info("TaskQueueWatcher: Maintenance mode is enabled, new tasks will not be loaded into the queue until next restart");
+            return;
+        }
         shouldRun.set(true);
 
         if (LOG.isDebugEnabled()) {

diff --git a/webapp/src/main/java/org/apache/atlas/web/filters/ActiveServerFilter.java b/webapp/src/main/java/org/apache/atlas/web/filters/ActiveServerFilter.java
@@ -18,6 +18,10 @@
 
 package org.apache.atlas.web.filters;
 
+import org.apache.atlas.AtlasConfiguration;
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.web.service.ActiveInstanceState;
 import org.apache.atlas.web.service.ServiceState;
 import org.slf4j.Logger;
@@ -37,6 +41,7 @@
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.HttpHeaders;
 import java.io.IOException;
+import java.util.HashMap;
 
 /**
  * A servlet {@link Filter} that redirects web requests from a passive Atlas server instance to an active one.
@@ -52,6 +57,9 @@ public class ActiveServerFilter implements Filter {
     private static final Logger LOG = LoggerFactory.getLogger(ActiveServerFilter.class);
     private static final String MIGRATION_STATUS_STATIC_PAGE = "migration-status.html";
 
+    private static final String[] WHITELISTED_APIS_SIGNATURE = {"search", "lineage", "auditSearch", "accessors"
+        , "evaluator"};
+
     private final ActiveInstanceState activeInstanceState;
     private ServiceState serviceState;
 
@@ -78,6 +86,18 @@ public void init(FilterConfig filterConfig) throws ServletException {
     @Override
     public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse,
                          FilterChain filterChain) throws IOException, ServletException {
+        // If maintenance mode is enabled, return a 503
+        if (AtlasConfiguration.ATLAS_MAINTENANCE_MODE.getBoolean()) {
+            // Block all the POST, PUT, DELETE operations
+            HttpServletRequest request = (HttpServletRequest) servletRequest;
+            HttpServletResponse response = (HttpServletResponse) servletResponse;
+            if (isBlockedMethod(request.getMethod()) && !isWhitelistedAPI(request.getRequestURI())) {
+                LOG.error("Maintenance mode enabled. Blocking request: {}", request.getRequestURI());
+                sendMaintenanceModeResponse(response);
+                return; // Stop further processing
+            }
+        }
+
         if (isFilteredURI(servletRequest)) {
             LOG.debug("Is a filtered URI: {}. Passing request downstream.",
                     ((HttpServletRequest)servletRequest).getRequestURI());
@@ -111,6 +131,7 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
 
     final String adminUriNotFiltered[] = { "/admin/export", "/admin/import", "/admin/importfile", "/admin/audits",
             "/admin/purge", "/admin/expimp/audit", "/admin/metrics", "/admin/server", "/admin/audit/", "admin/tasks"};
+
     private boolean isFilteredURI(ServletRequest servletRequest) {
         HttpServletRequest httpServletRequest = (HttpServletRequest) servletRequest;
         String requestURI = httpServletRequest.getRequestURI();
@@ -129,6 +150,43 @@ private boolean isFilteredURI(ServletRequest servletRequest) {
         }
     }
 
+    private boolean isWhitelistedAPI(String requestURI) {
+        for (String api : WHITELISTED_APIS_SIGNATURE) {
+            if (requestURI.contains(api)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private void sendMaintenanceModeResponse(HttpServletResponse response) throws IOException {
+        response.setStatus(HttpServletResponse.SC_SERVICE_UNAVAILABLE);
+        response.setContentType("application/json");
+        response.setCharacterEncoding("UTF-8");
+
+        AtlasBaseException serverException = new AtlasBaseException(AtlasErrorCode.MAINTENANCE_MODE_ENABLED,
+                AtlasErrorCode.MAINTENANCE_MODE_ENABLED.getFormattedErrorMessage());
+
+        HashMap<String, Object> errorMap = new HashMap<>();
+        errorMap.put("errorCode", serverException.getAtlasErrorCode().getErrorCode());
+        errorMap.put("errorMessage", serverException.getMessage());
+
+        String jsonResponse = AtlasType.toJson(errorMap);
+        response.getOutputStream().write(jsonResponse.getBytes());
+        response.getOutputStream().flush();
+    }
+
+    private boolean isBlockedMethod(String method) {
+        switch (method) {
+            case HttpMethod.POST:
+            case HttpMethod.PUT:
+            case HttpMethod.DELETE:
+                return true;
+            default:
+                return false;
+        }
+    }
+
     private boolean isRootURI(ServletRequest servletRequest) {
         HttpServletRequest httpServletRequest = (HttpServletRequest) servletRequest;
         String requestURI = httpServletRequest.getRequestURI();