avtomatika 1.0b4__py3-none-any.whl → 1.0b6__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,15 +1,18 @@
1
1
  from abc import ABC
2
+ from contextlib import suppress
3
+ from datetime import datetime
2
4
  from logging import getLogger
3
- from typing import Any, Dict, List
5
+ from typing import Any
4
6
  from uuid import uuid4
7
+ from zoneinfo import ZoneInfo
5
8
 
6
- from asyncpg import Pool, PostgresError, create_pool # type: ignore[import-untyped]
9
+ from asyncpg import Connection, Pool, PostgresError, create_pool # type: ignore[import-untyped]
10
+ from orjson import dumps, loads
7
11
 
8
12
  from .base import HistoryStorageBase
9
13
 
10
14
  logger = getLogger(__name__)
11
15
 
12
- # SQL queries to create tables, adapted for PostgreSQL
13
16
  CREATE_JOB_HISTORY_TABLE_PG = """
14
17
  CREATE TABLE IF NOT EXISTS job_history (
15
18
  event_id UUID PRIMARY KEY,
@@ -42,14 +45,24 @@ CREATE_JOB_ID_INDEX_PG = "CREATE INDEX IF NOT EXISTS idx_job_id ON job_history(j
42
45
  class PostgresHistoryStorage(HistoryStorageBase, ABC):
43
46
  """Implementation of the history store based on asyncpg for PostgreSQL."""
44
47
 
45
- def __init__(self, dsn: str):
48
+ def __init__(self, dsn: str, tz_name: str = "UTC"):
46
49
  self._dsn = dsn
47
50
  self._pool: Pool | None = None
51
+ self.tz_name = tz_name
52
+ self.tz = ZoneInfo(tz_name)
53
+
54
+ async def _setup_connection(self, conn: Connection):
55
+ """Configures the connection session with the correct timezone."""
56
+ try:
57
+ await conn.execute(f"SET TIME ZONE '{self.tz_name}'")
58
+ except PostgresError as e:
59
+ logger.error(f"Failed to set timezone '{self.tz_name}' for PG connection: {e}")
48
60
 
49
61
  async def initialize(self):
50
62
  """Initializes the connection pool to PostgreSQL and creates tables."""
51
63
  try:
52
- self._pool = await create_pool(dsn=self._dsn)
64
+ # We use init parameter to configure each new connection in the pool
65
+ self._pool = await create_pool(dsn=self._dsn, init=self._setup_connection)
53
66
  if not self._pool:
54
67
  raise RuntimeError("Failed to create a connection pool.")
55
68
 
@@ -57,7 +70,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
57
70
  await conn.execute(CREATE_JOB_HISTORY_TABLE_PG)
58
71
  await conn.execute(CREATE_WORKER_HISTORY_TABLE_PG)
59
72
  await conn.execute(CREATE_JOB_ID_INDEX_PG)
60
- logger.info("PostgreSQL history storage initialized.")
73
+ logger.info(f"PostgreSQL history storage initialized (TZ={self.tz_name}).")
61
74
  except (PostgresError, OSError) as e:
62
75
  logger.error(f"Failed to initialize PostgreSQL history storage: {e}")
63
76
  raise
@@ -68,21 +81,27 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
68
81
  await self._pool.close()
69
82
  logger.info("PostgreSQL history storage connection pool closed.")
70
83
 
71
- async def log_job_event(self, event_data: Dict[str, Any]):
84
+ async def log_job_event(self, event_data: dict[str, Any]):
72
85
  """Logs a job lifecycle event to PostgreSQL."""
73
86
  if not self._pool:
74
87
  raise RuntimeError("History storage is not initialized.")
75
88
 
76
89
  query = """
