TableInfo: ensure ExprColumn dependent columns are referenced

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

@@ -181,10 +181,10 @@ public enum MultiValuedFkType
     private final Map<String, CounterDefinition> _counterDefinitionMap = new CaseInsensitiveHashMap<>();    // Really only 1 for now, but could be more in future
 
     /* If a subclass generates a non-trivial FROM clause in getFromSQL(), it may need to track dependencies
-     * for calculated columns.  This is where we do that.  This is setup in loadAllButCustomizerFromXML(), and
-     * and used in getFromSQL(String alias, Set<FieldKey> cols)
+     * for calculated columns. Lazily built from each column's getReferencedFieldKeys() and used in
+     * getFromSQL(String alias, Set<FieldKey> cols). Cleared whenever the column shape changes.
      */
-    protected final HashMap<FieldKey, HashSet<FieldKey>> _referencedColumns = new HashMap<>();
+    private volatile Map<FieldKey, Set<FieldKey>> _referencedColumns;
 
     private boolean _initialColumnsAreAdded = false;
 
@@ -325,7 +325,6 @@ public void afterConstruct()
         }
     }
 
-
     protected Map<String, ColumnInfo> constructColumnMap()
     {
         if (isCaseSensitive())
@@ -392,25 +391,56 @@ public SQLFragment getFromSQL(String alias)
             return new SQLFragment().append("(").append(getFromSQL()).append(") ").appendIdentifier(alias);
     }
 
+    // Clear the cached resolved/referenced columns so we regenerate them when the shape of the table changes
+    private void clearColumnReferences()
+    {
+        _resolvedColumns.clear();
+        _referencedColumns = null;
+    }
+
     /** When a table a) overrides (String alias, Set<FieldKey> cols) b) has CalculatedColumns we need to make sure that
      * we include the dependent columns in the Set<>.
      */
     protected Set<FieldKey> expandColumns(Set<FieldKey> columns)
     {
-        if (null == columns || columns.isEmpty() || _referencedColumns.isEmpty())
+        if (null == columns || columns.isEmpty())
             return columns;
-        // We're not recursively expanding. However, if expressions can reference each other we'll have to.
+
+        var referenced = getReferencedColumns();
+        if (referenced.isEmpty())
+            return columns;
+
+        // We're not recursively expanding. However, if expressions can reference each other, we'll have to.
         HashSet<FieldKey> expanded = new HashSet<>();
         for (var fk : columns)
         {
             expanded.add(fk);
-            HashSet<FieldKey> refs = _referencedColumns.get(fk);
+            Set<FieldKey> refs = referenced.get(fk);
             if (null != refs)
                 expanded.addAll(refs);
         }
+
         return expanded;
     }
 
+    private Map<FieldKey, Set<FieldKey>> getReferencedColumns()
+    {
+        if (_referencedColumns == null)
+        {
+            var built = new HashMap<FieldKey, Set<FieldKey>>();
+            for (var col : getColumns())
+            {
+                var deps = col.getReferencedFieldKeys();
+                if (!deps.isEmpty())
+                    built.put(col.getFieldKey(), deps);
+            }
+
+            _referencedColumns = built.isEmpty() ? Map.of() : Collections.unmodifiableMap(built);
+        }
+
+        return _referencedColumns;
+    }
+
     @Override
     public final SQLFragment getFromSQL(String alias, Set<FieldKey> cols)
     {
@@ -439,8 +469,8 @@ protected SQLFragment getFromSQLExpanded(String alias, Set<FieldKey> cols)
             List<ColumnInfo> pkColumns = getPkColumns();
             if (pkColumns.size() != 1)
                 return new NamedObjectList();
-            else
-                return getSelectList(pkColumns.get(0), Collections.emptyList(), maxRows, titleColumnInfo);
+
+            return getSelectList(pkColumns.getFirst(), Collections.emptyList(), maxRows, titleColumnInfo);
         }
 
         ColumnInfo column = getColumn(columnName);
@@ -460,12 +490,12 @@ protected SQLFragment getFromSQLExpanded(String alias, Set<FieldKey> cols)
         final int titleIndex;
         if (!(firstColumn.equals(titleColumn)))
         {
-            cols = Arrays.asList(firstColumn, titleColumn);
+            cols = List.of(firstColumn, titleColumn);
             titleIndex = 2;
         }
         else
         {
-            cols = Arrays.asList(firstColumn);
+            cols = List.of(firstColumn);
             titleIndex = 1;
         }
 
@@ -563,7 +593,7 @@ public String getTitleColumn()
                 }
             }
             if (null == _titleColumn && !getColumns().isEmpty())
-                _titleColumn = getColumns().get(0).getName();
+                _titleColumn = getColumns().getFirst().getName();
         }
 
         return _titleColumn;
