Merge pull request #2998 from atlanhq/maintenance-mode-redis

feat: add feature flag storage based on Redis and maintenance mode based on feature flag
atlanhq · May 7, 2024 · 4356129 · 4356129
2 parents 6f22505 + 5f89496
commit 4356129
Show file tree

Hide file tree

Showing 14 changed files with 144 additions and 29 deletions.
diff --git a/addons/policies/bootstrap_admin_policies.json b/addons/policies/bootstrap_admin_policies.json
@@ -49,6 +49,31 @@
           "admin-task-cud"
         ]
       }
+    },
+    {
+      "typeName": "AuthPolicy",
+      "attributes": {
+        "name": "ADMIN_ALLOW_FEATURE_FLAG_CUD",
+        "qualifiedName": "ADMIN_ALLOW_FEATURE_FLAG_CUD",
+        "policyCategory": "bootstrap",
+        "policySubCategory": "default",
+        "policyServiceName": "atlas",
+        "policyType": "allow",
+        "policyPriority": 1,
+        "policyUsers": [
+          "service-account-atlan-argo",
+          "service-account-atlan-backend"
+        ],
+        "policyGroups": [],
+        "policyRoles": [],
+        "policyResourceCategory": "ADMIN",
+        "policyResources": [
+          "atlas-service:*"
+        ],
+        "policyActions": [
+          "admin-featureFlag-cud"
+        ]
+      }
     }
   ]
 }
diff --git a/auth-agents-common/src/main/resources/service-defs/atlas-servicedef-atlas.json b/auth-agents-common/src/main/resources/service-defs/atlas-servicedef-atlas.json
@@ -450,6 +450,11 @@
 			"itemId": 24,
 			"name": "admin-task-cud",
 			"label": "Admin task CUD API"
+		},
+		{
+			"itemId": 25,
+			"name": "admin-featureFlag-cud",
+			"label": "Admin featureflag CUD API"
 		}
 
 	],

