LPD-27982 When db partitioning is enabled, split the batch into each …

…company.
brianikim · Jun 8, 2024 · e90e7f3 · e90e7f3
1 parent a414aeb
commit e90e7f3
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/modules/apps/portal-security-audit/portal-security-audit-storage-service/build.gradle b/modules/apps/portal-security-audit/portal-security-audit-storage-service/build.gradle
@@ -12,6 +12,7 @@ dependencies {
 	compileOnly project(":apps:portal-security-audit:portal-security-audit-storage-api")
 	compileOnly project(":apps:portal:portal-aop-api")
 	compileOnly project(":apps:static:portal:portal-upgrade-api")
+	compileOnly project(":core:petra:petra-function")
 	compileOnly project(":core:petra:petra-lang")
 	compileOnly project(":core:petra:petra-sql-dsl-api")
 	compileOnly project(":core:petra:petra-string")

diff --git a/...c/main/java/com/liferay/portal/security/audit/storage/internal/AuditEventManagerImpl.java b/...c/main/java/com/liferay/portal/security/audit/storage/internal/AuditEventManagerImpl.java
@@ -6,6 +6,8 @@
 package com.liferay.portal.security.audit.storage.internal;
 
 import com.liferay.portal.kernel.audit.AuditMessage;
+import com.liferay.portal.kernel.db.partition.DBPartition;
+import com.liferay.portal.kernel.service.CompanyLocalService;
 import com.liferay.portal.kernel.util.OrderByComparator;
 import com.liferay.portal.security.audit.AuditEvent;
 import com.liferay.portal.security.audit.AuditEventManager;
@@ -14,7 +16,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.osgi.service.component.annotations.Component;
 import org.osgi.service.component.annotations.Reference;
@@ -33,7 +37,29 @@ public AuditEvent addAuditEvent(AuditMessage auditMessage) {
 
 	@Override
 	public void addAuditEvents(List<AuditMessage> auditMessages) {
-		_auditEventLocalService.addAuditEvents(auditMessages);
+		if (DBPartition.isPartitionEnabled()) {
+			Map<Long, List<AuditMessage>> auditMessagesMap = new HashMap<>();
+
+			for (AuditMessage auditMessage : auditMessages) {
+				List<AuditMessage> companyAuditMessages =
+					auditMessagesMap.computeIfAbsent(
+						auditMessage.getCompanyId(), key -> new ArrayList<>());
+
+				companyAuditMessages.add(auditMessage);
+			}
+
+			for (Map.Entry<Long, List<AuditMessage>> entry :
+					auditMessagesMap.entrySet()) {
+
+				_companyLocalService.forEachCompanyId(
+					companyId -> _auditEventLocalService.addAuditEvents(
+						entry.getValue()),
+					new long[] {entry.getKey()});
+			}
+		}
+		else {
+			_auditEventLocalService.addAuditEvents(auditMessages);
+		}
 	}
 
 	@Override
@@ -121,4 +147,7 @@ private List<AuditEvent> _translate(
 	@Reference
 	private AuditEventLocalService _auditEventLocalService;
 
+	@Reference
+	private CompanyLocalService _companyLocalService;
+
 }