77
90
  INSERT INTO job_history (
78
- event_id, job_id, state, event_type, duration_ms,
91
+ event_id, job_id, timestamp, state, event_type, duration_ms,
79
92
  previous_state, next_state, worker_id, attempt_number,
80
93
  context_snapshot
81
- ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
94
+ ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
82
95
  """
96
+ now = datetime.now(self.tz)
97
+
98
+ context_snapshot = event_data.get("context_snapshot")
99
+ context_snapshot_json = dumps(context_snapshot).decode("utf-8") if context_snapshot else None
100
+
83
101
  params = (
84
102
  uuid4(),
85
103
  event_data.get("job_id"),
104
+ now,
86
105
  event_data.get("state"),
87
106
  event_data.get("event_type"),
88
107
  event_data.get("duration_ms"),
@@ -90,7 +109,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
90
109
  event_data.get("next_state"),
91
110
  event_data.get("worker_id"),
92
111
  event_data.get("attempt_number"),
93
- event_data.get("context_snapshot"),
112
+ context_snapshot_json,
94
113
  )
95
114
  try:
96
115
  async with self._pool.acquire() as conn:
@@ -98,21 +117,27 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
98
117
  except PostgresError as e:
99
118
  logger.error(f"Failed to log job event to PostgreSQL: {e}")
100
119
 
101
- async def log_worker_event(self, event_data: Dict[str, Any]):
120
+ async def log_worker_event(self, event_data: dict[str, Any]):
102
121
  """Logs a worker lifecycle event to PostgreSQL."""
103
122
  if not self._pool:
104
123
  raise RuntimeError("History storage is not initialized.")
105
124
 
106
125
  query = """
107
126
  INSERT INTO worker_history (
108
- event_id, worker_id, event_type, worker_info_snapshot
109
- ) VALUES ($1, $2, $3, $4)
127
+ event_id, worker_id, timestamp, event_type, worker_info_snapshot
128
+ ) VALUES ($1, $2, $3, $4, $5)
110
129
  """
130
+ now = datetime.now(self.tz)
131
+
132
+ worker_info = event_data.get("worker_info_snapshot")
133
+ worker_info_json = dumps(worker_info).decode("utf-8") if worker_info else None
134
+
111
135
  params = (
112
136
  uuid4(),
113
137
  event_data.get("worker_id"),
138
+ now,
114
139
  event_data.get("event_type"),
115
- event_data.get("worker_info_snapshot"),
140
+ worker_info_json,
116
141
  )
117
142
  try:
118
143
  async with self._pool.acquire() as conn:
@@ -120,7 +145,24 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
120
145
  except PostgresError as e:
121
146
  logger.error(f"Failed to log worker event to PostgreSQL: {e}")
122
147
 
123
- async def get_job_history(self, job_id: str) -> List[Dict[str, Any]]:
148
+ def _format_row(self, row: dict[str, Any]) -> dict[str, Any]:
149
+ """Helper to format a row from DB: convert timestamp to local TZ and decode JSON."""
150
+ item = dict(row)
151
+
152
+ if isinstance(item.get("context_snapshot"), str):
153
+ with suppress(Exception):
154
+ item["context_snapshot"] = loads(item["context_snapshot"])
155
+
156
+ if isinstance(item.get("worker_info_snapshot"), str):
157
+ with suppress(Exception):
158
+ item["worker_info_snapshot"] = loads(item["worker_info_snapshot"])
159
+
160
+ if "timestamp" in item and isinstance(item["timestamp"], datetime):
161
+ item["timestamp"] = item["timestamp"].astimezone(self.tz)
162
+
163
+ return item
164
+
165
+ async def get_job_history(self, job_id: str) -> list[dict[str, Any]]:
124
166
  """Gets the full history for the specified job from PostgreSQL."""
125
167
  if not self._pool:
126
168
  raise RuntimeError("History storage is not initialized.")
@@ -129,15 +171,14 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
129
171
  try:
130
172
  async with self._pool.acquire() as conn:
131
173
  rows = await conn.fetch(query, job_id)
132
- # asyncpg.Record can be easily converted to a dict
133
- return [dict(row) for row in rows]
174
+ return [self._format_row(row) for row in rows]
134
175
  except PostgresError as e:
135
176
  logger.error(
136
177
  f"Failed to get job history for job_id {job_id} from PostgreSQL: {e}",
137
178
  )
138
179
  return []
139
180
 
140
- async def get_jobs(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
181
+ async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
141
182
  if not self._pool:
142
183
  raise RuntimeError("History storage is not initialized.")
143
184
 
@@ -156,12 +197,12 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
156
197
  try:
157
198
  async with self._pool.acquire() as conn:
158
199
  rows = await conn.fetch(query, limit, offset)
159
- return [dict(row) for row in rows]
200
+ return [self._format_row(row) for row in rows]
160
201
  except PostgresError as e:
161
202
  logger.error(f"Failed to get jobs list from PostgreSQL: {e}")
162
203
  return []
163
204
 
164
- async def get_job_summary(self) -> Dict[str, int]:
205
+ async def get_job_summary(self) -> dict[str, int]:
165
206
  if not self._pool:
166
207
  raise RuntimeError("History storage is not initialized.")
167
208
 
@@ -195,7 +236,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
195
236
  self,
196
237
  worker_id: str,
197
238
  since_days: int,
198
- ) -> List[Dict[str, Any]]:
239
+ ) -> list[dict[str, Any]]:
199
240
  if not self._pool:
200
241
  raise RuntimeError("History storage is not initialized.")
201
242
 
@@ -208,7 +249,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
208
249
  try:
209
250
  async with self._pool.acquire() as conn:
210
251
  rows = await conn.fetch(query, worker_id, since_days)
211
- return [dict(row) for row in rows]
252
+ return [self._format_row(row) for row in rows]
212
253
  except PostgresError as e:
213
254
  logger.error(f"Failed to get worker history for worker_id {worker_id} from PostgreSQL: {e}")
214
255
  return []
@@ -1,6 +1,10 @@
1
+ from contextlib import suppress
2
+ from datetime import datetime
1
3
  from logging import getLogger
2
- from typing import Any, Dict, List
4
+ from time import time
5
+ from typing import Any
3
6
  from uuid import uuid4
7
+ from zoneinfo import ZoneInfo
4
8
 
5
9
  from aiosqlite import Connection, Error, Row, connect
6
10
  from orjson import dumps, loads
@@ -9,12 +13,11 @@ from .base import HistoryStorageBase
9
13
 
10
14
  logger = getLogger(__name__)
11
15
 
12
- # SQL queries for creating tables
13
16
  CREATE_JOB_HISTORY_TABLE = """