diff --git a/authorization/src/main/java/org/apache/atlas/authorize/AtlasPrivilege.java b/authorization/src/main/java/org/apache/atlas/authorize/AtlasPrivilege.java
@@ -51,7 +51,9 @@ public enum AtlasPrivilege {
      ADMIN_ENTITY_AUDITS("admin-entity-audits"),
      ADMIN_REPAIR_INDEX("admin-repair-index"),
 
-     ADMIN_TASK_CUD("admin-task-cud");
+     ADMIN_TASK_CUD("admin-task-cud"),
+
+     ADMIN_FEATURE_FLAG_CUD("admin-featureFlag-cud");
      private final String type;
 
      AtlasPrivilege(String actionType){

diff --git a/common/src/main/java/org/apache/atlas/repository/Constants.java b/common/src/main/java/org/apache/atlas/repository/Constants.java
@@ -19,7 +19,9 @@
 package org.apache.atlas.repository;
 
 import org.apache.atlas.ApplicationProperties;
+import org.apache.atlas.AtlasConfiguration;
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.service.FeatureFlagStore;
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
@@ -450,6 +452,20 @@ private static String getEncodedTypePropertyKey(String defaultKey) {
         }
     }
 
+    public static String getESIndex() {
+        String indexSuffix  = null;
+        if(AtlasConfiguration.ATLAS_MAINTENANCE_MODE.getBoolean()) {
+            try {
+                if (FeatureFlagStore.evaluate("use_temp_es_index", "true")) {
+                    indexSuffix = "_temp";
+                }
+            } catch (Exception e) {
+                LOG.error("Failed to evaluate feature flag with error", e);
+            }
+        }
+        return indexSuffix == null ? VERTEX_INDEX_NAME : VERTEX_INDEX_NAME + indexSuffix;
+    }
+
     public static String getStaticFileAsString(String fileName) throws IOException {
         String atlasHomeDir  = System.getProperty("atlas.home");
         atlasHomeDir = StringUtils.isEmpty(atlasHomeDir) ? "." : atlasHomeDir;

diff --git a/common/src/main/java/org/apache/atlas/service/FeatureFlagStore.java b/common/src/main/java/org/apache/atlas/service/FeatureFlagStore.java
@@ -0,0 +1,45 @@
+package org.apache.atlas.service;
+
+import org.apache.atlas.service.redis.RedisService;
+import org.apache.commons.lang.StringUtils;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.stereotype.Component;
+
+@Component
+public class FeatureFlagStore {
+    private static RedisService redisService = null;
+    public FeatureFlagStore(@Qualifier("redisServiceImpl") RedisService redisService) {
+        FeatureFlagStore.redisService = redisService;
+    }
+
+    public static boolean evaluate(String key, String expectedValue) {
+        boolean ret = false;
+        try{
+            if (StringUtils.isEmpty(key) || StringUtils.isEmpty(expectedValue))
+                return ret;
+            String value = redisService.getValue(addFeatureFlagNamespace(key));
+            ret = StringUtils.equals(value, expectedValue);
+        } catch (Exception e) {
+            return ret;
+        }
+        return ret;
+    }
+
+    public static void setFlag(String key, String value) {
+        if (StringUtils.isEmpty(key) || StringUtils.isEmpty(value))
+            return;
+
+        redisService.putValue(addFeatureFlagNamespace(key), value);
+    }
+
+    public static void deleteFlag(String key) {
+        if (StringUtils.isEmpty(key))
+            return;
+
+        redisService.removeValue(addFeatureFlagNamespace(key));
+    }
+
+    private static String addFeatureFlagNamespace(String key) {
+        return "ff:"+key;
+    }
+}
diff --git a/common/src/main/java/org/apache/atlas/service/redis/AbstractRedisService.java b/common/src/main/java/org/apache/atlas/service/redis/AbstractRedisService.java
@@ -81,7 +81,14 @@ public String getValue(String key) {
     @Override
     public String putValue(String key, String value) {
         // Put the value in the redis cache with TTL
-        redisCacheClient.getBucket(convertToNamespace(key)).set(value, 30, TimeUnit.SECONDS);
+        redisCacheClient.getBucket(convertToNamespace(key)).set(value);
+        return value;
+    }
+
+    @Override
+    public String putValue(String key, String value, int timeout) {
+        // Put the value in the redis cache with TTL
+        redisCacheClient.getBucket(convertToNamespace(key)).set(value, timeout, TimeUnit.SECONDS);
         return value;
     }
 

diff --git a/common/src/main/java/org/apache/atlas/service/redis/NoRedisServiceImpl.java b/common/src/main/java/org/apache/atlas/service/redis/NoRedisServiceImpl.java
@@ -35,7 +35,7 @@ public String getValue(String key) {
     }
 
     @Override
-    public String putValue(String key, String value) {
+    public String putValue(String key, String value, int timeout) {
         return null;
     }
 

diff --git a/common/src/main/java/org/apache/atlas/service/redis/RedisService.java b/common/src/main/java/org/apache/atlas/service/redis/RedisService.java
@@ -12,6 +12,8 @@ public interface RedisService {
 
   String putValue(String key, String value);
 
+  String putValue(String key, String value, int timeout);
+
   void removeValue(String key);
 
   Logger getLogger();

diff --git a/common/src/main/java/org/apache/atlas/service/redis/RedisServiceLocalImpl.java b/common/src/main/java/org/apache/atlas/service/redis/RedisServiceLocalImpl.java
@@ -28,7 +28,7 @@ public String getValue(String key) {
     }
 
     @Override
-    public String putValue(String key, String value) {
+    public String putValue(String key, String value, int timeout) {
         return null;
     }
 

diff --git a/...hdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/SearchContextCache.java b/...hdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/SearchContextCache.java
@@ -24,7 +24,7 @@ public static void put(String key, Integer sequence, String esAsyncId) {
        try {
            // Build the string in format `sequence/esAsyncId` and store it in redis
            String val = sequence + "/" + esAsyncId;
-           redisService.putValue(key, val);
+           redisService.putValue(key, val, 30);
        } finally {
            RequestContext.get().endMetricRecord(metric);
        }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -45,6 +45,7 @@
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
 import org.apache.atlas.repository.userprofile.UserProfileService;
 import org.apache.atlas.searchlog.ESSearchLogger;
+import org.apache.atlas.service.FeatureFlagStore;
 import org.apache.atlas.stats.StatsClient;
 import org.apache.atlas.type.*;
 import org.apache.atlas.type.AtlasBuiltInTypes.AtlasObjectIdType;
@@ -75,9 +76,7 @@
 import static org.apache.atlas.SortOrder.ASCENDING;
 import static org.apache.atlas.model.instance.AtlasEntity.Status.ACTIVE;
 import static org.apache.atlas.model.instance.AtlasEntity.Status.DELETED;
-import static org.apache.atlas.repository.Constants.ASSET_ENTITY_TYPE;
-import static org.apache.atlas.repository.Constants.OWNER_ATTRIBUTE;
-import static org.apache.atlas.repository.Constants.VERTEX_INDEX_NAME;
+import static org.apache.atlas.repository.Constants.*;
 import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.BASIC_SEARCH_STATE_FILTER;
 import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.TO_RANGE_LIST;
 
@@ -1134,8 +1133,10 @@ public List<AtlasEntityHeader> searchUsingTermQualifiedName(int from, int size,
     }
 
     private String getIndexName(IndexSearchParams params) throws AtlasBaseException {
+        String vertexIndexName = getESIndex();
+
         if (StringUtils.isEmpty(params.getPersona()) && StringUtils.isEmpty(params.getPurpose())) {
-            return VERTEX_INDEX_NAME;
+            return vertexIndexName;
         }
 
         String qualifiedName = "";
@@ -1151,7 +1152,7 @@ private String getIndexName(IndexSearchParams params) throws AtlasBaseException
         if (StringUtils.isNotEmpty(aliasName)) {
             if(params.isAccessControlExclusive()) {
                 accessControlExclusiveDsl(params, aliasName);
-                aliasName = aliasName+","+VERTEX_INDEX_NAME;
+                aliasName = aliasName+","+vertexIndexName;
             }
             return aliasName;
         } else {

diff --git a/repository/src/main/java/org/apache/atlas/repository/util/AccessControlUtils.java b/repository/src/main/java/org/apache/atlas/repository/util/AccessControlUtils.java
@@ -45,15 +45,7 @@
 import static org.apache.atlas.AtlasErrorCode.ACCESS_CONTROL_ALREADY_EXISTS;
 import static org.apache.atlas.AtlasErrorCode.DISABLED_OPERATION;
 import static org.apache.atlas.AtlasErrorCode.OPERATION_NOT_SUPPORTED;
-import static org.apache.atlas.repository.Constants.ATTR_ADMIN_GROUPS;
-import static org.apache.atlas.repository.Constants.ATTR_ADMIN_ROLES;
-import static org.apache.atlas.repository.Constants.ATTR_ADMIN_USERS;
-import static org.apache.atlas.repository.Constants.ATTR_TENANT_ID;
-import static org.apache.atlas.repository.Constants.CONNECTION_ENTITY_TYPE;
-import static org.apache.atlas.repository.Constants.DEFAULT_TENANT_ID;
-import static org.apache.atlas.repository.Constants.NAME;
-import static org.apache.atlas.repository.Constants.QUALIFIED_NAME;
-import static org.apache.atlas.repository.Constants.VERTEX_INDEX_NAME;
+import static org.apache.atlas.repository.Constants.*;
 import static org.apache.atlas.repository.util.AtlasEntityUtils.getListAttribute;
 import static org.apache.atlas.repository.util.AtlasEntityUtils.getQualifiedName;
 import static org.apache.atlas.repository.util.AtlasEntityUtils.getStringAttribute;
@@ -379,7 +371,8 @@ public static void validateUniquenessByTags(AtlasGraph graph, List<String> tags,
 
     private static boolean hasMatchingVertex(AtlasGraph graph, List<String> newTags,
                                                IndexSearchParams indexSearchParams) throws AtlasBaseException {
-        AtlasIndexQuery indexQuery = graph.elasticsearchQuery(VERTEX_INDEX_NAME);
+        String vertexIndexName = getESIndex();
+        AtlasIndexQuery indexQuery = graph.elasticsearchQuery(vertexIndexName);
 
         DirectIndexQueryResult indexQueryResult = indexQuery.vertices(indexSearchParams);
         Iterator<AtlasIndexQuery.Result> iterator = indexQueryResult.getIterator();

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
@@ -21,6 +21,7 @@
 import org.apache.atlas.AtlasConfiguration;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.service.FeatureFlagStore;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.web.service.ActiveInstanceState;
 import org.apache.atlas.web.service.ServiceState;
@@ -56,9 +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"};
+        , "evaluator", "featureFlag"};
+    private static final String DISABLE_WRITE_FLAG = "disable_writes";
 
     private final ActiveInstanceState activeInstanceState;
     private ServiceState serviceState;
@@ -88,13 +89,15 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
                          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 (FeatureFlagStore.evaluate(DISABLE_WRITE_FLAG, "true")) {
+                // 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
+                }
             }
         }
 

diff --git a/webapp/src/main/java/org/apache/atlas/web/resources/AdminResource.java b/webapp/src/main/java/org/apache/atlas/web/resources/AdminResource.java
@@ -56,6 +56,7 @@
 import org.apache.atlas.repository.impexp.ZipSink;
 import org.apache.atlas.repository.patches.AtlasPatchManager;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.service.FeatureFlagStore;
 import org.apache.atlas.service.metrics.MetricsRegistry;
 import org.apache.atlas.services.MetricsService;
 import org.apache.atlas.tasks.TaskManagement;
@@ -930,6 +931,21 @@ public Map getDebugMetrics() {
         return debugMetricsRESTSink.getMetrics();
     }
 
+    @POST
+    @Path("featureFlag")
+    @Produces(MediaType.APPLICATION_JSON)
+    public void setFeatureFlag(@QueryParam("key") String key, @QueryParam("value") String value) throws AtlasBaseException {
+        AtlasAuthorizationUtils.verifyAccess(new AtlasAdminAccessRequest(AtlasPrivilege.ADMIN_FEATURE_FLAG_CUD), "featureFlag");
+        FeatureFlagStore.setFlag(key, value);
+    }
+
+    @DELETE
+    @Path("featureFlag/{flag}")
+    @Produces(MediaType.APPLICATION_JSON)
+    public void deleteFeatureFlag(@PathParam("flag") String key) throws AtlasBaseException {
+        AtlasAuthorizationUtils.verifyAccess(new AtlasAdminAccessRequest(AtlasPrivilege.ADMIN_FEATURE_FLAG_CUD), "featureFlag");
+        FeatureFlagStore.deleteFlag(key);
+    }
     private String getEditableEntityTypes(Configuration config) {
         String ret = DEFAULT_EDITABLE_ENTITY_TYPES;