fix: CitusDB for analytics tables Merge with 2.42/master (#14945) [DHIS2-15548]

dhis-2/dhis-api/src/main/java/org/hisp/dhis/analytics/AnalyticsTableType.java

@@ -28,6 +28,7 @@
 package org.hisp.dhis.analytics;
 
 import lombok.Getter;
+import org.apache.commons.lang3.StringUtils;
 
 /**
  * @author Lars Helge Overland
@@ -38,15 +39,18 @@ public enum AnalyticsTableType {
   COMPLETENESS("analytics_completeness", true, true),
   COMPLETENESS_TARGET("analytics_completenesstarget", false, false),
   ORG_UNIT_TARGET("analytics_orgunittarget", false, false),
-  EVENT("analytics_event", false, true),
-  ENROLLMENT("analytics_enrollment", false, false),
+  EVENT("analytics_event", false, true, "psi"),
+  ENROLLMENT("analytics_enrollment", false, false, "pi"),
   OWNERSHIP("analytics_ownership", false, false),
   VALIDATION_RESULT("analytics_validationresult", true, false),
-  TRACKED_ENTITY_INSTANCE_EVENTS("analytics_tei_events", false, true),
-  TRACKED_ENTITY_INSTANCE_ENROLLMENTS("analytics_tei_enrollments", false, false),
-  TRACKED_ENTITY_INSTANCE("analytics_tei", false, false);
+  TRACKED_ENTITY_INSTANCE_EVENTS("analytics_tei_events", false, true, "trackedentityinstanceuid"),
+  TRACKED_ENTITY_INSTANCE_ENROLLMENTS(
+      "analytics_tei_enrollments", false, false, "trackedentityinstanceuid"),
+  TRACKED_ENTITY_INSTANCE("analytics_tei", false, false, "trackedentityinstanceuid");
 
-  private final String tableName;
+  @Getter private String tableName;
+
+  @Getter private String distributionColumn;
 
   private final boolean periodDimension;
 
@@ -57,4 +61,17 @@ public enum AnalyticsTableType {
     this.periodDimension = periodDimension;
     this.latestPartition = latestPartition;
   }
+
+  AnalyticsTableType(
+      String tableName,
+      boolean periodDimension,
+      boolean latestPartition,
+      String distributionColumn) {
+    this(tableName, periodDimension, latestPartition);
+    this.distributionColumn = distributionColumn;
+  }
+
+  public boolean isDistributed() {
+    return StringUtils.isNotBlank(distributionColumn);
+  }
 }

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/AnalyticsTableUpdateParams.java

@@ -81,6 +81,9 @@ public class AnalyticsTableUpdateParams {
   /** Current date, only used for testing */
   private Date today;
 
+  /** Is Citus enabled? */
+  @Getter private boolean citusExtensionEnabled = false;
+
   private final Map<String, Object> extraParameters = new HashMap<>();
 
   public void addExtraParam(String prefix, String key, Object value) {
@@ -171,6 +174,7 @@ public AnalyticsTableUpdateParams instance() {
     params.jobId = this.jobId;
     params.startTime = this.startTime;
     params.lastSuccessfulUpdate = this.lastSuccessfulUpdate;
+    params.citusExtensionEnabled = this.citusExtensionEnabled;
 
     return this;
   }
@@ -255,5 +259,10 @@ public AnalyticsTableUpdateParams build() {
       checkNotNull(this.params.startTime);
       return this.params;
     }
+
+    public Builder withCitusEnabled(boolean citusExtensionEnabled) {
+      this.params.citusExtensionEnabled = citusExtensionEnabled;
+      return this;
+    }
   }
 }

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/AbstractJdbcTableManager.java

