refactor(query-optimization): optimize entity_id condition for single…

… and multiple IDs
Permify · May 20, 2024 · 727fe46 · 727fe46
1 parent 9584af2
commit 727fe46
Showing 1 changed file with 78 additions and 22 deletions.
diff --git a/internal/storage/postgres/utils/filter.go b/internal/storage/postgres/utils/filter.go
@@ -14,8 +14,15 @@ func TuplesFilterQueryForSelectBuilder(sl squirrel.SelectBuilder, filter *base.T
 		eq["entity_type"] = filter.GetEntity().GetType()
 	}
 
-	if len(filter.GetEntity().GetIds()) > 0 {
-		eq["entity_id"] = filter.GetEntity().GetIds()
+	entityIds := filter.GetEntity().GetIds()
+	if len(entityIds) > 0 {
+		if len(entityIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["entity_id"] = entityIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["entity_id"] = entityIds
+		}
 	}
 
 	if filter.GetRelation() != "" {
@@ -26,16 +33,23 @@ func TuplesFilterQueryForSelectBuilder(sl squirrel.SelectBuilder, filter *base.T
 		eq["subject_type"] = filter.GetSubject().GetType()
 	}
 
-	if len(filter.GetSubject().GetIds()) > 0 {
-		eq["subject_id"] = filter.GetSubject().GetIds()
+	subjectIds := filter.GetSubject().GetIds()
+	if len(subjectIds) > 0 {
+		if len(subjectIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["subject_id"] = subjectIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["subject_id"] = subjectIds
+		}
 	}
 
 	if filter.GetSubject().GetRelation() != "" {
 		eq["subject_relation"] = filter.GetSubject().GetRelation()
 	}
 
-	// If eq is empty, return the original squirrel.UpdateBuilder without attaching a WHERE clause.
-	// This ensures we don't accidentally update every row in the table.
+	// If eq is empty, return the original squirrel.SelectBuilder without attaching a WHERE clause.
+	// This ensures we don't accidentally select every row in the table.
 	if len(eq) == 0 {
 		return sl
 	}
@@ -51,16 +65,30 @@ func AttributesFilterQueryForSelectBuilder(sl squirrel.SelectBuilder, filter *ba
 		eq["entity_type"] = filter.GetEntity().GetType()
 	}
 
-	if len(filter.GetEntity().GetIds()) > 0 {
-		eq["entity_id"] = filter.GetEntity().GetIds()
+	entityIds := filter.GetEntity().GetIds()
+	if len(entityIds) > 0 {
+		if len(entityIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["entity_id"] = entityIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["entity_id"] = entityIds
+		}
 	}
 
-	if len(filter.GetAttributes()) > 0 {
-		eq["attribute"] = filter.GetAttributes()
+	attributes := filter.GetAttributes()
+	if len(attributes) > 0 {
+		if len(attributes) == 1 {
+			// If there is only one attribute, use the = operator
+			eq["attribute"] = attributes[0]
+		} else {
+			// If there are multiple attributes, use the IN operator
+			eq["attribute"] = attributes
+		}
 	}
 
-	// If eq is empty, return the original squirrel.UpdateBuilder without attaching a WHERE clause.
-	// This ensures we don't accidentally update every row in the table.
+	// If eq is empty, return the original squirrel.SelectBuilder without attaching a WHERE clause.
+	// This ensures we don't accidentally select every row in the table.
 	if len(eq) == 0 {
 		return sl
 	}
@@ -76,8 +104,15 @@ func TuplesFilterQueryForUpdateBuilder(sl squirrel.UpdateBuilder, filter *base.T
 		eq["entity_type"] = filter.GetEntity().GetType()
 	}
 
-	if len(filter.GetEntity().GetIds()) > 0 {
-		eq["entity_id"] = filter.GetEntity().GetIds()
+	entityIds := filter.GetEntity().GetIds()
+	if len(entityIds) > 0 {
+		if len(entityIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["entity_id"] = entityIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["entity_id"] = entityIds
+		}
 	}
 
 	if filter.GetRelation() != "" {
@@ -88,8 +123,15 @@ func TuplesFilterQueryForUpdateBuilder(sl squirrel.UpdateBuilder, filter *base.T
 		eq["subject_type"] = filter.GetSubject().GetType()
 	}
 
-	if len(filter.GetSubject().GetIds()) > 0 {
-		eq["subject_id"] = filter.GetSubject().GetIds()
+	subjectIds := filter.GetSubject().GetIds()
+	if len(subjectIds) > 0 {
+		if len(subjectIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["subject_id"] = subjectIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["subject_id"] = subjectIds
+		}
 	}
 
 	if filter.GetSubject().GetRelation() != "" {
@@ -113,12 +155,26 @@ func AttributesFilterQueryForUpdateBuilder(sl squirrel.UpdateBuilder, filter *ba
 		eq["entity_type"] = filter.GetEntity().GetType()
 	}
 
-	if len(filter.GetEntity().GetIds()) > 0 {
-		eq["entity_id"] = filter.GetEntity().GetIds()
-	}
-
-	if len(filter.GetAttributes()) > 0 {
-		eq["attribute"] = filter.GetAttributes()
+	entityIds := filter.GetEntity().GetIds()
+	if len(entityIds) > 0 {
+		if len(entityIds) == 1 {
+			// If there is only one ID, use the = operator
+			eq["entity_id"] = entityIds[0]
+		} else {
+			// If there are multiple IDs, use the IN operator
+			eq["entity_id"] = entityIds
+		}
+	}
+
+	attributes := filter.GetAttributes()
+	if len(attributes) > 0 {
+		if len(attributes) == 1 {
+			// If there is only one attribute, use the = operator
+			eq["attribute"] = attributes[0]
+		} else {
+			// If there are multiple attributes, use the IN operator
+			eq["attribute"] = attributes
+		}
 	}
 
 	// If eq is empty, return the original squirrel.UpdateBuilder without attaching a WHERE clause.