feat (database): first and last formulas (baserow#5066)

* feat: add first(), last() and generalized index() formula functions

* fix: address PR review feedback and handle NaN index argument

- Rename _BaserowIndexShortcut to BaserowIndexShortcut (drop leading _)
- Remove unnecessary comments and section separators
- Rename variable c to clone in JSONBArrayGetElement
- Rename test to test_first_and_last_return_scalar_values, drop docstring
- Remove step comments in test_index_generalized.py
- Handle NaN/div-by-zero as index() argument (returns null instead of crash)

* feat: update index() formula examples to include array fields

Author: silvestrid

backend/src/baserow/contrib/database/formula/ast/function_defs.py

@@ -76,6 +76,7 @@
     BaserowExpressionContext,
     BaserowFunctionCall,
     BaserowIntegerLiteral,
+    BaserowStringLiteral,
 )
 from baserow.contrib.database.formula.expression_generator.django_expressions import (
     AndExpr,
@@ -84,6 +85,7 @@
     GreaterThanExpr,
     GreaterThanOrEqualExpr,
     IsNullExpr,
+    JSONBArrayGetElement,
     JSONBArrayJoinValues,
     JSONBArraySlice,
     JSONBArrayUniqueByValue,
@@ -265,6 +267,8 @@ def register_formula_functions(registry):
     registry.register(BaserowArrayLength())
     registry.register(BaserowArrayJoinValues())
     registry.register(BaserowArraySlice())
+    registry.register(BaserowFirst())
+    registry.register(BaserowLast())
     # ManyToMany functions
     registry.register(BaserowStringAggManyToManyValues())
     registry.register(BaserowManyToManyCount())
@@ -2607,6 +2611,45 @@ def to_django_expression(
         )
 
 
+class BaserowIndexShortcut(OneArgumentBaserowFunction):
+    arg_type = [BaserowFormulaValidType]
+    _index: int
+
+    def type_function(
+        self,
+        func_call: BaserowFunctionCall[UnTyped],
+        arg: BaserowExpression[BaserowFormulaValidType],
+    ) -> BaserowExpression[BaserowFormulaType]:
+        if arg.many:
+            arg = arg.expression_type.collapse_many(arg)
+
+        if not isinstance(arg.expression_type, BaserowFormulaArrayType):
+            return func_call.with_invalid_type(f"{self.type} requires an array input.")
+
+        from baserow.contrib.database.formula.registries import (
+            formula_function_registry,
+        )
+
+        num_type = BaserowFormulaNumberType(0)
+        index_func = formula_function_registry.get("index")
+        return index_func.call_and_type_with_args(
+            [arg, BaserowIntegerLiteral(self._index, num_type)]
+        )
+
+    def to_django_expression(self, arg: Expression) -> Expression:
+        raise NotImplementedError("type_function delegates to index")
+
+
+class BaserowFirst(BaserowIndexShortcut):
+    type = "first"
+    _index = 0
+
+
+class BaserowLast(BaserowIndexShortcut):
+    type = "last"
+    _index = -1
+
+
 class BaserowArrayLength(OneArgumentBaserowFunction):
     type = "array_length"
     arg_type = [BaserowFormulaArrayType]
@@ -2979,36 +3022,105 @@ def to_django_expression(self, arg: Expression) -> Expression:
         )
 
 
-class BaserowIndex(TwoArgumentBaserowFunction):
-    arg1_type = [BaserowFormulaArrayType]
-    arg2_type = [BaserowFormulaNumberType]
+def _index_output_field(mode):
+    """Return a fresh Django output_field for the given extraction mode."""
+
+    from baserow.contrib.database.formula.types.formula_types import (
+        _lookup_formula_type_from_string,
+    )
+
+    try:
+        return _lookup_formula_type_from_string(mode).output_field_class()
+    except Exception:
+        return fields.TextField()
+
+
+def _unwrap_literal_value(django_expr):
+    """
+    Extract the Python value from a Django expression that wraps a
+    ``Value(...)`` — e.g. ``Cast(Value('x'), TextField())``.
+    """
+
+    while not hasattr(django_expr, "value"):
+        if (
+            hasattr(django_expr, "source_expressions")
+            and django_expr.source_expressions
+        ):
+            django_expr = django_expr.source_expressions[0]
+        else:
+            return None
+    return django_expr.value
+
 
+class BaserowIndex(BaserowFunctionDefinition):
     type = "index"
+    num_args = NumOfArgsBetween(2, 4)
 