14
17
  CREATE TABLE IF NOT EXISTS job_history (
15
18
  event_id TEXT PRIMARY KEY,
16
19
  job_id TEXT NOT NULL,
17
- timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
20
+ timestamp REAL NOT NULL,
18
21
  state TEXT,
19
22
  event_type TEXT NOT NULL,
20
23
  duration_ms INTEGER,
@@ -30,7 +33,7 @@ CREATE_WORKER_HISTORY_TABLE = """
30
33
  CREATE TABLE IF NOT EXISTS worker_history (
31
34
  event_id TEXT PRIMARY KEY,
32
35
  worker_id TEXT NOT NULL,
33
- timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
36
+ timestamp REAL NOT NULL,
34
37
  event_type TEXT NOT NULL,
35
38
  worker_info_snapshot TEXT
36
39
  );
@@ -40,11 +43,15 @@ CREATE_JOB_ID_INDEX = "CREATE INDEX IF NOT EXISTS idx_job_id ON job_history(job_
40
43
 
41
44
 
42
45
  class SQLiteHistoryStorage(HistoryStorageBase):
43
- """Implementation of the history store based on aiosqlite."""
46
+ """Implementation of the history store based on aiosqlite.
47
+ Stores timestamps as Unix time (UTC) for correct sorting,
48
+ and converts them to the configured timezone upon retrieval.
49
+ """
44
50
 
45
- def __init__(self, db_path: str):
51
+ def __init__(self, db_path: str, tz_name: str = "UTC"):
46
52
  self._db_path = db_path
47
53
  self._conn: Connection | None = None
54
+ self.tz = ZoneInfo(tz_name)
48
55
 
49
56
  async def initialize(self):
50
57
  """Initializes the database connection and creates tables if they don't exist."""
@@ -67,22 +74,44 @@ class SQLiteHistoryStorage(HistoryStorageBase):
67
74
  await self._conn.close()
68
75
  logger.info("SQLite history storage connection closed.")
69
76
 
70
- async def log_job_event(self, event_data: Dict[str, Any]):
77
+ def _format_row(self, row: dict[str, Any]) -> dict[str, Any]:
78
+ """Helper to format a row from DB: decode JSON and convert timestamp."""
79
+ item = dict(row)
80
+
81
+ if item.get("context_snapshot"):
82
+ with suppress(Exception):
83
+ item["context_snapshot"] = loads(item["context_snapshot"])
84
+
85
+ if item.get("worker_info_snapshot"):
86
+ with suppress(Exception):
87
+ item["worker_info_snapshot"] = loads(item["worker_info_snapshot"])
88
+
89
+ if "timestamp" in item and isinstance(item["timestamp"], (int, float)):
90
+ item["timestamp"] = datetime.fromtimestamp(item["timestamp"], self.tz)
91
+
92
+ return item
93
+
94
+ async def log_job_event(self, event_data: dict[str, Any]):
71
95
  """Logs a job lifecycle event to the job_history table."""
72
96
  if not self._conn:
73
97
  raise RuntimeError("History storage is not initialized.")
74
98
 
75
99
  query = """
76
100
  INSERT INTO job_history (
77
- event_id, job_id, state, event_type, duration_ms,
101
+ event_id, job_id, timestamp, state, event_type, duration_ms,
78
102
  previous_state, next_state, worker_id, attempt_number,
79
103
  context_snapshot
80
- ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
104
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
81
105
  """
82
- # Ensure all keys are present to avoid errors
106
+ now_ts = time()
107
+
108
+ context_snapshot = event_data.get("context_snapshot")
109
+ context_snapshot_json = dumps(context_snapshot).decode("utf-8") if context_snapshot else None
110
+
83
111
  params = (
84
112
  str(uuid4()),
85
113
  event_data.get("job_id"),
114
+ now_ts,
86
115
  event_data.get("state"),
87
116
  event_data.get("event_type"),
88
117
  event_data.get("duration_ms"),
@@ -90,7 +119,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
90
119
  event_data.get("next_state"),
91
120
  event_data.get("worker_id"),
92
121
  event_data.get("attempt_number"),
93
- (dumps(event_data.get("context_snapshot")) if event_data.get("context_snapshot") else None),
122
+ context_snapshot_json,
94
123
  )
95
124
 
96
125
  try:
@@ -98,23 +127,28 @@ class SQLiteHistoryStorage(HistoryStorageBase):
98
127
  await self._conn.commit()
99
128
  except Error as e:
100
129
  logger.error(f"Failed to log job event: {e}")
101
- # According to "Option A", we log the error but do not interrupt execution
102
130
 
103
- async def log_worker_event(self, event_data: Dict[str, Any]):
131
+ async def log_worker_event(self, event_data: dict[str, Any]):
104
132
  """Logs a worker lifecycle event to the worker_history table."""
105
133
  if not self._conn:
106
134
  raise RuntimeError("History storage is not initialized.")
107
135
 
108
136
  query = """
109
137
  INSERT INTO worker_history (
110
- event_id, worker_id, event_type, worker_info_snapshot
111
- ) VALUES (?, ?, ?, ?)
138
+ event_id, worker_id, timestamp, event_type, worker_info_snapshot
139
+ ) VALUES (?, ?, ?, ?, ?)
112
140
  """
141
+ now_ts = time()
142
+
143
+ worker_info = event_data.get("worker_info_snapshot")
144
+ worker_info_json = dumps(worker_info).decode("utf-8") if worker_info else None
145
+
113
146
  params = (
114
147
  str(uuid4()),
115
148
  event_data.get("worker_id"),
149
+ now_ts,
116
150
  event_data.get("event_type"),
117
- (dumps(event_data.get("worker_info_snapshot")) if event_data.get("worker_info_snapshot") else None),
151
+ worker_info_json,
118
152
  )
119
153
 
120
154
  try:
@@ -123,7 +157,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
123
157
  except Error as e:
124
158
  logger.error(f"Failed to log worker event: {e}")
125
159
 
126
- async def get_job_history(self, job_id: str) -> List[Dict[str, Any]]:
160
+ async def get_job_history(self, job_id: str) -> list[dict[str, Any]]:
127
161
  """Gets the full history for the specified job, sorted by time."""
128
162
  if not self._conn:
129
163
  raise RuntimeError("History storage is not initialized.")
@@ -133,19 +167,12 @@ class SQLiteHistoryStorage(HistoryStorageBase):
133
167
  self._conn.row_factory = Row
134
168
  async with self._conn.execute(query, (job_id,)) as cursor:
135
169
  rows = await cursor.fetchall()
136
- history = []
137
- for row in rows:
138
- item = dict(row)
139
- # Deserialize the JSON string back into a dict
140
- if item.get("context_snapshot"):
141
- item["context_snapshot"] = loads(item["context_snapshot"])
142
- history.append(item)
143
- return history
170
+ return [self._format_row(row) for row in rows]
144
171
  except Error as e:
145
172
  logger.error(f"Failed to get job history for job_id {job_id}: {e}")
146
173
  return []
147
174
 
148
- async def get_jobs(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
175
+ async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
149
176
  """Gets a list of the latest unique jobs with pagination."""
150
177
  if not self._conn:
151
178
  raise RuntimeError("History storage is not initialized.")
@@ -166,18 +193,12 @@ class SQLiteHistoryStorage(HistoryStorageBase):
166
193
  self._conn.row_factory = Row
167
194
  async with self._conn.execute(query, (limit, offset)) as cursor:
168
195
  rows = await cursor.fetchall()
169
- jobs = []
170
- for row in rows:
171
- item = dict(row)
172
- if item.get("context_snapshot"):
173
- item["context_snapshot"] = loads(item["context_snapshot"])
174
- jobs.append(item)
175
- return jobs
196
+ return [self._format_row(row) for row in rows]
176
197
  except Error as e:
177
198
  logger.error(f"Failed to get jobs list: {e}")
178
199
  return []
179
200
 
180
- async def get_job_summary(self) -> Dict[str, int]:
201
+ async def get_job_summary(self) -> dict[str, int]:
181
202
  """Returns a summary of job statuses."""
182
203
  if not self._conn:
183
204
  raise RuntimeError("History storage is not initialized.")
@@ -213,27 +234,23 @@ class SQLiteHistoryStorage(HistoryStorageBase):
213
234
  self,
214
235
  worker_id: str,
215
236
  since_days: int,
216
- ) -> List[Dict[str, Any]]:
237
+ ) -> list[dict[str, Any]]:
217
238
  if not self._conn:
