feat(firestore): add new pipeline expressions (#16151)

Adds the following expressions:
- cmp
- timestamp_trunc
- timestamp_diff
- timestamp_extract
- if_null
- type
- is_type

Author: daniel-sanche

packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py

@@ -48,6 +48,75 @@
 )
 
 
+classTimeUnit(str, Enum):
+"""Enumeration of the different time units supported by the Firestore backend."""
+
+MICROSECOND="microsecond"
+MILLISECOND="millisecond"
+SECOND="second"
+MINUTE="minute"
+HOUR="hour"
+DAY="day"
+
+
+classTimeGranularity(str, Enum):
+"""Enumeration of the different time granularities supported by the Firestore backend."""
+
+# Inherit from TimeUnit
+MICROSECOND=TimeUnit.MICROSECOND.value
+MILLISECOND=TimeUnit.MILLISECOND.value
+SECOND=TimeUnit.SECOND.value
+MINUTE=TimeUnit.MINUTE.value
+HOUR=TimeUnit.HOUR.value
+DAY=TimeUnit.DAY.value
+
+# Additional granularities
+WEEK="week"
+WEEK_MONDAY="week(monday)"
+WEEK_TUESDAY="week(tuesday)"
+WEEK_WEDNESDAY="week(wednesday)"
+WEEK_THURSDAY="week(thursday)"
+WEEK_FRIDAY="week(friday)"
+WEEK_SATURDAY="week(saturday)"
+WEEK_SUNDAY="week(sunday)"
+ISOWEEK="isoweek"
+MONTH="month"
+QUARTER="quarter"
+YEAR="year"
+ISOYEAR="isoyear"
+
+
+classTimePart(str, Enum):
+"""Enumeration of the different time parts supported by the Firestore backend."""
+
+# Inherit from TimeUnit
+MICROSECOND=TimeUnit.MICROSECOND.value
+MILLISECOND=TimeUnit.MILLISECOND.value
+SECOND=TimeUnit.SECOND.value
+MINUTE=TimeUnit.MINUTE.value
+HOUR=TimeUnit.HOUR.value
+DAY=TimeUnit.DAY.value
+
+# Inherit from TimeGranularity
+WEEK=TimeGranularity.WEEK.value
+WEEK_MONDAY=TimeGranularity.WEEK_MONDAY.value
+WEEK_TUESDAY=TimeGranularity.WEEK_TUESDAY.value
+WEEK_WEDNESDAY=TimeGranularity.WEEK_WEDNESDAY.value
+WEEK_THURSDAY=TimeGranularity.WEEK_THURSDAY.value
+WEEK_FRIDAY=TimeGranularity.WEEK_FRIDAY.value
+WEEK_SATURDAY=TimeGranularity.WEEK_SATURDAY.value
+WEEK_SUNDAY=TimeGranularity.WEEK_SUNDAY.value
+ISOWEEK=TimeGranularity.ISOWEEK.value
+MONTH=TimeGranularity.MONTH.value
+QUARTER=TimeGranularity.QUARTER.value
+YEAR=TimeGranularity.YEAR.value
+ISOYEAR=TimeGranularity.ISOYEAR.value
+
+# Additional parts
+DAY_OF_WEEK="dayofweek"
+DAY_OF_YEAR="dayofyear"
+
+
 classOrdering:
 """Represents the direction for sorting results in a pipeline."""
 
@@ -90,6 +159,31 @@ def _to_pb(self) -> Value:
  )
 
 
+classType(str, Enum):
+"""Enumeration of the different types generated by the Firestore backend."""
+
+NULL="null"
+ARRAY="array"
+BOOLEAN="boolean"
+BYTES="bytes"
+TIMESTAMP="timestamp"
+GEO_POINT="geo_point"
+NUMBER="number"
+INT32="int32"
+INT64="int64"
+FLOAT64="float64"
+DECIMAL128="decimal128"
+MAP="map"
+REFERENCE="reference"
+STRING="string"
+VECTOR="vector"
+MAX_KEY="max_key"
+MIN_KEY="min_key"
+OBJECT_ID="object_id"
+REGEX="regex"
+REQUEST_TIMESTAMP="request_timestamp"
+
+
 classExpression(ABC):
 """Represents an expression that can be evaluated to a value within the
  execution of a pipeline.
@@ -120,6 +214,8 @@ def _cast_to_expr_or_convert_to_constant(
 """Convert arbitrary object to an Expression."""
 ifisinstance(o, Expression):
 returno
+ifisinstance(o, Enum):
+o=o.value
 ifisinstance(o, dict):
 returnMap(o)
 ifisinstance(o, list):
@@ -2079,20 +2175,21 @@ def unix_seconds_to_timestamp(self) -> "Expression":
 
 @expose_as_static
 deftimestamp_add(
-self, unit: Expression|str, amount: Expression|float
+self, unit: TimeUnit|str|Expression, amount: Expression|float
  ) ->"Expression":
 """Creates an expression that adds a specified amount of time to this timestamp expression.
 
  Example:
