Consolidate Dataclass data update methods - use DIB for update only

api/src/org/labkey/api/audit/AuditHandler.java

@@ -6,6 +6,7 @@
 import org.labkey.api.data.ColumnInfo;
 import org.labkey.api.data.Container;
 import org.labkey.api.data.MultiValuedForeignKey;
+import org.labkey.api.data.MultiValuedRenderContext;
 import org.labkey.api.data.TableInfo;
 import org.labkey.api.dataiterator.DataIterator;
 import org.labkey.api.dataiterator.ExistingRecordDataIterator;
@@ -122,11 +123,22 @@ static Pair<Map<String, Object>, Map<String, Object>> getOldAndNewRecordForMerge
                 }
             }
 
+            boolean isAliasInput = row.containsKey(ExperimentService.ALIASCOLUMNALIAS) && "Alias".equalsIgnoreCase(lcName);
+
             boolean isExtraAuditField = extraFieldsToInclude != null && extraFieldsToInclude.contains(nameFromAlias);
-            if (!excludedFromDetailDiff.contains(nameFromAlias) && (row.containsKey(nameFromAlias) || isExpInput))
+            if (!excludedFromDetailDiff.contains(nameFromAlias) && (row.containsKey(nameFromAlias) || isExpInput || isAliasInput))
             {
                 Object oldValue = entry.getValue();
                 Object newValue = row.get(nameFromAlias);
+
+                // See ExpDataIterator: step1.addColumn(ExperimentService.ALIASCOLUMNALIAS, colNameMap.get(Alias.name()))
+                if (isAliasInput && newValue == null)
+                {
+                    newValue = row.get(ExperimentService.ALIASCOLUMNALIAS);
+                    if (oldValue instanceof String aliasStr)
+                        oldValue = Arrays.asList(aliasStr.split(MultiValuedRenderContext.VALUE_DELIMITER_REGEX));
+                }
+
                 // compare dates using string values to allow for both Date and Timestamp types
                 if (newValue instanceof Date && oldValue != null)
                 {
@@ -167,7 +179,7 @@ else if (!Objects.equals(oldValue, newValue) || isExtraAuditField)
                     // If multivalued columns change, the value in this table will remain the key to the junction table
                     // but at this point newValue will look like the newly chosen values not that key. So we skip
                     // this in the diff unless the value changes from non-null to null or vice versa.
-                    if (isMultiValued)
+                    if (isMultiValued && !isAliasInput)
                     {
                         if ((oldValue == null && newValue != null) || (newValue == null && oldValue != null))
                         {

api/src/org/labkey/api/data/AbstractTableInfo.java

@@ -2316,6 +2316,31 @@ public boolean allowRobotsIndex()
         return getUserSchema().allowRobotsIndex();
     }
 
+    @NotNull
+    private final Map<Enum, Object> _customTableConfigs = new HashMap<>();
+
+    public void putCustomTableConfig(Enum key, Object value)
+    {
+        _customTableConfigs.put(key, value);
+    }
+
+    @NotNull
+    public Map<Enum, Object> getCustomTableConfigs()
+    {
+        return _customTableConfigs;
+    }
+
+    @Nullable
+    public Object getCustomTableConfig(Enum key)
+    {
+        return getCustomTableConfigs().get(key);
+    }
+
+    public boolean getCustomTableConfigBoolean(Enum key)
+    {
+        return Boolean.TRUE == getCustomTableConfig(key);
+    }
+
     public static class TestCase extends Assert{
         @Test
         public void testEnum()

api/src/org/labkey/api/data/TableInfo.java

@@ -223,11 +223,6 @@ void addColumn(ColumnInfo column)
      */
     @NotNull List<ColumnInfo> getAlternateKeyColumns();
 
-    @NotNull default Set<String> getAltKeysForUpdate()
-    {
-       return Collections.emptySet();
-    }
-
     @Nullable default Set<String> getDisabledSystemFields()
     {
         return Collections.emptySet();

api/src/org/labkey/api/dataiterator/AbstractPrefetchingDataIterator.java

@@ -0,0 +1,132 @@
+package org.labkey.api.dataiterator;
+
+import org.apache.commons.lang3.StringUtils;
+import org.labkey.api.data.ColumnInfo;
+import org.labkey.api.data.JdbcType;
+import org.labkey.api.data.TableInfo;
+import org.labkey.api.query.BatchValidationException;
+
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import static org.labkey.api.util.IntegerUtils.asInteger;
+
+/**
+ * Abstract base for data iterators that prefetch data in batches using a single primary key column.
+ * Provides {@link #buildBatch()} to collect up to 50 rows and their keys into lookup maps,
+ * and a standard {@link #next()} implementation that marks, advances, prefetches, and checks errors.
+ * <p>
+ * Used by {@link SampleUpdateAddColumnsDataIterator} and {@link DataClassUpdateAddColumnsDataIterator}.
+ */
+public abstract class AbstractPrefetchingDataIterator extends WrapperDataIterator
+{
+    final DataIteratorContext _context;
+
+    // NOTE: in may be wrapped with a LoggingDataIterator by WrapperDataIterator; _unwrapped retains the original
+    final CachingDataIterator _unwrapped;
+
+    final TableInfo target;
+    final ColumnInfo pkColumn;
+    final Supplier<Object> pkSupplier;
+
+    int lastPrefetchRowNumber = -1;
+
+    /**
+     * @param in            the DataIterator to wrap; must be a {@link CachingDataIterator}
+     * @param context       the current data iterator context
+     * @param target        the target table; used to resolve the key column
+     * @param keyColumnName name of the primary key column present in {@code in}
+     */
+    AbstractPrefetchingDataIterator(CachingDataIterator in, DataIteratorContext context, TableInfo target, String keyColumnName)
+    {
+        super(in);
+        this._unwrapped = in;
+        this._context = context;
+        this.target = target;
+
+        var map = DataIteratorUtil.createColumnNameMap(in);
+        Integer index = map.get(keyColumnName);
+        ColumnInfo col = target.getColumn(keyColumnName);
+        if (null == index || null == col)
+            throw new IllegalArgumentException("Key column not found: " + keyColumnName);
+        pkSupplier = in.getSupplier(index);
+        pkColumn = col;
+    }
+
+    abstract void prefetchExisting() throws BatchValidationException;
+
+    /**
+     * Holds the result of collecting a batch of rows for prefetching.
+     * <ul>
+     *   <li>{@code rowKeyMap} — maps each row number to its key value</li>
+     *   <li>{@code keyRowMap} — maps each key value to the list of row numbers that carry it
+     *       (normally one row per key, but supports duplicates gracefully)</li>
+     * </ul>
+     */
+    record BatchResult(Map<Integer, Object> rowKeyMap, Map<Object, List<Integer>> keyRowMap) {}
+
+    /**
+     * Collects up to 50 rows from the current position into key-lookup maps.
+     * Validates that each key value is non-null (numeric keys) or non-blank (string keys).
+     * Updates {@link #lastPrefetchRowNumber}.
+     * Call {@link #resetAfterBatch()} after processing to rewind the iterator.
+     */
+    BatchResult buildBatch() throws BatchValidationException
+    {
+        int rowsToFetch = 50;
+        String keyFieldName = pkColumn.getName();
+        boolean numericKey = pkColumn.isNumericType();
+        JdbcType jdbcType = pkColumn.getJdbcType();
+        Map<Integer, Object> rowKeyMap = new LinkedHashMap<>();
+        Map<Object, List<Integer>> keyRowMap = new LinkedHashMap<>();
+        do
+        {
+            lastPrefetchRowNumber = asInteger(_delegate.get(0));
+            Object keyObj = pkSupplier.get();
+            Object key = jdbcType.convert(keyObj);
+
+            if (numericKey)
+            {
+                if (null == key)
+                    throw new IllegalArgumentException(keyFieldName + " value not provided on row " + lastPrefetchRowNumber);
+            }
+            else if (StringUtils.isEmpty((String) key))
+                throw new IllegalArgumentException(keyFieldName + " value not provided on row " + lastPrefetchRowNumber);
+
+            rowKeyMap.put(lastPrefetchRowNumber, key);
+            keyRowMap.computeIfAbsent(key, ignored -> new ArrayList<>()).add(lastPrefetchRowNumber);
+        }
+        while (--rowsToFetch > 0 && _delegate.next());
+        return new BatchResult(rowKeyMap, keyRowMap);
+    }
+
+    /**
+     * Rewinds the iterator to the start of the current batch so the caller can process rows one at a time.
+     */
+    void resetAfterBatch() throws BatchValidationException
+    {
+        _unwrapped.reset(); // unwrapped _delegate
+        _delegate.next();
+    }
+
+    @Override
+    public boolean next() throws BatchValidationException
+    {
+        if (_context.getErrors().hasErrors())
+            return false;
+
+        // NOTE: we have to call mark() before we call next() if we want the 'next' row to be cached
+        _unwrapped.mark();  // unwrapped _delegate
+        boolean ret = super.next();
+        if (ret && !_context.getErrors().hasErrors())
+        {
+            prefetchExisting();
+            if (_context.getErrors().hasErrors())
+                return false;
+        }
+        return ret;
+    }
+}

api/src/org/labkey/api/dataiterator/DataClassUpdateAddColumnsDataIterator.java

@@ -0,0 +1,132 @@
+package org.labkey.api.dataiterator;
+
+import org.jetbrains.annotations.NotNull;
+import org.labkey.api.collections.IntHashMap;
+import org.labkey.api.collections.Sets;
+import org.labkey.api.data.CompareType;
+import org.labkey.api.data.Container;
+import org.labkey.api.data.SimpleFilter;
+import org.labkey.api.data.TableInfo;
+import org.labkey.api.data.TableSelector;
+import org.labkey.api.exp.api.ExperimentService;
+import org.labkey.api.exp.query.ExpDataTable;
+import org.labkey.api.query.BatchValidationException;
+import org.labkey.api.query.FieldKey;
+import org.labkey.api.query.ValidationException;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+
+import static org.labkey.api.exp.query.ExpDataTable.Column.LSID;
+import static org.labkey.api.exp.query.ExpDataTable.Column.ClassId;
+import static org.labkey.api.util.IntegerUtils.asInteger;
+
+/**
+ * DataIterator that adds the LSID column for DataClass update operations.
+ * Queries the LSID from exp.data based on the provided key (rowId or name) and dataClassId.
+ * The LSID is needed downstream for attachment handling.
+ */
+public class DataClassUpdateAddColumnsDataIterator extends AbstractPrefetchingDataIterator
+{
+    private final Container _targetContainer;
+    private final long _dataClassId;
+    final int _lsidColIndex;
+
+    // prefetch of existing records
+    final IntHashMap<String> lsids = new IntHashMap<>();
+
+    public DataClassUpdateAddColumnsDataIterator(CachingDataIterator in, @NotNull DataIteratorContext context, TableInfo target, Container container, long dataClassId, String keyColumnName)
+    {
+        super(in, context, target, keyColumnName);
+        _targetContainer = container;
+        _dataClassId = dataClassId;
+
+        var map = DataIteratorUtil.createColumnNameMap(in);
+        Integer lsidIdx = map.get(ExpDataTable.Column.LSID.name());
+        if (lsidIdx == null)
+            throw new IllegalStateException("LSID column not found in input.");
+        this._lsidColIndex = lsidIdx;
+    }
+
+    @Override
+    public Supplier<Object> getSupplier(int i)
+    {
+        if (i != _lsidColIndex)
+            return _delegate.getSupplier(i);
+        return () -> get(i);
+    }
+
+    @Override
+    public Object get(int i)
+    {
+        Integer rowNumber = asInteger(_delegate.get(0));
+
+        if (i == _lsidColIndex)
+            return lsids.get(rowNumber);
+
+        return _delegate.get(i);
+    }
+
+    @Override
+    public boolean isConstant(int i)
+    {
+        if (i != _lsidColIndex)
+            return _delegate.isConstant(i);
+        return false;
+    }
+
+    @Override
+    public Object getConstantValue(int i)
+    {
+        if (i != _lsidColIndex)
+            return _delegate.getConstantValue(i);
+        return null;
+    }
+
+    @Override
+    protected void prefetchExisting() throws BatchValidationException
+    {
+        Integer rowNumber = asInteger(_delegate.get(0));
+        if (rowNumber <= lastPrefetchRowNumber)
+            return;
+
+        lsids.clear();
+
+        BatchResult batch = buildBatch();
+        Map<Integer, Object> rowKeyMap = batch.rowKeyMap();
+        Map<Object, List<Integer>> keyRowMap = batch.keyRowMap();
+        Set<Object> notFoundKeys = new HashSet<>(keyRowMap.keySet());
+
+        for (Integer rowInd : rowKeyMap.keySet())
+            lsids.put(rowInd, null);
+
+        String keyFieldName = pkColumn.getName();
+        SimpleFilter filter = new SimpleFilter(ClassId.fieldKey(), _dataClassId);
+        filter.addCondition(pkColumn.getFieldKey(), rowKeyMap.values(), CompareType.IN);
+        filter.addCondition(FieldKey.fromParts("Container"), _targetContainer);
+
+        Set<String> columns = Sets.newCaseInsensitiveHashSet(keyFieldName, LSID.name());
+        Map<String, Object>[] results = new TableSelector(ExperimentService.get().getTinfoData(), columns, filter, null).getMapArray();
+
+        for (Map<String, Object> result : results)
+        {
+            Object key = result.get(keyFieldName);
+            Object lsidObj = result.get(LSID.name());
+
+            if (lsidObj != null)
+            {
+                for (Integer rowInd : keyRowMap.get(key))
+                    lsids.put(rowInd, (String) lsidObj);
+                notFoundKeys.remove(key);
+            }
+        }
+
+        if (!notFoundKeys.isEmpty())
+            _context.getErrors().addRowError(new ValidationException("Data not found for " + notFoundKeys));
+
+        resetAfterBatch();
+    }
+}

api/src/org/labkey/api/dataiterator/ExistingRecordDataIterator.java

@@ -154,6 +154,8 @@ public Supplier<Object> getSupplier(int i)
     @Override
     public Object get(int i)
     {
+        assert(i <= existingColIndex) : "ExistingCol should be the last column.";
+
         if (i<existingColIndex)
             return _delegate.get(i);
         Integer rowNumber = asInteger(_delegate.get(0));