218
239
  raise RuntimeError("History storage is not initialized.")
219
240
 
241
+ threshold_ts = time() - (since_days * 86400)
242
+
220
243
  query = """
221
244
  SELECT * FROM job_history
222
245
  WHERE worker_id = ?
223
- AND timestamp >= date('now', '-' || ? || ' days')
246
+ AND timestamp >= ?
224
247
  ORDER BY timestamp DESC
225
248
  """
226
249
  try:
227
250
  self._conn.row_factory = Row
228
- async with self._conn.execute(query, (worker_id, since_days)) as cursor:
251
+ async with self._conn.execute(query, (worker_id, threshold_ts)) as cursor:
229
252
  rows = await cursor.fetchall()
230
- history = []
231
- for row in rows:
232
- item = dict(row)
233
- if item.get("context_snapshot"):
234
- item["context_snapshot"] = loads(item["context_snapshot"])
235
- history.append(item)
236
- return history
253
+ return [self._format_row(row) for row in rows]
237
254
  except Error as e:
238
255
  logger.error(f"Failed to get worker history for worker_id {worker_id}: {e}")
239
256
  return []
@@ -1,25 +1,67 @@
1
+ from datetime import datetime
1
2
  from logging import DEBUG, Formatter, StreamHandler, getLogger
2
3
  from sys import stdout
4
+ from zoneinfo import ZoneInfo
3
5
 
4
6
  from pythonjsonlogger import json
5
7
 
6
8
 
7
- def setup_logging(log_level: str = "INFO", log_format: str = "json"):
9
+ class TimezoneFormatter(Formatter):
10
+ """Formatter that respects a custom timezone."""
11
+
12
+ def __init__(self, fmt=None, datefmt=None, style="%", validate=True, *, tz_name="UTC"):
13
+ super().__init__(fmt, datefmt, style, validate)
14
+ self.tz = ZoneInfo(tz_name)
15
+
16
+ def converter(self, timestamp):
17
+ return datetime.fromtimestamp(timestamp, self.tz)
18
+
19
+ def formatTime(self, record, datefmt=None):
20
+ dt = self.converter(record.created)
21
+ if datefmt:
22
+ s = dt.strftime(datefmt)
23
+ else:
24
+ try:
25
+ s = dt.isoformat(timespec="milliseconds")
26
+ except TypeError:
27
+ s = dt.isoformat()
28
+ return s
29
+
30
+
31
+ class TimezoneJsonFormatter(json.JsonFormatter):
32
+ """JSON Formatter that respects a custom timezone."""
33
+
34
+ def __init__(self, *args, tz_name="UTC", **kwargs):
35
+ super().__init__(*args, **kwargs)
36
+ self.tz = ZoneInfo(tz_name)
37
+
38
+ def formatTime(self, record, datefmt=None):
39
+ # Override formatTime to use timezone-aware datetime
40
+ dt = datetime.fromtimestamp(record.created, self.tz)
41
+ if datefmt:
42
+ return dt.strftime(datefmt)
43
+ # Return ISO format with offset
44
+ return dt.isoformat()
45
+
46
+
47
+ def setup_logging(log_level: str = "INFO", log_format: str = "json", tz_name: str = "UTC"):
8
48
  """Configures structured logging for the entire application."""