-    def type_function(
+    @property
+    def arg_types(self) -> BaserowArgumentTypeChecker:
+        def type_checker(arg_index, arg_types):
+            if arg_index == 0:
+                return [BaserowFormulaValidType]
+            elif arg_index == 1:
+                return [BaserowFormulaNumberType]
+            else:
+                return [BaserowFormulaTextType]  # mode + sql literals
+
+        return type_checker
+
+    def type_function_given_valid_args(
         self,
+        args: List[BaserowExpression[BaserowFormulaValidType]],
         func_call: BaserowFunctionCall[UnTyped],
-        arg1: BaserowExpression[BaserowFormulaValidType],
-        arg2: BaserowExpression[BaserowFormulaValidType],
     ) -> BaserowExpression[BaserowFormulaType]:
-        if not isinstance(arg1.expression_type.sub_type, BaserowFormulaSingleFileType):
+        if len(args) not in (2, 4):
             return func_call.with_invalid_type(
-                "index only currently supports indexing file fields."
-            )
-        else:
-            if arg1.many:
-                arg1 = arg1.expression_type.collapse_many(arg1)
-            return func_call.with_args([arg1, arg2]).with_valid_type(
-                arg1.expression_type.sub_type
+                "index requires exactly 2 arguments: an array and an index."
             )
 
-    def to_django_expression(self, arg1: Expression, arg2: Expression) -> Expression:
-        return Func(
-            arg1,
-            Cast(arg2, fields.TextField()),
-            function="jsonb_extract_path",
-            output_field=JSONField(),
+        arg1, arg2 = args[0], args[1]
+
+        if arg1.many:
+            arg1 = arg1.expression_type.collapse_many(arg1)
+
+        if not isinstance(arg1.expression_type, BaserowFormulaArrayType):
+            return func_call.with_invalid_type("index requires an array input.")
+
+        sub_type = arg1.expression_type.sub_type
+
+        if len(args) == 4:
+            return func_call.with_args(list(args)).with_valid_type(sub_type)
+
+        mode_literal = BaserowStringLiteral(
+            sub_type.array_index_mode, BaserowFormulaTextType()
         )
+        sql_literal = BaserowStringLiteral(
+            sub_type.array_index_sql, BaserowFormulaTextType()
+        )
+
+        return func_call.with_args(
+            [arg1, arg2, mode_literal, sql_literal]
+        ).with_valid_type(sub_type)
+
+    def to_django_expression_given_args(
+        self,
+        args: List["WrappedExpressionWithMetadata"],
+        context: BaserowExpressionContext,
+    ) -> "WrappedExpressionWithMetadata":
+        mode = _unwrap_literal_value(args[2].expression) or "text"
+        value_sql = _unwrap_literal_value(args[3].expression) or "{elem} ->> 'value'"
+        safe_index = handle_arg_being_nan(
+            args[1].expression,
+            Value(None, output_field=fields.IntegerField()),
+            args[1].expression,
+        )
+        expr = JSONBArrayGetElement(
+            args[0].expression,
+            safe_index,
+            value_sql,
+            _index_output_field(mode),
+        )
+        return WrappedExpressionWithMetadata.from_args(expr, args)
 
 
 class BaserowJsonbExtractPathText(BaserowFunctionDefinition):

backend/src/baserow/contrib/database/formula/expression_generator/django_expressions.py

@@ -163,6 +163,47 @@ def as_sql(self, compiler, connection, **extra_context):
         return sql, (*separator_params, *params)
 
 
+class JSONBArrayGetElement(Expression):
+    """
+    Extract a single element from a JSONB array by 0-based index (negative
+    counts from end) and optionally unwrap / cast the ``value`` key.
+
+    *value_sql* is a SQL template with an ``{elem}`` placeholder that controls
+    how the element is extracted (e.g. ``({elem} ->> 'value')::numeric``).
+    Each formula type provides its own template via ``array_index_sql``.
+
+    PostgreSQL's ``->`` operator natively handles negative indices and returns
+    NULL for out-of-bounds, so no CASE expression is needed.
+    """
+
+    def __init__(self, array_expr, index_expr, value_sql, output_field):
+        super().__init__(output_field=output_field)
+        self.array_expr = array_expr
+        self.index_expr = index_expr
+        self.value_sql = value_sql
+
+    def resolve_expression(
+        self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False
+    ):
+        clone = self.copy()
+        clone.is_summary = summarize
+        clone.array_expr = self.array_expr.resolve_expression(
+            query, allow_joins, reuse, summarize, for_save
+        )
+        clone.index_expr = self.index_expr.resolve_expression(
+            query, allow_joins, reuse, summarize, for_save
+        )
+        return clone
+
+    def as_sql(self, compiler, connection):
+        arr_sql, arr_params = compiler.compile(self.array_expr)
+        idx_sql, idx_params = compiler.compile(self.index_expr)
+
+        elem_sql = f"({arr_sql}) -> ({idx_sql})::int"
+        sql = f"({self.value_sql.format(elem=elem_sql)})"
+        return sql, list(arr_params) + list(idx_params)
+
+
 class JSONBArraySlice(Expression):
     """
     Slice a JSONB array with offset, limit, and optional reverse.

backend/src/baserow/contrib/database/formula/types/formula_type.py

@@ -1,6 +1,7 @@
 import abc
 from typing import TYPE_CHECKING, List, Type, TypeVar
 
+from django.db import models
 from django.db.models import Expression, F, Model, Value
 from django.utils.functional import classproperty
 
@@ -285,6 +286,18 @@ def can_represent_collaborators(self) -> bool:
     def item_is_in_nested_value_object_when_in_array(self) -> bool:
         return True
 
+    @property
+    def array_index_mode(self) -> str:
+        return self.type
+
+    @property
+    def array_index_sql(self) -> str:
+        if not self.item_is_in_nested_value_object_when_in_array:
+            return "{elem}"
+        return "{elem} ->> 'value'"
+
+    output_field_class = models.TextField
+
     @property
     def can_have_db_index(self) -> bool:
         return False

backend/src/baserow/contrib/database/formula/types/formula_types.py

@@ -80,6 +80,9 @@
 
 
 class BaserowJSONBObjectBaseType(BaserowFormulaValidType, ABC):
+    array_index_sql = "{elem} -> 'value'"
+    output_field_class = JSONField
+
     def parse_filter_value(self, field, model_field, value):
         """
         Since the subclasses don't have a baserow_field_type or data might be stored
@@ -361,6 +364,8 @@ class BaserowFormulaNumberType(
 ):
     type = "number"
     baserow_field_type = "number"
+    array_index_sql = "({elem} ->> 'value')::numeric"
+    output_field_class = models.DecimalField
     user_overridable_formatting_option_fields = [
         "number_decimal_places",
         "number_prefix",
@@ -516,6 +521,8 @@ class BaserowFormulaBooleanType(
 ):
     type = "boolean"
     baserow_field_type = "boolean"
+    array_index_sql = "({elem} ->> 'value')::boolean"
+    output_field_class = models.BooleanField
     can_order_by_in_array = True
     can_group_by = True
     can_have_db_index = True
@@ -715,6 +722,8 @@ class BaserowFormulaDurationType(
 ):
     type = "duration"
     baserow_field_type = "duration"
+    array_index_sql = "({elem} ->> 'value')::interval"
+    output_field_class = models.DurationField
     user_overridable_formatting_option_fields = ["duration_format"]
     can_group_by = True
     can_order_by_in_array = True
@@ -856,6 +865,7 @@ class BaserowFormulaDateType(
     can_order_by_in_array = True
     can_group_by = True
     can_have_db_index = True
+    output_field_class = models.DateTimeField
 
     def __init__(
         self,
@@ -873,6 +883,11 @@ def __init__(
         self.date_show_tzinfo = date_show_tzinfo
         self.date_force_timezone = date_force_timezone
 
+    @property
+    def array_index_sql(self) -> str:
+        cast = "::timestamptz" if self.date_include_time else "::date"
+        return f"({{elem}} ->> 'value'){cast}"
+
     @property
     def comparable_types(self) -> List[Type["BaserowFormulaValidType"]]:
         return [
@@ -1003,6 +1018,7 @@ class BaserowFormulaSingleFileType(
     can_order_by_in_array = False
     baserow_field_type = None
     item_is_in_nested_value_object_when_in_array = False
+    array_index_sql = "{elem}"
     can_represent_files = True
 
     def is_searchable(self, field):

backend/tests/baserow/contrib/automation/history/utils.py

@@ -6,7 +6,11 @@ def assert_history(
 ):
     """Helper to test AutomationWorkflowHistory objects."""
 
-    histories = list(AutomationWorkflowHistory.objects.filter(workflow=workflow))
+    histories = list(
+        AutomationWorkflowHistory.objects.filter(workflow=workflow).order_by(
+            "started_on", "id"
+        )
+    )
     assert len(histories) == expected_count
     if expected_count > 0:
         history = histories[history_index]

backend/tests/baserow/contrib/database/formula/test_baserow_formula_results.py

@@ -3063,6 +3063,43 @@ def test_array_slice_empty_array(data_fixture):
     assert getattr(result, ref_slice_field.db_column) == []
 
 
+@pytest.mark.django_db
+def test_first_and_last_return_scalar_values(data_fixture):
+    user = data_fixture.create_user()
+    table_a, table_b, link_field = data_fixture.create_two_linked_tables(user=user)
+    text_field, b_rows, row_a1 = _setup_text_5_rows(
+        data_fixture, table_a, table_b, link_field, user
+    )
+
+    lookup_field = FieldHandler().create_field(
+        user,
+        table_a,
+        "formula",
+        name="lookup",
+        formula=f"lookup('{link_field.name}', '{text_field.name}')",
+    )
+    first_field = FieldHandler().create_field(
+        user,
+        table_a,
+        "formula",
+        name="first_val",
+        formula="first(field('lookup'))",
+    )
+    last_field = FieldHandler().create_field(
+        user,
+        table_a,
+        "formula",
+        name="last_val",
+        formula="last(field('lookup'))",
+    )
+
+    table_a_model = table_a.get_model()
+    result = table_a_model.objects.get(id=row_a1.id)
+
+    assert getattr(result, first_field.db_column) == "A"
+    assert getattr(result, last_field.db_column) == "E"
+
+
 @pytest.mark.django_db
 def test_array_slice_rejects_non_array_input(data_fixture):
     user = data_fixture.create_user()