InfluxDB3Sink: improve timestamp type handling (#811)

* improve timestamp type handling

* added some extra validation

* cleaned up a little

* remove extra comment

* add timezone info to default min/max

* add way to pass custom timestamp setter

* fix a bug with int checking accidentally matching on bools

* update docs and finish adding tests

* Update quixstreams/sinks/core/influxdb3.py

Co-authored-by: Remy Gwaramadze <[email protected]>

---------

Co-authored-by: Remy Gwaramadze <[email protected]>

Author: tim-quix

docs/connectors/sinks/influxdb3-sink.md

@@ -96,16 +96,27 @@ InfluxDB3Sink accepts the following configuration parameters:
 
 - `measurement` - a measurement name, required.
 
-- `fields_keys` - a list of keys to be used as "fields" when writing to InfluxDB.  
+- `fields_keys` - an iterable (list) of strings used as InfluxDB "fields".  
+  Also accepts a single-argument callable that receives the current message data as a dict and returns an iterable of strings.
+  - If present, it must not overlap with "tags_keys".
+  - If empty, the whole record value will be used.  
 See the [What data can be sent to InfluxDB](#what-data-can-be-sent-to-influxdb) for more info.
 
-- `tags_keys` - a list of keys to be used as "tags" when writing to InfluxDB.  
+- `tags_keys` - an iterable (list) of strings used as InfluxDB "tags".  
+  Also accepts a single-argument callable that receives the current message data as a 
+  dict and returns an iterable of strings.
+  - If present, it must not overlap with "fields_keys".
+  - Given keys are popped from the value dictionary since the same key
+    cannot be both a tag and field.
+  - If empty, no tags will be sent.  
 See the [What data can be sent to InfluxDB](#what-data-can-be-sent-to-influxdb) for more info.
 
-            
-- `time_key` - a key to be used as "time" when writing to InfluxDB.  
-By default, the record timestamp will be used with millisecond time precision.
-When using a custom key, you may need to adjust the `time_precision` setting to match.
+- `time_setter` - an optional column name to use as "time" for InfluxDB.  
+  Also accepts a callable which receives the current message data and
+  returns either the desired time or `None` (use default).  
+  - The time can be an `int`, `string` (RFC3339 format), or `datetime`.
+  - The time must match the `time_precision` argument if not a `datetime` object, else raises.
+  - By default, a record's kafka timestamp with `"ms"` time precision is used.
 
 - `time_precision` - a time precision to use when writing to InfluxDB.  
 Default - `ms`.

quixstreams/sinks/core/influxdb3.py

@@ -1,6 +1,7 @@
 import logging
 import sys
 import time
+from datetime import datetime, timezone
 from typing import Any, Callable, Iterable, Literal, Mapping, Optional, Union, get_args
 
 from quixstreams.models import HeadersTuples
@@ -27,25 +28,32 @@
 
 
 TimePrecision = Literal["ms", "ns", "us", "s"]
+TIME_PRECISION_LEN = {
+    "s": 10,
+    "ms": 13,
+    "ns": 16,
+    "us": 19,
+}
 
 InfluxDBValueMap = dict[str, Union[str, int, float, bool]]
 
 FieldsCallable = Callable[[InfluxDBValueMap], Iterable[str]]
 MeasurementCallable = Callable[[InfluxDBValueMap], str]
 TagsCallable = Callable[[InfluxDBValueMap], Iterable[str]]
-
+TimeCallable = Callable[[InfluxDBValueMap], Optional[Union[str, int, datetime]]]
 
 FieldsSetter = Union[Iterable[str], FieldsCallable]
 MeasurementSetter = Union[str, MeasurementCallable]
 TagsSetter = Union[Iterable[str], TagsCallable]
+TimeSetter = Union[str, TimeCallable]
 
 
 class InfluxDB3Sink(BatchingSink):
     _TIME_PRECISIONS = {
+        "s": WritePrecision.S,
         "ms": WritePrecision.MS,
         "ns": WritePrecision.NS,
         "us": WritePrecision.US,
-        "s": WritePrecision.S,
     }
 
     def __init__(
@@ -57,7 +65,7 @@ def __init__(
         measurement: MeasurementSetter,
         fields_keys: FieldsSetter = (),
         tags_keys: TagsSetter = (),
-        time_key: Optional[str] = None,
+        time_setter: Optional[TimeSetter] = None,
         time_precision: TimePrecision = "ms",
         allow_missing_fields: bool = False,
         include_metadata_tags: bool = False,
@@ -108,10 +116,12 @@ def __init__(
             - If empty, no tags will be sent.
             >***NOTE***: InfluxDB client always converts tag values to strings.
             Default - `()`.
-        :param time_key: a key to be used as "time" when writing to InfluxDB.
-            By default, the record timestamp will be used with "ms" time precision.
-            When using a custom key, you may need to adjust the `time_precision` setting
-            to match.
+        :param time_setter: an optional column name to use as "time" for InfluxDB.
+            Also accepts a callable which receives the current message data and
+            returns either the desired time or `None` (use default).
+            The time can be an `int`, `string` (RFC3339 format), or `datetime`.
+            The time must match the `time_precision` argument if not a `datetime` object, else raises.
+            By default, a record's kafka timestamp with "ms" time precision is used.
         :param time_precision: a time precision to use when writing to InfluxDB.
             Possible values: "ms", "ns", "us", "s".
             Default - `"ms"`.
@@ -173,31 +183,16 @@ def __init__(
             },
         }
         self._client: Optional[InfluxDBClient3] = None
-        self._measurement = self._measurement_callable(measurement)
-        self._fields_keys = self._fields_callable(fields_keys)
-        self._tags_keys = self._tags_callable(tags_keys)
+        self._measurement = _measurement_callable(measurement)
+        self._fields_keys = _fields_callable(fields_keys)
+        self._tags_keys = _tags_callable(tags_keys)
+        self._time_setter = _time_callable(time_setter)
         self._include_metadata_tags = include_metadata_tags
-        self._time_key = time_key
         self._write_precision = self._TIME_PRECISIONS[time_precision]
         self._batch_size = batch_size
         self._allow_missing_fields = allow_missing_fields
         self._convert_ints_to_floats = convert_ints_to_floats
 
-    def _measurement_callable(self, setter: MeasurementSetter) -> MeasurementCallable:
-        if callable(setter):
-            return setter
-        return lambda value: setter
-
-    def _fields_callable(self, setter: FieldsSetter) -> FieldsCallable:
-        if callable(setter):
-            return setter
-        return lambda value: setter
-
-    def _tags_callable(self, setter: TagsSetter) -> TagsCallable:
-        if callable(setter):
-            return setter
-        return lambda value: setter
-
     def setup(self):
         self._client = InfluxDBClient3(**self._client_args)
         try:
@@ -241,19 +236,20 @@ def write(self, batch: SinkBatch):
         measurement = self._measurement
         fields_keys = self._fields_keys
         tags_keys = self._tags_keys
-        time_key = self._time_key
+        time_setter = self._time_setter
         for write_batch in batch.iter_chunks(n=self._batch_size):
             records = []
 
-            min_timestamp = sys.maxsize
-            max_timestamp = -1
+            min_timestamp = None
+            max_timestamp = None
 
             for item in write_batch:
                 value = item.value
                 # Evaluate these before we alter the value
                 _measurement = measurement(value)
                 _tags_keys = tags_keys(value)
                 _fields_keys = fields_keys(value)
+                ts = time_setter(value)
 
                 tags = {}
                 for tag_key in _tags_keys:
@@ -265,6 +261,24 @@ def write(self, batch: SinkBatch):
                         tag = value.pop(tag_key)
                         tags[tag_key] = tag
 
+                if ts is None:
+                    ts = item.timestamp
+                # Note: currently NOT validating the timestamp itself is valid
+                elif not isinstance(ts, valid := (str, int, datetime)):
+                    raise TypeError(
+                        f'InfluxDB3 "time" field expects: {valid}, got {type(ts)}'
+                    )
+
+                if isinstance(ts, int):
+                    time_len = len(str(ts))
+                    expected = TIME_PRECISION_LEN[self._write_precision]
+                    if time_len != expected:
+                        raise ValueError(
+                            f'`time_precision` of "{self._write_precision}" '
+                            f"expects a {expected}-digit integer epoch, "
+                            f"got {time_len} (timestamp: {ts})."
+                        )
+
                 if self._include_metadata_tags:
                     tags["__key"] = item.key
                     tags["__topic"] = batch.topic
@@ -281,20 +295,19 @@ def write(self, batch: SinkBatch):
 
                 if self._convert_ints_to_floats:
                     fields = {
-                        k: float(v) if isinstance(v, int) else v
+                        k: float(v) if type(v) is int else v  # avoids bool matching
                         for k, v in fields.items()
                     }
 
-                ts = value[time_key] if time_key is not None else item.timestamp
                 record = {
                     "measurement": _measurement,
                     "tags": tags,
                     "fields": fields,
                     "time": ts,
                 }
                 records.append(record)
-                min_timestamp = min(ts, min_timestamp)
-                max_timestamp = max(ts, max_timestamp)
+                min_timestamp = min(ts, min_timestamp or _ts_min_default(ts))
+                max_timestamp = max(ts, max_timestamp or _ts_max_default(ts))
 
             try:
                 _start = time.monotonic()
@@ -317,3 +330,47 @@ def write(self, batch: SinkBatch):
                         retry_after=int(exc.retry_after)
                     ) from exc
                 raise
+
+
+def _ts_min_default(timestamp: Union[int, str, datetime]):
+    if isinstance(timestamp, int):
+        return sys.maxsize
+    elif isinstance(timestamp, str):
+        return "~"  # lexicographically largest ASCII char
+    elif isinstance(timestamp, datetime):
+        return datetime.max.replace(tzinfo=timezone.utc)
+
+
+def _ts_max_default(timestamp: Union[int, str, datetime]):
+    if isinstance(timestamp, int):
+        return -1
+    elif isinstance(timestamp, str):
+        return ""
+    elif isinstance(timestamp, datetime):
+        return datetime.min.replace(tzinfo=timezone.utc)
+
+
+def _measurement_callable(setter: MeasurementSetter) -> MeasurementCallable:
+    if callable(setter):
+        return setter
+    return lambda value: setter
+
+
+def _fields_callable(setter: FieldsSetter) -> FieldsCallable:
+    if callable(setter):
+        return setter
+    return lambda value: setter
+
+
+def _tags_callable(setter: TagsSetter) -> TagsCallable:
+    if callable(setter):
+        return setter
+    return lambda value: setter
+
+
+def _time_callable(setter: Optional[TimeSetter]) -> TimeCallable:
+    if callable(setter):
+        return setter
+    if isinstance(setter, str):
+        return lambda value: value[setter]  # type:  ignore
+    return lambda value: None  # the kafka timestamp will be used