9
- logger = getLogger("orchestrator")
49
+ logger = getLogger("avtomatika")
10
50
  logger.setLevel(log_level)
11
51
 
12
52
  handler = StreamHandler(stdout)
13
53
  formatter: Formatter
14
54
  if log_format.lower() == "json":
15
55
  # Formatter for JSON logs
16
- formatter = json.JsonFormatter(
56
+ formatter = TimezoneJsonFormatter(
17
57
  "%(asctime)s %(name)s %(levelname)s %(message)s %(pathname)s %(lineno)d",
58
+ tz_name=tz_name,
18
59
  )
19
60
  else:
20
61
  # Standard text formatter
21
- formatter = Formatter(
62
+ formatter = TimezoneFormatter(
22
63
  "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
64
+ tz_name=tz_name,
23
65
  )
24
66
 
25
67
  handler.setFormatter(formatter)
@@ -30,12 +72,21 @@ def setup_logging(log_level: str = "INFO", log_format: str = "json"):
30
72
 
31
73
  # Configure the root logger to see logs from libraries (aiohttp, etc.)
32
74
  root_logger = getLogger()
33
- # Set the root logger level so as not to filter messages
34
- # for child loggers prematurely.
35
75
  root_logger.setLevel(DEBUG)
76
+
36
77
  if not root_logger.handlers:
37
78
  root_handler = StreamHandler(stdout)
38
- root_handler.setFormatter(
39
- Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"),
79
+ root_formatter = TimezoneFormatter(
80
+ "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
81
+ tz_name=tz_name,
40
82
  )
83
+ root_handler.setFormatter(root_formatter)
41
84
  root_logger.addHandler(root_handler)
85
+ else:
86
+ for h in root_logger.handlers:
87
+ h.setFormatter(
88
+ TimezoneFormatter(
89
+ "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
90
+ tz_name=tz_name,
91
+ )
92
+ )
@@ -0,0 +1,119 @@
1
+ from asyncio import CancelledError, sleep
2
+ from datetime import datetime
3
+ from logging import getLogger
4
+ from typing import TYPE_CHECKING
5
+ from zoneinfo import ZoneInfo
6
+
7
+ from .scheduler_config_loader import ScheduledJobConfig, load_schedules_from_file
8
+
9
+ if TYPE_CHECKING:
10
+ from .engine import OrchestratorEngine
11
+
12
+ logger = getLogger(__name__)
13
+
14
+
15
+ class Scheduler:
16
+ def __init__(self, engine: "OrchestratorEngine"):
17
+ self.engine = engine
18
+ self.config = engine.config
19
+ self.storage = engine.storage
20
+ self._running = False
21
+ self.schedules: list[ScheduledJobConfig] = []
22
+ self.timezone = ZoneInfo(self.config.TZ)
23
+
24
+ def load_config(self):
25
+ if not self.config.SCHEDULES_CONFIG_PATH:
26
+ logger.info("No SCHEDULES_CONFIG_PATH set. Scheduler will not run any jobs.")
27
+ return
28
+
29
+ try:
30
+ self.schedules = load_schedules_from_file(self.config.SCHEDULES_CONFIG_PATH)
31
+ logger.info(f"Loaded {len(self.schedules)} scheduled jobs.")
32
+ except Exception as e:
33
+ logger.error(f"Failed to load schedules config: {e}")
34
+
35
+ async def run(self):
36
+ self.load_config()
37
+ if not self.schedules:
38
+ logger.info("No schedules found. Scheduler loop will not start.")
39
+ return
40
+
41
+ logger.info("Scheduler started.")
42
+ self._running = True
43
+
44
+ while self._running:
45
+ try:
46
+ now_utc = datetime.now(ZoneInfo("UTC"))
47
+ now_tz = now_utc.astimezone(self.timezone)
48
+
49
+ for job in self.schedules:
50
+ await self._process_job(job, now_tz)
51
+
52
+ await sleep(1)
53
+
54
+ except CancelledError:
55
+ break
56
+ except Exception as e:
57
+ logger.error(f"Error in scheduler loop: {e}", exc_info=True)
58
+ await sleep(30)
59
+
60
+ logger.info("Scheduler stopped.")
61
+
62
+ def stop(self):
63
+ self._running = False
64
+
65
+ async def _process_job(self, job: ScheduledJobConfig, now_tz: datetime):
66
+ if job.interval_seconds:
67
+ await self._process_interval_job(job, now_tz)
68
+ else:
69
+ await self._process_calendar_job(job, now_tz)
70
+
71
+ async def _process_interval_job(self, job: ScheduledJobConfig, now_tz: datetime):
72
+ last_run_key = f"scheduler:last_run:{job.name}"
73
+ last_run_ts = await self.storage.get_str(last_run_key)
74
+
75
+ now_ts = now_tz.timestamp()
76
+
77
+ if last_run_ts and now_ts - float(last_run_ts) < job.interval_seconds:
78
+ return
79
+
80
+ lock_key = f"scheduler:lock:interval:{job.name}"
81
+ if await self.storage.set_nx_ttl(lock_key, "locked", ttl=5):
82
+ try:
83
+ await self._trigger_job(job)
84
+ await self.storage.set_str(last_run_key, str(now_ts))
85
+ except Exception as e:
86
+ logger.error(f"Failed to trigger interval job {job.name}: {e}")
87
+
88
+ async def _process_calendar_job(self, job: ScheduledJobConfig, now_tz: datetime):
89
+ target_time_str = job.daily_at or job.time
90
+ if not target_time_str:
91
+ return
92
+
93
+ current_time_str = now_tz.strftime("%H:%M")
94
+
95
+ if current_time_str != target_time_str:
96
+ return
97
+
98
+ if job.weekly_days:
99
+ current_day_str = now_tz.strftime("%a").lower()
100
+ if current_day_str not in [d.lower() for d in job.weekly_days]:
101
+ return
102
+
103
+ if job.monthly_dates and now_tz.day not in job.monthly_dates:
104
+ return
105
+
106
+ date_str = now_tz.strftime("%Y-%m-%d")
107
+ lock_key = f"scheduler:lock:{job.name}:{date_str}"
108
+
109
+ if await self.storage.set_nx_ttl(lock_key, "locked", ttl=86400):
110
+ logger.info(f"Triggering scheduled job {job.name}")
111
+ await self._trigger_job(job)
112
+
113
+ async def _trigger_job(self, job: ScheduledJobConfig):
114
+ try:
115
+ await self.engine.create_background_job(
116
+ blueprint_name=job.blueprint, initial_data=job.input_data, source=f"scheduler:{job.name}"
117
+ )
118
+ except Exception as e:
119
+ logger.error(f"Failed to create background job {job.name}: {e}")