@@ -181,6 +181,42 @@ public void removeUpdatedData(List<AnalyticsTable> tables) {}
   public void createTable(AnalyticsTable table) {
     createAnalyticsTable(table);
     createAnalyticsTablePartitions(table);
+
+    if (analyticsTableSettings.isCitusExtensionEnabled()) {
+      createDistributedCitusTable(table);
+    }
+  }
+
+  /**
+   * Create a distributed Citus table.
+   *
+   * @param table the {@link AnalyticsTable}.
+   */
+  private void createDistributedCitusTable(AnalyticsTable table) {
+    if (!table.isTableTypeDistributed()) {
+      log.warn(
+          "No distribution column defined for table "
+              + table.getMainName()
+              + " so it won't be distributed");
+      return;
+    }
+
+    String tableName = table.getName();
+    String distributionColumn = table.getTableType().getDistributionColumn();
+
+    try {
+      jdbcTemplate.query(
+          "select create_distributed_table( :1, :2 )",
+          ps -> {
+            ps.setString(1, tableName);
+            ps.setString(2, distributionColumn);
+          },
+          rs -> {});
+      log.info("Successfully distributed table " + tableName + " on column " + distributionColumn);
+    } catch (Exception e) {
+      log.warn(
+          "Failed to distribute table " + table.getName() + " on column " + distributionColumn, e);
+    }
   }
 
   /**

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/DefaultAnalyticsTableGenerator.java

@@ -45,6 +45,7 @@
 import org.hisp.dhis.analytics.AnalyticsTableUpdateParams;
 import org.hisp.dhis.analytics.cache.AnalyticsCache;
 import org.hisp.dhis.analytics.cache.OutliersCache;
+import org.hisp.dhis.analytics.table.setting.AnalyticsTableSettings;
 import org.hisp.dhis.resourcetable.ResourceTableService;
 import org.hisp.dhis.scheduling.JobProgress;
 import org.hisp.dhis.setting.SettingKey;
@@ -69,6 +70,8 @@ public class DefaultAnalyticsTableGenerator implements AnalyticsTableGenerator {
 
   private final OutliersCache outliersCache;
 
+  private final AnalyticsTableSettings analyticsTableSettings;
+
   // TODO introduce last successful timestamps per table type
 
   @Override
@@ -85,6 +88,7 @@ public void generateAnalyticsTables(AnalyticsTableUpdateParams params0, JobProgr
     AnalyticsTableUpdateParams params =
         AnalyticsTableUpdateParams.newBuilder(params0)
             .withLastSuccessfulUpdate(lastSuccessfulUpdate)
+            .withCitusEnabled(analyticsTableSettings.isCitusExtensionEnabled())
             .build();
 
     log.info("Found {} analytics table types: {}", availableTypes.size(), availableTypes);

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/DefaultAnalyticsTableService.java

@@ -105,7 +105,7 @@ public void create(AnalyticsTableUpdateParams params, JobProgress progress) {
       return;
     }
 
-    List<AnalyticsTable> tables = tableManager.getAnalyticsTables(params);
+    List<AnalyticsTable> tables = getAnalyticsTables(params);
 
     if (tables.isEmpty()) {
       clock.logTime(
@@ -172,6 +172,19 @@ public void create(AnalyticsTableUpdateParams params, JobProgress progress) {
     clock.logTime("Table update done: '{}'", tableType.getTableName());
   }
 
+  private List<AnalyticsTable> getAnalyticsTables(AnalyticsTableUpdateParams params) {
+    return tableManager.getAnalyticsTables(params).stream()
+        .map(analyticsTable -> withDistributedFlag(analyticsTable, params))
+        .toList();
+  }
+
+  private AnalyticsTable withDistributedFlag(
+      AnalyticsTable analyticsTable, AnalyticsTableUpdateParams params) {
+    analyticsTable.setTableDistributed(
+        params.isCitusExtensionEnabled() && analyticsTable.isTableTypeDistributed());
+    return analyticsTable;
+  }
+
   @Override
   public void dropTables() {
     Set<String> tables = tableManager.getExistingDatabaseTables();

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/JdbcTeiEventsAnalyticsTableManager.java

@@ -172,9 +172,15 @@ public AnalyticsTableType getAnalyticsTableType() {
   @Override
   @Transactional
   public List<AnalyticsTable> getAnalyticsTables(AnalyticsTableUpdateParams params) {
+    List<AnalyticsTable> tables = new ArrayList<>();
+
+    // Skipping if the latest partition is requested since it's not supported.
+    if (params.isLatestUpdate()) {
+      return tables;
+    }
+
     Calendar calendar = PeriodType.getCalendar();
     List<TrackedEntityType> trackedEntityTypes = trackedEntityTypeService.getAllTrackedEntityType();
-    List<AnalyticsTable> tables = new ArrayList<>();
     Logged logged = analyticsTableSettings.getTableLogged();
 
     for (TrackedEntityType tet : trackedEntityTypes) {
@@ -189,7 +195,7 @@ public List<AnalyticsTable> getAnalyticsTables(AnalyticsTableUpdateParams params
             List.of(), year, getStartDate(calendar, year), getEndDate(calendar, year));
       }
 
-      if (table.hasTablePartitions()) {
+      if (table.hasTablePartitions() || params.isCitusExtensionEnabled()) {
         tables.add(table);
       }
     }
@@ -266,6 +272,10 @@ public void populateTable(AnalyticsTableUpdateParams params, AnalyticsTableParti
     List<AnalyticsTableColumn> columns = partition.getMasterTable().getAnalyticsTableColumns();
     String partitionClause = getPartitionClause(partition);
 
+    if (params.isCitusExtensionEnabled()) {
+      partitionClause = eventDateExpression + " is not null";
+    }
+
     StringBuilder sql = new StringBuilder("insert into " + tableName + " (");
 
     for (AnalyticsTableColumn col : columns) {

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/model/AnalyticsTable.java

@@ -31,6 +31,7 @@
 import java.util.List;
 import java.util.Objects;
 import lombok.Getter;
+import lombok.Setter;
 import org.apache.commons.lang3.StringUtils;
 import org.hisp.dhis.analytics.AnalyticsTableType;
 import org.hisp.dhis.commons.collection.UniqueArrayList;
@@ -62,6 +63,8 @@ public class AnalyticsTable extends Table {
   /** Tracked entity type of enrollments in analytics table. */
   private TrackedEntityType trackedEntityType;
 
