-
Notifications
You must be signed in to change notification settings - Fork 45
feat: add timezone support for CRON triggers (#438) #463
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 5 commits
db9e0a0
93071da
770eb4e
4b065e4
9d55322
421c4bf
077ca50
8142f52
b425b31
530b136
7594831
2437624
cacf060
e7fc2f7
9510f29
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,9 +8,13 @@ | |
| from pymongo import ReturnDocument | ||
| from pymongo.errors import DuplicateKeyError | ||
| from app.config.settings import get_settings | ||
| from zoneinfo import ZoneInfo | ||
| import croniter | ||
| import asyncio | ||
|
|
||
| # Cache UTC timezone at module level to avoid repeated instantiation | ||
| UTC = ZoneInfo("UTC") | ||
|
|
||
| logger = LogsManager().get_logger() | ||
|
|
||
| async def get_due_triggers(cron_time: datetime) -> DatabaseTriggers | None: | ||
|
|
@@ -42,15 +46,26 @@ async def mark_as_failed(trigger: DatabaseTriggers): | |
|
|
||
| async def create_next_triggers(trigger: DatabaseTriggers, cron_time: datetime): | ||
| assert trigger.expression is not None | ||
| iter = croniter.croniter(trigger.expression, trigger.trigger_time) | ||
|
|
||
| # Use the trigger's timezone, defaulting to UTC if not specified | ||
| tz = ZoneInfo(trigger.timezone or "UTC") | ||
|
|
||
| # Convert trigger_time to the specified timezone for croniter | ||
| trigger_time_tz = trigger.trigger_time.replace(tzinfo=UTC).astimezone(tz) | ||
| iter = croniter.croniter(trigger.expression, trigger_time_tz) | ||
|
|
||
|
Comment on lines
+55
to
61
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Back-compat: handle missing timezone attribute and avoid Tests and older records may not have - # Use the trigger's timezone, defaulting to UTC if not specified
- tz = ZoneInfo(trigger.timezone or "UTC")
+ # Resolve timezone safely for old records/tests without the attribute
+ tz_name = getattr(trigger, "timezone", None) or "UTC"
+ tz = ZoneInfo(tz_name)
- # Convert trigger_time to the specified timezone for croniter
- trigger_time_tz = trigger.trigger_time.replace(tzinfo=UTC).astimezone(tz)
- iter = croniter.croniter(trigger.expression, trigger_time_tz)
+ # Convert trigger_time to the specified timezone for croniter
+ trigger_time_tz = trigger.trigger_time.replace(tzinfo=UTC).astimezone(tz)
+ cron_iter = croniter.croniter(trigger.expression, trigger_time_tz)Also avoids shadowing built-in 🧰 Tools🪛 GitHub Actions: State Manager Unit Tests[error] 56-56: pytest failed with AttributeError: Mock object has no attribute 'timezone' during create_next_triggers; the test mocks for DatabaseTriggers without providing a timezone attribute on the trigger mock. Command: uv run pytest tests/ --cov=app --cov-report=xml --cov-report=term-missing --cov-report=html -v --junitxml=full-pytest-report.xml 🤖 Prompt for AI Agents |
||
| while True: | ||
| next_trigger_time = iter.get_next(datetime) | ||
| # Get next trigger time in the specified timezone | ||
| next_trigger_time_tz = iter.get_next(datetime) | ||
|
|
||
| # Convert back to UTC for storage | ||
| next_trigger_time = next_trigger_time_tz.astimezone(UTC).replace(tzinfo=None) | ||
|
|
||
| try: | ||
| await DatabaseTriggers( | ||
| type=TriggerTypeEnum.CRON, | ||
| expression=trigger.expression, | ||
| timezone=trigger.timezone, | ||
| graph_name=trigger.graph_name, | ||
| namespace=trigger.namespace, | ||
| trigger_time=next_trigger_time, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,6 +3,7 @@ | |
|
|
||
| from datetime import datetime | ||
| from json_schema_to_pydantic import create_model | ||
| from zoneinfo import ZoneInfo | ||
|
|
||
| from app.models.db.graph_template_model import GraphTemplate | ||
| from app.models.graph_template_validation_status import GraphTemplateValidationStatus | ||
|
|
@@ -11,6 +12,9 @@ | |
| from app.models.trigger_models import TriggerStatusEnum, TriggerTypeEnum | ||
| from app.models.db.trigger import DatabaseTriggers | ||
|
|
||
| # Cache UTC timezone at module level to avoid repeated instantiation | ||
| UTC = ZoneInfo("UTC") | ||
|
|
||
| logger = LogsManager().get_logger() | ||
|
|
||
| async def verify_node_exists(graph_template: GraphTemplate, registered_nodes: list[RegisteredNode]) -> list[str]: | ||
|
|
@@ -101,20 +105,37 @@ async def verify_inputs(graph_template: GraphTemplate, registered_nodes: list[Re | |
| return errors | ||
|
|
||
| async def create_crons(graph_template: GraphTemplate): | ||
| expressions_to_create = set([trigger.value["expression"] for trigger in graph_template.triggers if trigger.type == TriggerTypeEnum.CRON]) | ||
| # Build a map of (expression, timezone) -> validated CronTrigger for deduplication | ||
| triggers_to_create = {} | ||
| for trigger in graph_template.triggers: | ||
| if trigger.type == TriggerTypeEnum.CRON: | ||
| # Validate through CronTrigger model to normalize timezone (None -> "UTC") | ||
| from app.models.trigger_models import CronTrigger | ||
|
spa-raj marked this conversation as resolved.
Outdated
|
||
| cron_trigger = CronTrigger.model_validate(trigger.value) | ||
|
spa-raj marked this conversation as resolved.
Outdated
|
||
| triggers_to_create[(cron_trigger.expression, cron_trigger.timezone)] = cron_trigger | ||
|
|
||
| current_time = datetime.now(UTC).replace(tzinfo=None) | ||
|
|
||
| current_time = datetime.now() | ||
|
|
||
| new_db_triggers = [] | ||
| for expression in expressions_to_create: | ||
| iter = croniter.croniter(expression, current_time) | ||
| for (expression, timezone), cron_trigger in triggers_to_create.items(): | ||
| # Use the validated timezone (guaranteed to be valid IANA timezone, never None) | ||
| tz = ZoneInfo(timezone) | ||
|
|
||
| # Get current time in the specified timezone | ||
| current_time_tz = current_time.replace(tzinfo=UTC).astimezone(tz) | ||
| iter = croniter.croniter(expression, current_time_tz) | ||
|
|
||
| # Get next trigger time in the specified timezone | ||
| next_trigger_time_tz = iter.get_next(datetime) | ||
|
|
||
| # Convert back to UTC for storage (remove timezone info for storage) | ||
| next_trigger_time = next_trigger_time_tz.astimezone(UTC).replace(tzinfo=None) | ||
|
Comment on lines
+127
to
+135
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧹 Nitpick | 🔵 Trivial Name nit + consistency for Avoid shadowing - current_time_tz = current_time.replace(tzinfo=UTC).astimezone(tz)
- iter = croniter.croniter(expression, current_time_tz)
+ current_time_tz = current_time.replace(tzinfo=UTC).astimezone(tz)
+ cron_iter = croniter.croniter(expression, current_time_tz)
@@
- next_trigger_time_tz = iter.get_next(datetime)
+ next_trigger_time_tz = cron_iter.get_next(datetime)
@@
- next_trigger_time = next_trigger_time_tz.astimezone(UTC).replace(tzinfo=None)
- expires_at = next_trigger_time + timedelta(hours=settings.trigger_retention_hours)
+ next_trigger_time = next_trigger_time_tz.astimezone(UTC).replace(tzinfo=None)
+ expires_at = next_trigger_time.replace(tzinfo=timezone.utc) + timedelta(hours=settings.trigger_retention_hours)
🤖 Prompt for AI Agents |
||
|
|
||
| next_trigger_time = iter.get_next(datetime) | ||
|
|
||
| new_db_triggers.append( | ||
| DatabaseTriggers( | ||
| type=TriggerTypeEnum.CRON, | ||
| expression=expression, | ||
| expression=cron_trigger.expression, | ||
| timezone=cron_trigger.timezone, | ||
| graph_name=graph_template.name, | ||
| namespace=graph_template.namespace, | ||
| trigger_status=TriggerStatusEnum.PENDING, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,122 @@ | ||
| import pytest | ||
| from pydantic import ValidationError | ||
| from app.models.trigger_models import CronTrigger, Trigger, TriggerTypeEnum | ||
|
|
||
|
|
||
| class TestCronTrigger: | ||
| """Test cases for CronTrigger model""" | ||
|
|
||
| def test_valid_cron_trigger_with_timezone(self): | ||
| """Test creating a valid cron trigger with timezone""" | ||
| trigger = CronTrigger(expression="0 9 * * *", timezone="America/New_York") | ||
| assert trigger.expression == "0 9 * * *" | ||
| assert trigger.timezone == "America/New_York" | ||
|
|
||
| def test_valid_cron_trigger_without_timezone(self): | ||
| """Test creating a valid cron trigger without timezone defaults to UTC""" | ||
| trigger = CronTrigger(expression="0 9 * * *") | ||
| assert trigger.expression == "0 9 * * *" | ||
| assert trigger.timezone == "UTC" | ||
|
|
||
| def test_valid_cron_trigger_with_utc_timezone(self): | ||
| """Test creating a valid cron trigger with UTC timezone""" | ||
| trigger = CronTrigger(expression="0 9 * * *", timezone="UTC") | ||
| assert trigger.expression == "0 9 * * *" | ||
| assert trigger.timezone == "UTC" | ||
|
|
||
| def test_valid_cron_trigger_with_europe_london(self): | ||
| """Test creating a valid cron trigger with Europe/London timezone""" | ||
| trigger = CronTrigger(expression="0 17 * * *", timezone="Europe/London") | ||
| assert trigger.expression == "0 17 * * *" | ||
| assert trigger.timezone == "Europe/London" | ||
|
|
||
| def test_valid_cron_trigger_with_asia_tokyo(self): | ||
| """Test creating a valid cron trigger with Asia/Tokyo timezone""" | ||
| trigger = CronTrigger(expression="30 8 * * 1-5", timezone="Asia/Tokyo") | ||
| assert trigger.expression == "30 8 * * 1-5" | ||
| assert trigger.timezone == "Asia/Tokyo" | ||
|
|
||
| def test_invalid_cron_expression(self): | ||
| """Test creating a cron trigger with invalid expression""" | ||
| with pytest.raises(ValidationError) as exc_info: | ||
| CronTrigger(expression="invalid cron", timezone="UTC") | ||
|
|
||
| errors = exc_info.value.errors() | ||
| assert len(errors) == 1 | ||
| # Check the error message (Pydantic v2 doesn't always populate ctx) | ||
| error_msg = errors[0].get("msg") or str(errors[0]) | ||
| assert "Invalid cron expression" in error_msg | ||
|
|
||
| def test_invalid_timezone(self): | ||
| """Test creating a cron trigger with invalid timezone""" | ||
| with pytest.raises(ValidationError) as exc_info: | ||
| CronTrigger(expression="0 9 * * *", timezone="Invalid/Timezone") | ||
|
|
||
| errors = exc_info.value.errors() | ||
| assert len(errors) == 1 | ||
| # Check the error message (Pydantic v2 doesn't always populate ctx) | ||
| error_msg = errors[0].get("msg") or str(errors[0]) | ||
| assert "Invalid timezone" in error_msg | ||
| assert "Invalid/Timezone" in error_msg | ||
|
|
||
| def test_none_timezone_defaults_to_utc(self): | ||
| """Test that None timezone defaults to UTC""" | ||
| trigger = CronTrigger(expression="0 9 * * *", timezone=None) | ||
| assert trigger.timezone == "UTC" | ||
|
|
||
| def test_complex_cron_expression_with_timezone(self): | ||
| """Test complex cron expression with timezone""" | ||
| trigger = CronTrigger(expression="0 0 1,15 * *", timezone="America/Los_Angeles") | ||
| assert trigger.expression == "0 0 1,15 * *" | ||
| assert trigger.timezone == "America/Los_Angeles" | ||
|
|
||
| def test_every_15_minutes_cron_with_timezone(self): | ||
| """Test every 15 minutes cron with timezone""" | ||
| trigger = CronTrigger(expression="*/15 * * * *", timezone="Europe/Paris") | ||
| assert trigger.expression == "*/15 * * * *" | ||
| assert trigger.timezone == "Europe/Paris" | ||
|
|
||
|
|
||
| class TestTrigger: | ||
| """Test cases for Trigger model""" | ||
|
|
||
| def test_valid_trigger_with_cron_and_timezone(self): | ||
| """Test creating a valid trigger with CRON type and timezone""" | ||
| trigger = Trigger( | ||
| type=TriggerTypeEnum.CRON, | ||
| value={"expression": "0 9 * * *", "timezone": "America/New_York"} | ||
| ) | ||
| assert trigger.type == TriggerTypeEnum.CRON | ||
| assert trigger.value["expression"] == "0 9 * * *" | ||
| assert trigger.value["timezone"] == "America/New_York" | ||
|
|
||
| def test_valid_trigger_with_cron_without_timezone(self): | ||
| """Test creating a valid trigger with CRON type without timezone""" | ||
| trigger = Trigger( | ||
| type=TriggerTypeEnum.CRON, | ||
| value={"expression": "0 9 * * *"} | ||
| ) | ||
| assert trigger.type == TriggerTypeEnum.CRON | ||
| assert trigger.value["expression"] == "0 9 * * *" | ||
|
|
||
| def test_invalid_trigger_with_invalid_cron_expression(self): | ||
| """Test creating a trigger with invalid cron expression""" | ||
| with pytest.raises(ValidationError) as exc_info: | ||
| Trigger( | ||
| type=TriggerTypeEnum.CRON, | ||
| value={"expression": "invalid cron"} | ||
| ) | ||
|
|
||
| errors = exc_info.value.errors() | ||
| assert len(errors) > 0 | ||
|
|
||
| def test_invalid_trigger_with_invalid_timezone(self): | ||
| """Test creating a trigger with invalid timezone""" | ||
| with pytest.raises(ValidationError) as exc_info: | ||
| Trigger( | ||
| type=TriggerTypeEnum.CRON, | ||
| value={"expression": "0 9 * * *", "timezone": "Invalid/Zone"} | ||
| ) | ||
|
|
||
| errors = exc_info.value.errors() | ||
| assert len(errors) > 0 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we should add validation for timezone in SDK also, allowing to detect failures early. We can also take this as a separate PR and issue.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Feel free to create an issue if we are moving this to different PR.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure