From c703b8871f0eb7bc561c8600910e83a07cf32fe8 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Mon, 4 Dec 2023 18:37:19 +0530 Subject: [PATCH] helper to set tzinfo --- .../source_kobotoolbox/source.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/airbyte-integrations/connectors/source-kobotoolbox/source_kobotoolbox/source.py b/airbyte-integrations/connectors/source-kobotoolbox/source_kobotoolbox/source.py index a32d0b5746f8..692e4b6bcb98 100644 --- a/airbyte-integrations/connectors/source-kobotoolbox/source_kobotoolbox/source.py +++ b/airbyte-integrations/connectors/source-kobotoolbox/source_kobotoolbox/source.py @@ -95,6 +95,17 @@ def state(self, value: Mapping[str, Any]): """setter for state""" self._cursor_value = value[self.cursor_field] + def mk_tzaware_utc(self, dt): + """ + add a utc-tzinfo object to the dt if it doesn't have tzinfo + if it has a tzinfo, convert to utc + """ + from datetime import timezone + + if dt.tzinfo is None: + return dt.replace(tzinfo=timezone.utc) + return dt.astimezone(timezone.utc) + def mk_query(self): """query using endtime""" if self.cursor_field == "_submission_time": @@ -102,15 +113,8 @@ def mk_query(self): else: start_sub_time = datetime.fromisoformat(self.state[self.cursor_field]) start_sub_time -= timedelta(days=self.max_days_to_close) - from datetime import timezone - - tzaware_start_time = datetime.fromisoformat(self.start_time) - if tzaware_start_time.tzinfo is None: - # interpret as utc - tzaware_start_time = tzaware_start_time.replace(tzinfo=timezone.utc) - else: - # convert to utc if necessary - tzaware_start_time = tzaware_start_time.astimezone(timezone.utc) + start_sub_time = self.mk_tzaware_utc(start_sub_time) + tzaware_start_time = self.mk_tzaware_utc(datetime.fromisoformat(self.start_time)) start_sub_time = max(start_sub_time, tzaware_start_time) return {"_submission_time": {"$gte": start_sub_time.isoformat()}, self.cursor_field: {"$gte": self.state[self.cursor_field]}} @@ -168,10 +172,8 @@ def parse_response(self, response: requests.Response, **kwargs) -> Iterable[Mapp retval["_submission_time"] = record["_submission_time"] retval["endtime"] = record.get("endtime") if retval["endtime"]: - from datetime import timezone - # endtime is in utc - endtime = datetime.fromisoformat(retval["endtime"]).astimezone(timezone.utc) + endtime = self.mk_tzaware_utc(datetime.fromisoformat(retval["endtime"])) retval["endtime"] = endtime.isoformat() yield retval