+ >>> # Add 1.5 days to the 'timestamp' field using TimeUnit enum.
+ >>> Field.of("timestamp").timestamp_add(TimeUnit.DAY, 1.5)
  >>> # Add a duration specified by the 'unit' and 'amount' fields to the 'timestamp' field.
  >>> Field.of("timestamp").timestamp_add(Field.of("unit"), Field.of("amount"))
- >>> # Add 1.5 days to the 'timestamp' field.
+ >>> # Add 1.5 days to the 'timestamp' field using a string.
  >>> Field.of("timestamp").timestamp_add("day", 1.5)
 
  Args:
- unit: The expression or string evaluating to the unit of time to add, must be one of
- 'microsecond', 'millisecond', 'second', 'minute', 'hour', 'day'.
- amount: The expression or float representing the amount of time to add.
+ unit: The unit of time to add.
+ amount: The amount of time to add.
 
  Returns:
  A new `Expression` representing the resulting timestamp.
@@ -2108,20 +2205,21 @@ def timestamp_add(
 
 @expose_as_static
 deftimestamp_subtract(
-self, unit: Expression|str, amount: Expression|float
+self, unit: TimeUnit|str|Expression, amount: Expression|float
  ) ->"Expression":
 """Creates an expression that subtracts a specified amount of time from this timestamp expression.
 
  Example:
+ >>> # Subtract 2.5 hours from the 'timestamp' field using TimeUnit enum.
+ >>> Field.of("timestamp").timestamp_subtract(TimeUnit.HOUR, 2.5)
  >>> # Subtract a duration specified by the 'unit' and 'amount' fields from the 'timestamp' field.
  >>> Field.of("timestamp").timestamp_subtract(Field.of("unit"), Field.of("amount"))
- >>> # Subtract 2.5 hours from the 'timestamp' field.
+ >>> # Subtract 2.5 hours from the 'timestamp' field using a string.
  >>> Field.of("timestamp").timestamp_subtract("hour", 2.5)
 
  Args:
- unit: The expression or string evaluating to the unit of time to subtract, must be one of
- 'microsecond', 'millisecond', 'second', 'minute', 'hour', 'day'.
- amount: The expression or float representing the amount of time to subtract.
+ unit: The unit of time to subtract.
+ amount: The amount of time to subtract.
 
  Returns:
  A new `Expression` representing the resulting timestamp.
@@ -2206,6 +2304,163 @@ def as_(self, alias: str) -> "AliasedExpression":
  """
 returnAliasedExpression(self, alias)
 
+@expose_as_static
+defcmp(self, other: Expression|CONSTANT_TYPE) ->"Expression":
+"""Creates an expression that compares this expression to another expression.
+
+ Returns an integer:
+ * -1 if this expression is less than the other
+ * 0 if they are equal
+ * 1 if this expression is greater than the other
+
+ Example:
+ >>> # Compare the 'price' field to 10
+ >>> Field.of("price").cmp(10)
+
+ Args:
+ other: The value to compare against.
+
+ Returns:
+ A new `Expression` representing the comparison operation.
+ """
+returnFunctionExpression(
+"cmp", [self, self._cast_to_expr_or_convert_to_constant(other)]
+ )
+
+@expose_as_static
+deftimestamp_trunc(
+self,
+granularity: TimeGranularity|Expression|str,
+timezone: Expression|str|None=None,
+ ) ->"Expression":
+"""Creates an expression that truncates a timestamp to a specified granularity.
+
+ Example:
+ >>> # Truncate the 'createdAt' field to the day using TimeGranularity enum
+ >>> Field.of("createdAt").timestamp_trunc(TimeGranularity.DAY)
+ >>> # Truncate the 'createdAt' field to the day in the 'America/Los_Angeles' timezone
+ >>> Field.of("createdAt").timestamp_trunc(TimeGranularity.DAY, "America/Los_Angeles")
+ >>> # Truncate the 'createdAt' field to the day using a string
+ >>> Field.of("createdAt").timestamp_trunc("day")
+
+ Args:
+ granularity: The granularity to truncate to.
+ timezone: The optional timezone.
+
+ Returns:
+ A new `Expression` representing the timestamp_trunc operation.
+ """
+args= [self, self._cast_to_expr_or_convert_to_constant(granularity)]
+iftimezoneisnotNone:
+args.append(self._cast_to_expr_or_convert_to_constant(timezone))
+returnFunctionExpression("timestamp_trunc", args)
+
+@expose_as_static
+deftimestamp_extract(
+self,
+part: TimePart|str|Expression,
+timezone: str|Expression|None=None,
+ ) ->"Expression":
+"""Creates an expression that extracts a part of a timestamp.
+
+ Example:
+ >>> # Extract the year from the 'createdAt' field using TimePart enum
+ >>> Field.of("createdAt").timestamp_extract(TimePart.YEAR)
+ >>> # Extract the year from the 'createdAt' field in the 'America/Los_Angeles' timezone
+ >>> Field.of("createdAt").timestamp_extract(TimePart.YEAR, "America/Los_Angeles")
+ >>> # Extract the year from the 'createdAt' field using a string
+ >>> Field.of("createdAt").timestamp_extract("year")
+
+ Args:
+ part: The part to extract.
+ timezone: The optional timezone.
+
+ Returns:
+ A new `Expression` representing the timestamp_extract operation.
+ """
+args= [self, self._cast_to_expr_or_convert_to_constant(part)]
+iftimezoneisnotNone:
+args.append(self._cast_to_expr_or_convert_to_constant(timezone))
+returnFunctionExpression("timestamp_extract", args)
+
+@expose_as_static
+deftimestamp_diff(
+self, start: Expression|datetime.datetime, unit: TimeUnit|str|Expression
+ ) ->"Expression":
+"""Creates an expression that computes the difference between two timestamps in the specified unit.
+
+ Example:
+ >>> # Compute the difference in days between the 'end' field and the 'start' field using TimeUnit enum
+ >>> Field.of("end").timestamp_diff(Field.of("start"), TimeUnit.DAY)
+ >>> # Compute the difference in days using a string
+ >>> Field.of("end").timestamp_diff(Field.of("start"), "day")
+
+ Args:
+ start: The start timestamp.
+ unit: The unit of time.
+
+ Returns:
+ A new `Expression` representing the timestamp_diff operation.
+ """
+returnFunctionExpression(
+"timestamp_diff",
+ [
+self,
+self._cast_to_expr_or_convert_to_constant(start),
+self._cast_to_expr_or_convert_to_constant(unit),
+ ],
+ )
+
+@expose_as_static
+defif_null(self, *others: Expression|CONSTANT_TYPE) ->"Expression":
+"""Creates an expression that returns the first non-null expression from the provided arguments.
+
+ Example:
+ >>> # Return the 'nickname' field if not null, otherwise return 'firstName'
+ >>> Field.of("nickname").if_null(Field.of("firstName"))
+
+ Args:
+ *others: Additional expressions or constants to evaluate if the previous ones are null.
+
+ Returns:
+ A new `Expression` representing the if_null operation.
+ """
+returnFunctionExpression(
+"if_null",
+ [self] + [self._cast_to_expr_or_convert_to_constant(o) foroinothers],
+ )
+
+@expose_as_static
+deftype(self) ->"Expression":
+"""Creates an expression that returns the data type of this expression's result as a string.
+
+ Example:
+ >>> # Get the type of the 'title' field
+ >>> Field.of("title").type()
+
+ Returns:
+ A new `Expression` representing the type operation.
+ """
+returnFunctionExpression("type", [self])
+
+@expose_as_static
+defis_type(self, type_val: Type|str|Expression) ->"BooleanExpression":
+"""Creates an expression that checks if the result is of the specified type.
+
+ Example:
+ >>> # Check if the 'price' field is a number
+ >>> Field.of("price").is_type("number")
+
+ Args:
+ type_val: The type string or expression to check against.
+
+ Returns:
+ A new `BooleanExpression` representing the is_type operation.
+ """
+returnBooleanExpression(
+"is_type", [self, self._cast_to_expr_or_convert_to_constant(type_val)]
+ )
+
 
 classConstant(Expression, Generic[CONSTANT_TYPE]):
 """Represents a constant literal value in an expression."""