fix typing issues

Author: aaronsteers

airbyte_cdk/sources/declarative/async_job/timer.py

@@ -2,6 +2,8 @@
 from datetime import datetime, timedelta, timezone
 from typing import Optional
 
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
+
 
 class Timer:
     def __init__(self, timeout: timedelta) -> None:
@@ -36,4 +38,4 @@ def has_timed_out(self) -> bool:
 
     @staticmethod
     def _now() -> datetime:
-        return datetime.now(tz=timezone.utc)
+        return ab_datetime_now()

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -14,6 +14,7 @@
 from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
 from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 
 
 class JwtAlgorithm(str):
@@ -127,7 +128,7 @@ def _get_jwt_payload(self) -> dict[str, Any]:
         """
         Builds and returns the payload used when signing the JWT.
         """
-        now = int(datetime.now().timestamp())
+        now = ab_datetime_now().timestamp()
         exp = now + self._token_duration if isinstance(self._token_duration, int) else now
         nbf = now
 

airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py

@@ -21,7 +21,11 @@
 )
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
-from airbyte_cdk.utils.datetime_helpers import ab_datetime_format, ab_datetime_parse
+from airbyte_cdk.utils.datetime_helpers import (
+    ab_datetime_format,
+    ab_datetime_now,
+    ab_datetime_parse,
+)
 from airbyte_cdk.utils.mapping_helpers import _validate_component_request_option_paths
 
 
@@ -90,7 +94,6 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
             None if not self.end_datetime else MinMaxDatetime.create(self.end_datetime, parameters)
         )
 