@@ -773,17 +803,18 @@ public boolean removeColumn(ColumnInfo column)
     {
         checkLocked();
         ensureInitialColumnsAreAdded();
-        // Clear the cached resolved columns so we regenerate it if the shape of the table changes
-        _resolvedColumns.clear();
-        return _columnMap.remove(column.getName()) != null;
+
+        boolean removed = _columnMap.remove(column.getName()) != null;
+        if (removed)
+            clearColumnReferences();
+
+        return removed;
     }
 
     public MutableColumnInfo addColumn(MutableColumnInfo column)
     {
         checkLocked();
         ensureInitialColumnsAreAdded();
-        // Not true if this is a VirtualTableInfo
-        // assert column.getParentTable() == this;
         if (_columnMap.containsKey(column.getName()))
         {
             String message = "Column " + column.getName() + " already exists for table " + getName() + ". Full set of existing columns: " + _columnMap.keySet();
@@ -792,9 +823,9 @@ public MutableColumnInfo addColumn(MutableColumnInfo column)
             assert false : message;
         }
         _columnMap.put(column.getName(), column);
-        // Clear the cached resolved columns so we regenerate it if the shape of the table changes
-        _resolvedColumns.clear();
+        clearColumnReferences();
         assert !(column instanceof BaseColumnInfo) || ((BaseColumnInfo)column).lockName();
+
         return column;
     }
 
@@ -803,8 +834,6 @@ public MutableColumnInfo addColumn(MutableColumnInfo column)
      * This is usually only done in TableInfo.afterConstruct() to modify the behavior of a column.
      * Because the ColumnInfo implementation can change in afterConstruct(), TableInfo implementations
      * should hold onto columnInfo references by FieldKey, and not by reference.
-
-     * during construction.
      */
     public ColumnInfo replaceColumn(ColumnInfo updated, ColumnInfo existing)
     {
@@ -819,23 +848,22 @@ public ColumnInfo replaceColumn(ColumnInfo updated, ColumnInfo existing)
             throw new IllegalStateException("Column must have the same name");
 
         _columnMap.put(updated.getName(), updated);
-        // Clear the cached resolved columns so we regenerate it if the shape of the table changes
-        _resolvedColumns.clear();
+        clearColumnReferences();
+
         return updated;
     }
 
-
     protected ColumnInfo transformColumn(MutableColumnInfo existing, @Nullable ColumnInfoTransformer t)
     {
         checkLocked();
         existing.checkLocked();
         if (null == t)
             return existing;
+
         MutableColumnInfo updated = t.apply(existing);
         return replaceColumn(updated, existing);
     }
 