+  @Getter @Setter private boolean tableDistributed = false;
+
   /** Analytics table partitions for this base analytics table. */
   private List<AnalyticsTablePartition> tablePartitions = new UniqueArrayList<>();
 
@@ -241,6 +244,10 @@ public boolean hasTablePartitions() {
     return !tablePartitions.isEmpty();
   }
 
+  public boolean isTableTypeDistributed() {
+    return tableType.isDistributed();
+  }
+
   /**
    * Returns the latest partition, or null if no latest partition exists.
    *

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/table/setting/AnalyticsTableSettings.java

@@ -36,8 +36,10 @@
 import static org.hisp.dhis.util.ObjectUtils.isNull;
 
 import lombok.RequiredArgsConstructor;
+import lombok.experimental.Delegate;
 import org.apache.commons.lang3.EnumUtils;
 import org.apache.commons.lang3.StringUtils;
+import org.hisp.dhis.configuration.CitusSettings;
 import org.hisp.dhis.db.model.Database;
 import org.hisp.dhis.db.model.Logged;
 import org.hisp.dhis.external.conf.DhisConfigurationProvider;
@@ -57,6 +59,8 @@ public class AnalyticsTableSettings {
 
   private final SystemSettingManager systemSettings;
 
+  @Delegate private final CitusSettings citusSettings;
+
   /**
    * Returns the setting indicating whether resource and analytics tables should be logged or
    * unlogged.

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/db/sql/PostgreSqlBuilder.java

@@ -29,7 +29,9 @@
 
 import static org.hisp.dhis.commons.util.TextUtils.removeLastComma;
 
+import java.util.List;
 import java.util.stream.Collectors;
+import org.hisp.dhis.analytics.table.model.AnalyticsTable;
 import org.hisp.dhis.db.model.Collation;
 import org.hisp.dhis.db.model.Column;
 import org.hisp.dhis.db.model.Index;
@@ -239,52 +241,91 @@ public String createTable(Table table) {
             .append(" (");
 
     // Columns
-
     if (table.hasColumns()) {
-      for (Column column : table.getColumns()) {
-        String dataType = getDataTypeName(column.getDataType());
-        String nullable = column.getNullable() == Nullable.NOT_NULL ? " not null" : " null";
-        String collation = column.getCollation() == Collation.C ? (" collate " + quote("C")) : "";
-
-        sql.append(quote(column.getName()) + " ")
-            .append(dataType)
-            .append(nullable)
-            .append(collation)
-            .append(COMMA);
-      }
+      sql.append(columnsStatement(table.getColumns()));
     }
 
     // Primary key
-
     if (table.hasPrimaryKey()) {
-      sql.append("primary key (");
-
-      for (String columnName : table.getPrimaryKey()) {
-        sql.append(quote(columnName)).append(COMMA);
-      }
-
-      removeLastComma(sql).append(")").append(COMMA);
+      sql.append(primaryKeysStatement(table.getPrimaryKey()));
     }
 
     // Checks
-
     if (table.hasChecks()) {
-      for (String check : table.getChecks()) {
-        sql.append("check(" + check + ")").append(COMMA);
-      }
+      sql.append(checksStatement(table.getChecks()));
     }
 
     removeLastComma(sql).append(")");
 
     // Parent
-
-    if (table.hasParent()) {
+    // Only use partitioned (inherited) tables when we should not distribute the table.
+    if (table instanceof AnalyticsTable analyticsTable
+        && !analyticsTable.isTableDistributed()
+        && table.hasParent()) {
       sql.append(" inherits (").append(quote(table.getParent().getName())).append(")");
     }
 
     return sql.append(";").toString();
   }
 
+  /**
+   * Returns the SQL statement for the given list of checks.
+   *
+   * @param checks the list of checks to append to the SQL statement.
+   * @return the SQL statement.
+   */
+  private static String checksStatement(List<String> checks) {
+    StringBuilder sql = new StringBuilder();
+
+    for (String check : checks) {
+      sql.append("check(" + check + ")").append(COMMA);
+    }
+
+    return sql.toString();
+  }
+
+  /**
+   * Returns the SQL statement for the given list of primaryKeys.
+   *
+   * @param primaryKeys the list of primaryKeys to append to the SQL statement.
+   * @return the SQL statement.
+   */
+  private String primaryKeysStatement(List<String> primaryKeys) {
+    StringBuilder sql = new StringBuilder("primary key (");
+
+    for (String columnName : primaryKeys) {
+      sql.append(quote(columnName)).append(COMMA);
+    }
+
+    removeLastComma(sql).append(")").append(COMMA);
+
+    return sql.toString();
+  }
+
+  /**
+   * Returns the SQL statement for the given list of columns.
+   *
+   * @param columns the list of columns to append to the SQL statement.
+   * @return the SQL statement.
+   */
+  private String columnsStatement(List<Column> columns) {
+    StringBuilder sql = new StringBuilder();
+
+    for (Column column : columns) {
+      String dataType = getDataTypeName(column.getDataType());
+      String nullable = column.getNullable() == Nullable.NOT_NULL ? " not null" : " null";
+      String collation = column.getCollation() == Collation.C ? (" collate " + quote("C")) : "";
+
+      sql.append(quote(column.getName()) + " ")
+          .append(dataType)
+          .append(nullable)
+          .append(collation)
+          .append(COMMA);
+    }
+
+    return sql.toString();
+  }
+
   @Override
   public String analyzeTable(String name) {
     return String.format("analyze %s;", quote(name));