-        self._timezone = datetime.timezone.utc
         self._interpolation = JinjaInterpolation()
 
         self._step = (
@@ -240,10 +243,10 @@ def select_best_end_datetime(self) -> datetime.datetime:
 
         :return datetime.datetime: The best end datetime, which is either the current datetime or the pre-configured end datetime, whichever is earlier.
         """
-        now = datetime.datetime.now(tz=self._timezone)
+        now = ab_datetime_now()
         if not self._end_datetime:
             return now
-        return min(self._end_datetime.get_datetime(self.config), now)
+        return min(self._end_datetime.get_datetime(self.config), now.to_datetime())
 
     def _calculate_cursor_datetime_from_state(
         self, stream_state: Mapping[str, Any]

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -546,6 +546,7 @@
 from airbyte_cdk.sources.streams.http.error_handlers.response_models import ResponseAction
 from airbyte_cdk.sources.types import Config
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 
 ComponentDefinition = Mapping[str, Any]
 
@@ -3532,7 +3533,7 @@ def create_fixed_window_call_rate_policy(
         # Set the initial reset timestamp to 10 days from now.
         # This value will be updated by the first request.
         return FixedWindowCallRatePolicy(
-            next_reset_ts=datetime.datetime.now() + datetime.timedelta(days=10),
+            next_reset_ts=ab_datetime_now() + datetime.timedelta(days=10),
             period=parse_duration(model.period),
             call_limit=model.call_limit,
             matchers=matchers,

airbyte_cdk/sources/file_based/file_types/unstructured_parser.py

@@ -37,6 +37,7 @@
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 from airbyte_cdk.sources.file_based.schema_helpers import SchemaType
 from airbyte_cdk.utils import is_cloud_environment
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 unstructured_partition_pdf = None
@@ -300,7 +301,7 @@ def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[st
                 format_config.processing,
                 FileType.MD,
                 "auto",
-                RemoteFile(uri="test", last_modified=datetime.now()),
+                RemoteFile(uri="test", last_modified=ab_datetime_now()),
             )
         except Exception:
             return False, "".join(traceback.format_exc())

airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py

@@ -20,7 +20,11 @@
 from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.types import Record
-from airbyte_cdk.utils.datetime_helpers import ab_datetime_format, ab_datetime_parse
+from airbyte_cdk.utils.datetime_helpers import (
+    ab_datetime_format,
+    ab_datetime_now,
+    ab_datetime_parse,
+)
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition
@@ -310,7 +314,7 @@ def _compute_start_time(self) -> datetime:
                 disallow_other_formats=False,
             )
             if self._is_history_full():
-                time_window = datetime.now() - self._time_window_if_history_is_full
+                time_window = ab_datetime_now() - self._time_window_if_history_is_full
                 earliest_dt = min(earliest_dt, time_window)
             return earliest_dt
 

airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py

@@ -12,7 +12,11 @@
     AbstractFileBasedCursor,
 )
 from airbyte_cdk.sources.file_based.types import StreamState
-from airbyte_cdk.utils.datetime_helpers import ab_datetime_format, ab_datetime_parse
+from airbyte_cdk.utils.datetime_helpers import (
+    ab_datetime_format,
+    ab_datetime_now,
+    ab_datetime_parse,
+)
 
 
 class DefaultFileBasedCursor(AbstractFileBasedCursor):
@@ -156,6 +160,6 @@ def _compute_start_time(self) -> datetime:
                 disallow_other_formats=False,
             )
             if self._is_history_full():
-                time_window = datetime.now() - self._time_window_if_history_is_full
+                time_window = ab_datetime_now() - self._time_window_if_history_is_full
                 earliest_dt = min(earliest_dt, time_window)
             return earliest_dt

airbyte_cdk/sources/streams/call_rate.py

@@ -19,7 +19,7 @@
 from pyrate_limiter import Rate as PyRateRate
 from pyrate_limiter.exceptions import BucketFullException
 
-from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime
+from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now
 
 # prevents mypy from complaining about missing session attributes in LimiterMixin
 if TYPE_CHECKING:
@@ -282,7 +282,7 @@ class UnlimitedCallRatePolicy(BaseCallRatePolicy):
             ),
             FixedWindowCallRatePolicy(
                 matchers=[HttpRequestMatcher(url="/some/method")],
-                next_reset_ts=datetime.now(),
+                next_reset_ts=ab_datetime_now(),
                 period=timedelta(hours=1)
                 call_limit=1000,
             ),
@@ -387,7 +387,7 @@ def update(
                 self._next_reset_ts = call_reset_ts
 
     def _update_current_window(self) -> None:
-        now = datetime.datetime.now()
+        now = ab_datetime_now()
         if now > self._next_reset_ts:
             logger.debug("started new window, %s calls available now", self._call_limit)
             self._next_reset_ts = self._next_reset_ts + self._offset

airbyte_cdk/utils/datetime_helpers.py

@@ -81,8 +81,9 @@
 ```
 """
 
+import decimal
 from datetime import datetime, timedelta, timezone
-from typing import Any, Optional, Union, overload
+from typing import Any, Optional, Union, cast, overload
 
 from dateutil import parser
 from typing_extensions import Never
@@ -360,8 +361,8 @@ def ab_datetime_now() -> AirbyteDateTime:
 
 
 def ab_datetime_parse(
-    dt_str: str | int,
-    formats: list[str | None] | None = None,
+    dt_str: str | int | float,
+    formats: list[str | None] | list[str] | None = None,
     disallow_other_formats: bool = False,
 ) -> AirbyteDateTime:
     """Parses a datetime string or timestamp into an AirbyteDateTime with timezone awareness.
@@ -404,6 +405,27 @@ def ab_datetime_parse(
         # Remove None values from formats list, and coalesce to None if empty
         formats = [f for f in formats or [] if f] or None
 
+        if isinstance(dt_str, str):
+            if dt_str.startswith("-"):
+                raise ValueError("Timestamp cannot be negative: " + dt_str)
+
+            if dt_str[1:].replace(".", "").isdigit():
+                # Handle floats and ints as strings
+                if "." in dt_str:
+                    dt_str = float(dt_str)
+                else:
+                    dt_str = int(dt_str)
+
+        if isinstance(dt_str, float):
+            # Handle float values as Unix timestamps (UTC)
+            if dt_str < 0:
+                raise ValueError("Timestamp cannot be negative")
+            if len(str(abs(int(dt_str)))) > 10:
+                raise ValueError("Timestamp value too large")
+
+            instant = Instant.from_timestamp(dt_str)
+            return AirbyteDateTime.from_datetime(instant.py_datetime())
+
         # Handle numeric values as Unix timestamps (UTC)
         if isinstance(dt_str, int) or (
             isinstance(dt_str, str)

airbyte_cdk/utils/stream_status_utils.py

@@ -17,6 +17,7 @@
     TraceType,
 )
 from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 
 
 def as_airbyte_message(
@@ -28,7 +29,7 @@ def as_airbyte_message(
     Builds an AirbyteStreamStatusTraceMessage for the provided stream
     """
 
-    now_millis = datetime.now().timestamp() * 1000.0
+    now_millis = ab_datetime_now().to_epoch_millis()
 
     trace_message = AirbyteTraceMessage(
         type=TraceType.STREAM_STATUS,