-
     public void addCounterDefinition(@NotNull CounterDefinition counterDef)
     {
         boolean valid = true;
@@ -1380,7 +1408,6 @@ protected void loadAllButCustomizerFromXML(QuerySchema schema, @Nullable TableTy
         if (xmlTable.isSetTableUrl())
             _detailsURL = DetailsURL.fromXML(xmlTable.getTableUrl(), errors);
 
-
         if (xmlTable.isSetCacheSize())
             _cacheSize = xmlTable.getCacheSize();
 
@@ -1465,9 +1492,7 @@ else if (xmlColumn.isSetWrappedColumnName() && isNotBlank(xmlColumn.getWrappedCo
                 try
                 {
                     var wrappedColumn = QueryService.get().createQueryExpressionColumn(this, FieldKey.fromParts(xmlColumn.getColumnName()), sql, xmlColumn);
-                    HashSet<FieldKey> referencedColumns = new HashSet<>();
-                    QueryService.get().bindQueryExpressionColumn(wrappedColumn, existingColumns, true, referencedColumns);
-                    _referencedColumns.put(wrappedColumn.getFieldKey(), referencedColumns);
+                    QueryService.get().bindQueryExpressionColumn(wrappedColumn, existingColumns, true, null);
                     calculatedFieldKeys.add(wrappedColumn.getFieldKey());
                     addColumn(wrappedColumn);
                 }
@@ -1503,7 +1528,6 @@ else if (!calculatedFieldKeys.isEmpty())
             {
                 warnings.add(new QueryParseWarning("Invalid AuditLogging: " + auditBehavior,null,0,0));
             }
-
         }
 
         _warnings.addAll(warnings);
@@ -2140,7 +2164,7 @@ public FieldKey getAuditRowPk()
         {
             List<String> pks = getPkColumnNames();
             if (pks.size() == 1)
-                _auditRowPk = FieldKey.fromParts(pks.get(0));
+                _auditRowPk = FieldKey.fromParts(pks.getFirst());
             else if (getColumn(FieldKey.fromParts("EntityId")) != null)
                 _auditRowPk = FieldKey.fromParts("EntityId");
             else if (getColumn(FieldKey.fromParts("RowId")) != null)

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

@@ -35,8 +35,10 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public interface ColumnInfo extends ColumnRenderProperties
 {
@@ -371,6 +373,12 @@ default boolean inferIsShownInInsertView()
 
     boolean isCalculated();
 
+    /** FieldKeys this column depends on at SQL generation time. */
+    default Set<FieldKey> getReferencedFieldKeys()
+    {
+        return Collections.emptySet();
+    }
+
     default boolean isValueExpressionColumn()
     {
         return getValueExpression() != null && getWrappedColumnName() == null;
@@ -405,6 +413,7 @@ static boolean booleanFromString(String str)
     {
         return BaseColumnInfo.booleanFromString(str);
     }
+
     static boolean booleanFromObj(Object o)
     {
         return BaseColumnInfo.booleanFromObj(o);
@@ -483,5 +492,3 @@ static SimpleConvert getDefaultConvertFn(ColumnInfo col)
         return ColumnRenderProperties.getDefaultConvertFn(col);
     }
 }
-
-

api/src/org/labkey/api/query/ExprColumn.java

@@ -16,15 +16,18 @@
 
 package org.labkey.api.query;
 
-import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.lang3.Strings;
 import org.labkey.api.data.BaseColumnInfo;
 import org.labkey.api.data.ColumnInfo;
 import org.labkey.api.data.JdbcType;
 import org.labkey.api.data.MutableColumnInfo;
 import org.labkey.api.data.SQLFragment;
 import org.labkey.api.data.TableInfo;
 
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Supplier;
 
 /**
@@ -70,12 +73,11 @@ public ExprColumn(TableInfo parent, String name, SQLFragment sql, JdbcType type,
         this(parent, FieldKey.fromParts(name), sql, type, dependentColumns);
     }
 
-    public static MutableColumnInfo create(TableInfo parent, String name, JdbcType type, SQLFragment sqlf, ColumnInfo ... dependentColumns)
+    public static MutableColumnInfo create(TableInfo parent, String name, JdbcType type, SQLFragment sql, ColumnInfo ... dependentColumns)
     {
-        return new ExprColumn(parent, name, sqlf, type, dependentColumns);
+        return new ExprColumn(parent, name, sql, type, dependentColumns);
     }
 
-
     /*
      * Use this 'constructor' to create an ExprColumn whose SQL is generated on demand.  This is useful if generating
      * the SQL is at all expensive.  Most TableInfo objects are not used for generating SQL and do not need to get this value.
@@ -94,24 +96,21 @@ public SQLFragment getValueSql(String tableAlias)
         };
     }
 
-
     @Override
     public SQLFragment getValueSql(String tableAlias)
     {
         if (tableAlias.equals(STR_TABLE_ALIAS))
             return _sql;
         SQLFragment ret = new SQLFragment(_sql);
-        ret.setSqlUnsafe(StringUtils.replace(_sql.getSQL(), STR_TABLE_ALIAS, tableAlias));
+        ret.setSqlUnsafe(Strings.CS.replace(_sql.getSQL(), STR_TABLE_ALIAS, tableAlias));
         return ret;
     }
 
-
     public void setValueSQL(SQLFragment sql)
     {
         _sql = sql;
     }
 
-
     @Override
     public void declareJoins(String parentAlias, Map<String, SQLFragment> map)
     {
@@ -123,4 +122,17 @@ public void declareJoins(String parentAlias, Map<String, SQLFragment> map)
             }
         }
     }
+
+    @Override
+    public Set<FieldKey> getReferencedFieldKeys()
+    {
+        if (_dependentColumns == null || _dependentColumns.length == 0)
+            return Collections.emptySet();
+
+        var keys = new HashSet<FieldKey>();
+        for (var c : _dependentColumns)
+            keys.add(c.getFieldKey());
+
+        return Collections.unmodifiableSet(keys);
+    }
 }

query/src/org/labkey/query/QueryServiceImpl.java

@@ -3040,7 +3040,7 @@ public MutableColumnInfo createQueryExpressionColumn(TableInfo table, FieldKey k
         return ret;
     }
 
-     /** Compute and set the metadata for this column based on the source expressoin and the xml override */
+     /** Compute and set the metadata for this column based on the source expression and the xml override */
     @Override
     public void bindQueryExpressionColumn(ColumnInfo col, Map<FieldKey,ColumnInfo> columns, boolean validateOnly, @Nullable Set<FieldKey> referencedKeys) throws QueryParseException
     {

query/src/org/labkey/query/sql/CalculatedExpressionColumn.java

@@ -23,8 +23,8 @@
 import org.labkey.data.xml.ColumnType;
 import org.labkey.data.xml.TableType;
 
-import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -97,6 +97,14 @@ public void validate(Map<FieldKey, ColumnInfo> columnMap, @Nullable Set<FieldKey
         }
     }
 
+    @Override
+    public Set<FieldKey> getReferencedFieldKeys()
+    {
+        if (_allFieldKeys == null)
+            return Collections.emptySet();
+
+        return Collections.unmodifiableSet(_allFieldKeys);
+    }
 
     public void computeMetaData(Map<FieldKey,ColumnInfo> columns)
     {
@@ -397,7 +405,7 @@ public _SchemaTableInfo(DbSchema schema, DatabaseTableType tableType, String tab
         }
 
         @Override
-        protected SchemaColumnMetaData createSchemaColumnMetaData() throws SQLException
+        protected SchemaColumnMetaData createSchemaColumnMetaData()
         {
             return new SchemaColumnMetaData(this, colsToAdd, xmlToApply);
         }