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/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/DataClassUpdateAddColumnsDataIterator.java

@@ -0,0 +1,194 @@
+package org.labkey.api.dataiterator;
+
+import org.apache.commons.lang3.StringUtils;
+import org.jetbrains.annotations.NotNull;
+import org.labkey.api.collections.IntHashMap;
+import org.labkey.api.collections.Sets;
+import org.labkey.api.data.ColumnInfo;
+import org.labkey.api.data.CompareType;
+import org.labkey.api.data.Container;
+import org.labkey.api.data.JdbcType;
+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.LinkedHashMap;
+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 WrapperDataIterator
+{
+    private final Container _targetContainer;
+    final CachingDataIterator _unwrapped;
+
+    private final long _dataClassId;
+    final int _lsidColIndex;
+    final ColumnInfo pkColumn;
+    final Supplier<Object> pkSupplier;
+
+    int lastPrefetchRowNumber = -1;
+    final IntHashMap<String> lsids = new IntHashMap<>();
+    final DataIteratorContext _context;
+
+    public DataClassUpdateAddColumnsDataIterator(DataIterator in, @NotNull DataIteratorContext context, TableInfo target, Container container, long dataClassId, String keyColumnName)
+    {
+        super(in);
+        this._unwrapped = (CachingDataIterator)in;
+        _context = context;
+        _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;
+
+        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;
+    }
+
+    @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;
+    }
+
+    protected void prefetchExisting() throws BatchValidationException
+    {
+        Integer rowNumber = asInteger(_delegate.get(0));
+        if (rowNumber <= lastPrefetchRowNumber)
+            return;
+
+        lsids.clear();
+
+        int rowsToFetch = 50;
+        String keyFieldName = pkColumn.getName();
+        boolean numericKey = pkColumn.isNumericType();
+        JdbcType jdbcType = pkColumn.getJdbcType();
+        Map<Integer, Object> rowKeyMap = new LinkedHashMap<>();
+        Map<Object, Set<Integer>> keyRowMap = new LinkedHashMap<>();
+        Set<Object> notFoundKeys = new HashSet<>();
+
+        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);
+            notFoundKeys.add(key);
+            // if keyRowMap doesn't contain key, add new set, then add row number to set for this key
+            if (!keyRowMap.containsKey(key))
+                keyRowMap.put(key, new HashSet<>());
+            keyRowMap.get(key).add(lastPrefetchRowNumber);
+            lsids.put(lastPrefetchRowNumber, null);
+        }
+        while (--rowsToFetch > 0 && _delegate.next());
+
+        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());
+
+            Set<Integer> rowInds = keyRowMap.get(key);
+            if (lsidObj != null)
+            {
+                for (Integer rowInd : rowInds)
+                    lsids.put(rowInd, (String) lsidObj);
+                notFoundKeys.remove(key);
+            }
+        }
+
+        if (!notFoundKeys.isEmpty())
+            _context.getErrors().addRowError(new ValidationException("Data not found for " + notFoundKeys));
+
+        // backup to where we started so caller can iterate through them one at a time
+        _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 (!_context.getErrors().hasErrors() && ret)
+        {
+            prefetchExisting();
+            if (_context.getErrors().hasErrors())
+                return false;
+        }
+
+        return ret;
+    }
+}

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);
+
         if (i<existingColIndex)
             return _delegate.get(i);
         Integer rowNumber = asInteger(_delegate.get(0));

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

@@ -1,6 +1,7 @@
 package org.labkey.api.dataiterator;
 
 import org.apache.commons.lang3.StringUtils;
+import org.jetbrains.annotations.NotNull;
 import org.labkey.api.collections.IntHashMap;
 import org.labkey.api.collections.Sets;
 import org.labkey.api.data.ColumnInfo;
@@ -33,17 +34,19 @@ public class SampleUpdateAddColumnsDataIterator extends WrapperDataIterator
     final int _aliquotedFromColIndex;
     final int _rootMaterialRowIdColIndex;
     final int _currentSampleStateColIndex;
+    final DataIteratorContext _context;
 
     // prefetch of existing records
     int lastPrefetchRowNumber = -1;
     final IntHashMap<String> aliquotParents = new IntHashMap<>();
     final IntHashMap<Integer> aliquotRoots = new IntHashMap<>();
     final IntHashMap<Integer> sampleState = new IntHashMap<>();
 
-    public SampleUpdateAddColumnsDataIterator(DataIterator in, TableInfo target, long sampleTypeId, String keyColumnName)
+    public SampleUpdateAddColumnsDataIterator(DataIterator in, @NotNull DataIteratorContext context, TableInfo target, long sampleTypeId, String keyColumnName)
     {
         super(in);
         this._unwrapped = (CachingDataIterator)in;
+        this._context = context;
         this.target = target;
         this._sampleTypeId = sampleTypeId;
 
@@ -167,11 +170,19 @@ else if (StringUtils.isEmpty((String) key))
     @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)
+        if (!_context.getErrors().hasErrors() && ret)
+        {
             prefetchExisting();
+            if (_context.getErrors().hasErrors())
+                return false;
+        }
+
         return ret;
     }
 }

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

@@ -538,6 +538,7 @@ public Object get(int i)
     {
         if (null != _keyColumnInfo.get(i))
             return _keyValues.get(i);
+
         return _data.get(i);
     }
 

api/src/org/labkey/api/exp/api/ExperimentService.java

@@ -130,6 +130,8 @@ public interface ExperimentService extends ExperimentRunTypeSource
 
     String EXPERIMENTAL_FEATURE_FROM_EXPANCESTORS = "org.labkey.api.exp.api.ExperimentService#FROM_EXPANCESTORS";
 
+    String EXPERIMENTAL_FEATURE_ALLOW_ROW_ID_MERGE = "org.labkey.experiment.api.SampleTypeUpdateServiceDI#ALLOW_ROW_ID_SAMPLE_MERGE";
+
     int SIMPLE_PROTOCOL_FIRST_STEP_SEQUENCE = 1;
     int SIMPLE_PROTOCOL_CORE_STEP_SEQUENCE = 10;
     int SIMPLE_PROTOCOL_EXTRA_STEP_SEQUENCE = 15;
@@ -149,7 +151,7 @@ static void setInstance(ExperimentService impl)
 
     enum QueryOptions
     {
-        UseLsidForUpdate,
+        UseProvidedLsidForXarImport,
         GetSampleRecomputeCol,
         SkipBulkRemapCache,
         DeferRequiredLineageValidation,