add indexes

akto-api-security · Nov 18, 2023 · 69cacfa · 69cacfa
1 parent a036725
commit 69cacfa
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 52 deletions.
diff --git a/libs/dao/src/main/java/com/akto/dao/ApiInfoDao.java b/libs/dao/src/main/java/com/akto/dao/ApiInfoDao.java
@@ -3,10 +3,11 @@
 import com.akto.dao.context.Context;
 import com.akto.dto.ApiInfo;
 import com.akto.dto.type.SingleTypeInfo;
-import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.MongoCursor;
 import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.Indexes;
 
+import org.bson.Document;
 import org.bson.conversions.Bson;
 
 import java.util.ArrayList;
@@ -19,17 +20,46 @@ public class ApiInfoDao extends AccountsContextDao<ApiInfo>{
 
     public void createIndicesIfAbsent() {
 
-        String dbName = Context.accountId.get()+"";
-        createCollectionIfAbsent(dbName, getCollName(), new CreateCollectionOptions());
+        boolean exists = false;
+        for (String col: clients[0].getDatabase(Context.accountId.get()+"").listCollectionNames()){
+            if (getCollName().equalsIgnoreCase(col)){
+                exists = true;
+                break;
+            }
+        };
 
-        List<Bson> indices = new ArrayList<>(Arrays.asList(
-                Indexes.ascending(new String[] { ApiInfo.ID_API_COLLECTION_ID }),
-                Indexes.ascending(new String[] { ApiInfo.ID_API_COLLECTION_ID, ApiInfo.ID_URL }),
-                Indexes.ascending(new String[] { ApiInfo.ID_URL }),
-                Indexes.ascending(new String[] { SingleTypeInfo._COLLECTION_IDS, ApiInfo.ID_URL })
-        ));
+        if (!exists) {
+            clients[0].getDatabase(Context.accountId.get()+"").createCollection(getCollName());
+        }
+
+        MongoCursor<Document> cursor = instance.getMCollection().listIndexes().cursor();
+        int counter = 0;
+        while (cursor.hasNext()) {
+            counter++;
+            cursor.next();
+        }
+
+        if (counter == 1) {
+            String[] fieldNames = {"_id.apiCollectionId"};
+            ApiInfoDao.instance.getMCollection().createIndex(Indexes.ascending(fieldNames));    
+            counter++;
+        }
+
+        if (counter == 2) {
+            String[] fieldNames = {"_id.url"};
+            ApiInfoDao.instance.getMCollection().createIndex(Indexes.ascending(fieldNames));    
+            counter++;
+        }
+
+        if (counter == 3) {
+            String[] fieldNames = {"_id.apiCollectionId", "_id.url"};
+            ApiInfoDao.instance.getMCollection().createIndex(Indexes.ascending(fieldNames));    
+            counter++;
+        }
+
+        MCollection.createIndexIfAbsent(getDBName(), getCollName(),
+                new String[] { SingleTypeInfo._COLLECTION_IDS, ApiInfo.ID_URL }, true);
 
-        createIndices(indices);
     }
 
     @Override

diff --git a/libs/dao/src/main/java/com/akto/dao/MCollection.java b/libs/dao/src/main/java/com/akto/dao/MCollection.java
@@ -191,39 +191,6 @@ public Logger getLogger() {
         return logger;
     }
 
-    public void dropIndexes(List<Bson> indexes){
-        try {
-            for(Bson index: indexes) {
-                this.getMCollection().dropIndex(index);
-            }
-        } catch (Exception e) {
-            // error outs on non-existent indices
-            e.printStackTrace();
-        }
-    }
-
-    public String getIndexName(Bson idx){
-        String ret = "";
-        BsonDocument id = idx.toBsonDocument();
-        int c = 0;
-        for(String key : id.keySet()){
-            if(c > 0){
-                ret += "_";
-            }
-            int value = id.getInt32(key).getValue();
-            ret += key + "_" + value;
-            c++;
-        }
-        return ret;
-    }
-
-    public void createIndices(List<Bson> indices){
-        for (Bson index : indices) {
-            String name = this.getIndexName(index);
-            createIndexIfAbsent(this.getDBName(), this.getCollName(), index, new IndexOptions().name(name));
-        }
-    }
-
     public static boolean createCollectionIfAbsent(String dbName, String collName, CreateCollectionOptions options){
         try{
             boolean exists = false;

diff --git a/libs/dao/src/main/java/com/akto/dao/SampleDataDao.java b/libs/dao/src/main/java/com/akto/dao/SampleDataDao.java
@@ -5,17 +5,14 @@
 import com.akto.dto.traffic.SampleData;
 import com.akto.dto.type.SingleTypeInfo;
 import com.mongodb.client.MongoCursor;
-import com.mongodb.client.model.CreateCollectionOptions;
 import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.Indexes;
 import com.mongodb.client.model.Projections;
 import com.mongodb.client.model.Sorts;
-
 import org.bson.Document;
 import org.bson.conversions.Bson;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 
 public class SampleDataDao extends AccountsContextDao<SampleData> {
@@ -63,6 +60,9 @@ public void createIndicesIfAbsent() {
             instance.getMCollection().createIndex(Indexes.ascending("_id.apiCollectionId"));
         }
 
+        MCollection.createIndexIfAbsent(getDBName(), getCollName(),
+                new String[] { SingleTypeInfo._COLLECTION_IDS, ApiInfo.ID_URL, ApiInfo.ID_METHOD }, true);
+
     }
 
     public List<SampleData> fetchSampleDataPaginated(int apiCollectionId, String lastFetchedUrl,

diff --git a/libs/dao/src/main/java/com/akto/dao/SensitiveSampleDataDao.java b/libs/dao/src/main/java/com/akto/dao/SensitiveSampleDataDao.java
@@ -5,15 +5,10 @@
 import com.akto.dto.SensitiveSampleData;
 import com.akto.dto.type.SingleTypeInfo;
 import com.mongodb.client.MongoCursor;
-import com.mongodb.client.model.CreateCollectionOptions;
 import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.Indexes;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-
 import org.bson.Document;
+
 import org.bson.conversions.Bson;
 
 public class SensitiveSampleDataDao extends AccountsContextDao<SensitiveSampleData>{
@@ -65,5 +60,12 @@ public void createIndicesIfAbsent() {
             String[] fieldNames = {"_id.url", "_id.apiCollectionId", "_id.method"};
             MCollection.createIndexIfAbsent(getDBName(), getCollName(), fieldNames, true);
         }
+
+        MCollection.createIndexIfAbsent(getDBName(), getCollName(),
+                new String[] { ApiInfo.ID_URL, SingleTypeInfo._COLLECTION_IDS, ApiInfo.ID_METHOD }, true);
+
+        MCollection.createIndexIfAbsent(getDBName(), getCollName(),
+                new String[] { SingleTypeInfo._COLLECTION_IDS }, true);
+
     }
 }
diff --git a/libs/dao/src/main/java/com/akto/dao/SingleTypeInfoDao.java b/libs/dao/src/main/java/com/akto/dao/SingleTypeInfoDao.java
@@ -86,6 +86,10 @@ public void createIndicesIfAbsent() {
             String[] fieldNames = {"responseCode", "subType", "timestamp",};
             SingleTypeInfoDao.instance.getMCollection().createIndex(Indexes.ascending(fieldNames));    
         }
+
+        MCollection.createIndexIfAbsent(getDBName(), getCollName(),
+            new String[] { SingleTypeInfo._COLLECTION_IDS }, true);
+
     }
 
     public static Bson filterForHostHeader(int apiCollectionId, boolean useApiCollectionId) {