diracx-db 0.0.1a46__py3-none-any.whl → 0.0.1a48__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.
@@ -9,13 +9,19 @@ class JobParametersDB(BaseOSDB):
9
9
  fields = {
10
10
  "JobID": {"type": "long"},
11
11
  "timestamp": {"type": "date"},
12
+ "PilotAgent": {"type": "keyword"},
13
+ "Pilot_Reference": {"type": "keyword"},
14
+ "JobGroup": {"type": "keyword"},
12
15
  "CPUNormalizationFactor": {"type": "long"},
13
16
  "NormCPUTime(s)": {"type": "long"},
14
- "Memory(kB)": {"type": "long"},
17
+ "Memory(MB)": {"type": "long"},
18
+ "LocalAccount": {"type": "keyword"},
15
19
  "TotalCPUTime(s)": {"type": "long"},
16
- "MemoryUsed(kb)": {"type": "long"},
17
- "HostName": {"type": "keyword"},
20
+ "PayloadPID": {"type": "long"},
21
+ "HostName": {"type": "text"},
18
22
  "GridCE": {"type": "keyword"},
23
+ "CEQueue": {"type": "keyword"},
24
+ "BatchSystem": {"type": "keyword"},
19
25
  "ModelName": {"type": "keyword"},
20
26
  "Status": {"type": "keyword"},
21
27
  "JobType": {"type": "keyword"},
diracx/db/sql/dummy/db.py CHANGED
@@ -1,9 +1,9 @@
1
1
  from __future__ import annotations
2
2
 
3
- from sqlalchemy import func, insert, select
3
+ from sqlalchemy import insert
4
4
  from uuid_utils import UUID
5
5
 
6
- from diracx.db.sql.utils import BaseSQLDB, apply_search_filters
6
+ from diracx.db.sql.utils import BaseSQLDB
7
7
 
8
8
  from .schema import Base as DummyDBBase
9
9
  from .schema import Cars, Owners
@@ -22,18 +22,7 @@ class DummyDB(BaseSQLDB):
22
22
  metadata = DummyDBBase.metadata
23
23
 
24
24
  async def summary(self, group_by, search) -> list[dict[str, str | int]]:
25
- columns = [Cars.__table__.columns[x] for x in group_by]
26
-
27
- stmt = select(*columns, func.count(Cars.license_plate).label("count"))
28
- stmt = apply_search_filters(Cars.__table__.columns.__getitem__, stmt, search)
29
- stmt = stmt.group_by(*columns)
30
-
31
- # Execute the query
32
- return [
33
- dict(row._mapping)
34
- async for row in (await self.conn.stream(stmt))
35
- if row.count > 0 # type: ignore
36
- ]
25
+ return await self._summary(Cars, group_by, search)
37
26
 
38
27
  async def insert_owner(self, name: str) -> int:
39
28
  stmt = insert(Owners).values(name=name)
diracx/db/sql/job/db.py CHANGED
@@ -5,15 +5,16 @@ __all__ = ["JobDB"]
5
5
  from datetime import datetime, timezone
6
6
  from typing import TYPE_CHECKING, Any, Iterable
7
7
 
8
- from sqlalchemy import bindparam, case, delete, func, insert, select, update
8
+ from sqlalchemy import bindparam, case, delete, literal, select, update
9
9
 
10
10
  if TYPE_CHECKING:
11
11
  from sqlalchemy.sql.elements import BindParameter
12
+ from sqlalchemy.sql import expression
12
13
 
13
14
  from diracx.core.exceptions import InvalidQueryError
14
15
  from diracx.core.models import JobCommand, SearchSpec, SortSpec
15
16
 
16
- from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
17
+ from ..utils import BaseSQLDB, _get_columns
17
18
  from ..utils.functions import utcnow
18
19
  from .schema import (
19
20
  HeartBeatLoggingInfo,
@@ -25,17 +26,6 @@ from .schema import (
25
26
  )
26
27
 
27
28
 
28
- def _get_columns(table, parameters):
29
- columns = [x for x in table.columns]
30
- if parameters:
31
- if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
32
- raise InvalidQueryError(
33
- f"Unrecognised parameters requested {unrecognised_parameters}"
34
- )
35
- columns = [c for c in columns if c.name in parameters]
36
- return columns
37
-
38
-
39
29
  class JobDB(BaseSQLDB):
40
30
  metadata = JobDBBase.metadata
41
31
 
@@ -54,20 +44,11 @@ class JobDB(BaseSQLDB):
54
44
  # to find a way to make it dynamic
55
45
  jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
56
46
 
57
- async def summary(self, group_by, search) -> list[dict[str, str | int]]:
47
+ async def summary(
48
+ self, group_by: list[str], search: list[SearchSpec]
49
+ ) -> list[dict[str, str | int]]:
58
50
  """Get a summary of the jobs."""
59
- columns = _get_columns(Jobs.__table__, group_by)
60
-
61
- stmt = select(*columns, func.count(Jobs.job_id).label("count"))
62
- stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
63
- stmt = stmt.group_by(*columns)
64
-
65
- # Execute the query
66
- return [
67
- dict(row._mapping)
68
- async for row in (await self.conn.stream(stmt))
69
- if row.count > 0 # type: ignore
70
- ]
51
+ return await self._summary(table=Jobs, group_by=group_by, search=search)
71
52
 
72
53
  async def search(
73
54
  self,
@@ -80,34 +61,15 @@ class JobDB(BaseSQLDB):
80
61
  page: int | None = None,
81
62
  ) -> tuple[int, list[dict[Any, Any]]]:
82
63
  """Search for jobs in the database."""
83
- # Find which columns to select
84
- columns = _get_columns(Jobs.__table__, parameters)
85
-
86
- stmt = select(*columns)
87
-
88
- stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
89
- stmt = apply_sort_constraints(Jobs.__table__.columns.__getitem__, stmt, sorts)
90
-
91
- if distinct:
92
- stmt = stmt.distinct()
93
-
94
- # Calculate total count before applying pagination
95
- total_count_subquery = stmt.alias()
96
- total_count_stmt = select(func.count()).select_from(total_count_subquery)
97
- total = (await self.conn.execute(total_count_stmt)).scalar_one()
98
-
99
- # Apply pagination
100
- if page is not None:
101
- if page < 1:
102
- raise InvalidQueryError("Page must be a positive integer")
103
- if per_page < 1:
104
- raise InvalidQueryError("Per page must be a positive integer")
105
- stmt = stmt.offset((page - 1) * per_page).limit(per_page)
106
-
107
- # Execute the query
108
- return total, [
109
- dict(row._mapping) async for row in (await self.conn.stream(stmt))
110
- ]
64
+ return await self._search(
65
+ table=Jobs,
66
+ parameters=parameters,
67
+ search=search,
68
+ sorts=sorts,
69
+ distinct=distinct,
70
+ per_page=per_page,
71
+ page=page,
72
+ )
111
73
 
112
74
  async def create_job(self, compressed_original_jdl: str):
113
75
  """Used to insert a new job with original JDL. Returns inserted job id."""
@@ -167,27 +129,14 @@ class JobDB(BaseSQLDB):
167
129
  ],
168
130
  )
169
131
 
170
- @staticmethod
171
- def _set_job_attributes_fix_value(column, value):
172
- """Apply corrections to the values before inserting them into the database.
173
-
174
- TODO: Move this logic into the sqlalchemy model.
175
- """
176
- if column == "VerifiedFlag":
177
- value_str = str(value)
178
- if value_str in ("True", "False"):
179
- return value_str
180
- if column == "AccountedFlag":
181
- value_str = str(value)
182
- if value_str in ("True", "False", "Failed"):
183
- return value_str
184
- else:
185
- return value
186
- raise NotImplementedError(f"Unrecognized value for column {column}: {value}")
187
-
188
132
  async def set_job_attributes(self, job_data):
189
133
  """Update the parameters of the given jobs."""
190
134
  # TODO: add myDate and force parameters.
135
+
136
+ if not job_data:
137
+ # nothing to do!
138
+ raise ValueError("job_data is empty")
139
+
191
140
  for job_id in job_data.keys():
192
141
  if "Status" in job_data[job_id]:
193
142
  job_data[job_id].update(
@@ -199,7 +148,11 @@ class JobDB(BaseSQLDB):
199
148
  *[
200
149
  (
201
150
  Jobs.__table__.c.JobID == job_id,
202
- self._set_job_attributes_fix_value(column, attrs[column]),
151
+ # Since the setting of the new column value is obscured by the CASE statement,
152
+ # ensure that SQLAlchemy renders the new column value with the correct type
153
+ literal(attrs[column], type_=Jobs.__table__.c[column].type)
154
+ if not isinstance(attrs[column], expression.FunctionElement)
155
+ else attrs[column],
203
156
  )
204
157
  for job_id, attrs in job_data.items()
205
158
  if column in attrs
@@ -232,7 +185,7 @@ class JobDB(BaseSQLDB):
232
185
  async def set_job_commands(self, commands: list[tuple[int, str, str]]) -> None:
233
186
  """Store a command to be passed to the job together with the next heart beat."""
234
187
  await self.conn.execute(
235
- insert(JobCommands),
188
+ JobCommands.__table__.insert(),
236
189
  [
237
190
  {
238
191
  "JobID": job_id,
@@ -2,7 +2,6 @@ from __future__ import annotations
2
2
 
3
3
  import sqlalchemy.types as types
4
4
  from sqlalchemy import (
5
- DateTime,
6
5
  ForeignKey,
7
6
  Index,
8
7
  Integer,
@@ -11,6 +10,8 @@ from sqlalchemy import (
11
10
  )
12
11
  from sqlalchemy.orm import declarative_base
13
12
 
13
+ from diracx.db.sql.utils.types import SmarterDateTime
14
+
14
15
  from ..utils import Column, EnumBackedBool, NullColumn
15
16
 
16
17
  JobDBBase = declarative_base()
@@ -19,11 +20,8 @@ JobDBBase = declarative_base()
19
20
  class AccountedFlagEnum(types.TypeDecorator):
20
21
  """Maps a ``AccountedFlagEnum()`` column to True/False in Python."""
21
22
 
22
- impl = types.Enum
23
- cache_ok: bool = True
24
-
25
- def __init__(self) -> None:
26
- super().__init__("True", "False", "Failed")
23
+ impl = types.Enum("True", "False", "Failed", name="accounted_flag_enum")
24
+ cache_ok = True
27
25
 
28
26
  def process_bind_param(self, value, dialect) -> str:
29
27
  if value is True:
@@ -63,12 +61,30 @@ class Jobs(JobDBBase):
63
61
  owner = Column("Owner", String(64), default="Unknown")
64
62
  owner_group = Column("OwnerGroup", String(128), default="Unknown")
65
63
  vo = Column("VO", String(32))
66
- submission_time = NullColumn("SubmissionTime", DateTime)
67
- reschedule_time = NullColumn("RescheduleTime", DateTime)
68
- last_update_time = NullColumn("LastUpdateTime", DateTime)
69
- start_exec_time = NullColumn("StartExecTime", DateTime)
70
- heart_beat_time = NullColumn("HeartBeatTime", DateTime)
71
- end_exec_time = NullColumn("EndExecTime", DateTime)
64
+ submission_time = NullColumn(
65
+ "SubmissionTime",
66
+ SmarterDateTime(),
67
+ )
68
+ reschedule_time = NullColumn(
69
+ "RescheduleTime",
70
+ SmarterDateTime(),
71
+ )
72
+ last_update_time = NullColumn(
73
+ "LastUpdateTime",
74
+ SmarterDateTime(),
75
+ )
76
+ start_exec_time = NullColumn(
77
+ "StartExecTime",
78
+ SmarterDateTime(),
79
+ )
80
+ heart_beat_time = NullColumn(
81
+ "HeartBeatTime",
82
+ SmarterDateTime(),
83
+ )
84
+ end_exec_time = NullColumn(
85
+ "EndExecTime",
86
+ SmarterDateTime(),
87
+ )
72
88
  status = Column("Status", String(32), default="Received")
73
89
  minor_status = Column("MinorStatus", String(128), default="Unknown")
74
90
  application_status = Column("ApplicationStatus", String(255), default="Unknown")
@@ -143,7 +159,11 @@ class HeartBeatLoggingInfo(JobDBBase):
143
159
  )
144
160
  name = Column("Name", String(100), primary_key=True)
145
161
  value = Column("Value", Text)
146
- heart_beat_time = Column("HeartBeatTime", DateTime, primary_key=True)
162
+ heart_beat_time = Column(
163
+ "HeartBeatTime",
164
+ SmarterDateTime(),
165
+ primary_key=True,
166
+ )
147
167
 
148
168
 
149
169
  class JobCommands(JobDBBase):
@@ -154,5 +174,12 @@ class JobCommands(JobDBBase):
154
174
  command = Column("Command", String(100))
155
175
  arguments = Column("Arguments", String(100))
156
176
  status = Column("Status", String(64), default="Received")
157
- reception_time = Column("ReceptionTime", DateTime, primary_key=True)
158
- execution_time = NullColumn("ExecutionTime", DateTime)
177
+ reception_time = Column(
178
+ "ReceptionTime",
179
+ SmarterDateTime(),
180
+ primary_key=True,
181
+ )
182
+ execution_time = NullColumn(
183
+ "ExecutionTime",
184
+ SmarterDateTime(),
185
+ )
@@ -1,7 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  from sqlalchemy import (
4
- DateTime,
5
4
  Double,
6
5
  Index,
7
6
  Integer,
@@ -10,6 +9,8 @@ from sqlalchemy import (
10
9
  )
11
10
  from sqlalchemy.orm import declarative_base
12
11
 
12
+ from diracx.db.sql.utils.types import SmarterDateTime
13
+
13
14
  from ..utils import Column, EnumBackedBool, NullColumn
14
15
 
15
16
  PilotAgentsDBBase = declarative_base()
@@ -29,8 +30,8 @@ class PilotAgents(PilotAgentsDBBase):
29
30
  vo = Column("VO", String(128))
30
31
  grid_type = Column("GridType", String(32), default="LCG")
31
32
  benchmark = Column("BenchMark", Double, default=0.0)
32
- submission_time = NullColumn("SubmissionTime", DateTime)
33
- last_update_time = NullColumn("LastUpdateTime", DateTime)
33
+ submission_time = NullColumn("SubmissionTime", SmarterDateTime)
34
+ last_update_time = NullColumn("LastUpdateTime", SmarterDateTime)
34
35
  status = Column("Status", String(32), default="Unknown")
35
36
  status_reason = Column("StatusReason", String(255), default="Unknown")
36
37
  accounting_sent = Column("AccountingSent", EnumBackedBool(), default=False)
@@ -47,7 +48,7 @@ class JobToPilotMapping(PilotAgentsDBBase):
47
48
 
48
49
  pilot_id = Column("PilotID", Integer, primary_key=True)
49
50
  job_id = Column("JobID", Integer, primary_key=True)
50
- start_time = Column("StartTime", DateTime)
51
+ start_time = Column("StartTime", SmarterDateTime)
51
52
 
52
53
  __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
53
54
 
@@ -3,6 +3,7 @@ from __future__ import annotations
3
3
  from .base import (
4
4
  BaseSQLDB,
5
5
  SQLDBUnavailableError,
6
+ _get_columns,
6
7
  apply_search_filters,
7
8
  apply_sort_constraints,
8
9
  )
@@ -10,6 +11,7 @@ from .functions import hash, substract_date, utcnow
10
11
  from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn
11
12
 
12
13
  __all__ = (
14
+ "_get_columns",
13
15
  "utcnow",
14
16
  "Column",
15
17
  "NullColumn",
@@ -8,18 +8,23 @@ from abc import ABCMeta
8
8
  from collections.abc import AsyncIterator
9
9
  from contextvars import ContextVar
10
10
  from datetime import datetime
11
- from typing import Self, cast
11
+ from typing import Any, Self, cast
12
12
 
13
13
  from pydantic import TypeAdapter
14
- from sqlalchemy import DateTime, MetaData, select
14
+ from sqlalchemy import DateTime, MetaData, func, select
15
15
  from sqlalchemy.exc import OperationalError
16
16
  from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
17
17
 
18
18
  from diracx.core.exceptions import InvalidQueryError
19
19
  from diracx.core.extensions import select_from_extension
20
- from diracx.core.models import SortDirection
20
+ from diracx.core.models import (
21
+ SearchSpec,
22
+ SortDirection,
23
+ SortSpec,
24
+ )
21
25
  from diracx.core.settings import SqlalchemyDsn
22
26
  from diracx.db.exceptions import DBUnavailableError
27
+ from diracx.db.sql.utils.types import SmarterDateTime
23
28
 
24
29
  from .functions import date_trunc
25
30
 
@@ -227,6 +232,71 @@ class BaseSQLDB(metaclass=ABCMeta):
227
232
  except OperationalError as e:
228
233
  raise SQLDBUnavailableError("Cannot ping the DB") from e
229
234
 
235
+ async def _search(
236
+ self,
237
+ table: Any,
238
+ parameters: list[str] | None,
239
+ search: list[SearchSpec],
240
+ sorts: list[SortSpec],
241
+ *,
242
+ distinct: bool = False,
243
+ per_page: int = 100,
244
+ page: int | None = None,
245
+ ) -> tuple[int, list[dict[str, Any]]]:
246
+ """Search for elements in a table."""
247
+ # Find which columns to select
248
+ columns = _get_columns(table.__table__, parameters)
249
+
250
+ stmt = select(*columns)
251
+
252
+ stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
253
+ stmt = apply_sort_constraints(table.__table__.columns.__getitem__, stmt, sorts)
254
+
255
+ if distinct:
256
+ stmt = stmt.distinct()
257
+
258
+ # Calculate total count before applying pagination
259
+ total_count_subquery = stmt.alias()
260
+ total_count_stmt = select(func.count()).select_from(total_count_subquery)
261
+ total = (await self.conn.execute(total_count_stmt)).scalar_one()
262
+
263
+ # Apply pagination
264
+ if page is not None:
265
+ if page < 1:
266
+ raise InvalidQueryError("Page must be a positive integer")
267
+ if per_page < 1:
268
+ raise InvalidQueryError("Per page must be a positive integer")
269
+ stmt = stmt.offset((page - 1) * per_page).limit(per_page)
270
+
271
+ # Execute the query
272
+ return total, [
273
+ dict(row._mapping) async for row in (await self.conn.stream(stmt))
274
+ ]
275
+
276
+ async def _summary(
277
+ self, table: Any, group_by: list[str], search: list[SearchSpec]
278
+ ) -> list[dict[str, str | int]]:
279
+ """Get a summary of the elements of a table."""
280
+ columns = _get_columns(table.__table__, group_by)
281
+
282
+ pk_columns = list(table.__table__.primary_key.columns)
283
+ if not pk_columns:
284
+ raise ValueError(
285
+ "Model has no primary key and no count_column was provided."
286
+ )
287
+ count_col = pk_columns[0]
288
+
289
+ stmt = select(*columns, func.count(count_col).label("count"))
290
+ stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
291
+ stmt = stmt.group_by(*columns)
292
+
293
+ # Execute the query
294
+ return [
295
+ dict(row._mapping)
296
+ async for row in (await self.conn.stream(stmt))
297
+ if row.count > 0 # type: ignore
298
+ ]
299
+
230
300
 
231
301
  def find_time_resolution(value):
232
302
  if isinstance(value, datetime):
@@ -258,6 +328,17 @@ def find_time_resolution(value):
258
328
  raise InvalidQueryError(f"Cannot parse {value=}")
259
329
 
260
330
 
331
+ def _get_columns(table, parameters):
332
+ columns = [x for x in table.columns]
333
+ if parameters:
334
+ if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
335
+ raise InvalidQueryError(
336
+ f"Unrecognised parameters requested {unrecognised_parameters}"
337
+ )
338
+ columns = [c for c in columns if c.name in parameters]
339
+ return columns
340
+
341
+
261
342
  def apply_search_filters(column_mapping, stmt, search):
262
343
  for query in search:
263
344
  try:
@@ -265,7 +346,7 @@ def apply_search_filters(column_mapping, stmt, search):
265
346
  except KeyError as e:
266
347
  raise InvalidQueryError(f"Unknown column {query['parameter']}") from e
267
348
 
268
- if isinstance(column.type, DateTime):
349
+ if isinstance(column.type, (DateTime, SmarterDateTime)):
269
350
  if "value" in query and isinstance(query["value"], str):
270
351
  resolution, value = find_time_resolution(query["value"])
271
352
  if resolution:
@@ -300,6 +381,15 @@ def apply_search_filters(column_mapping, stmt, search):
300
381
  expr = column.like(query["value"])
301
382
  elif query["operator"] in "ilike":
302
383
  expr = column.ilike(query["value"])
384
+ elif query["operator"] == "not like":
385
+ expr = column.not_like(query["value"])
386
+ elif query["operator"] == "regex":
387
+ # We check the regex validity here
388
+ try:
389
+ re.compile(query["value"])
390
+ except re.error as e:
391
+ raise InvalidQueryError(f"Invalid regex {query['value']}") from e
392
+ expr = column.regexp_match(query["value"])
303
393
  else:
304
394
  raise InvalidQueryError(f"Unknown filter {query=}")
305
395
  stmt = stmt.where(expr)
@@ -1,6 +1,8 @@
1
1
  from __future__ import annotations
2
2
 
3
+ from datetime import datetime
3
4
  from functools import partial
5
+ from zoneinfo import ZoneInfo
4
6
 
5
7
  import sqlalchemy.types as types
6
8
  from sqlalchemy import Column as RawColumn
@@ -20,11 +22,8 @@ def EnumColumn(name, enum_type, **kwargs): # noqa: N802
20
22
  class EnumBackedBool(types.TypeDecorator):
21
23
  """Maps a ``EnumBackedBool()`` column to True/False in Python."""
22
24
 
23
- impl = types.Enum
24
- cache_ok: bool = True
25
-
26
- def __init__(self) -> None:
27
- super().__init__("True", "False")
25
+ impl = types.Enum("True", "False", name="enum_backed_bool")
26
+ cache_ok = True
28
27
 
29
28
  def process_bind_param(self, value, dialect) -> str:
30
29
  if value is True:
@@ -41,3 +40,98 @@ class EnumBackedBool(types.TypeDecorator):
41
40
  return False
42
41
  else:
43
42
  raise NotImplementedError(f"Unknown {value=}")
43
+
44
+
45
+ class SmarterDateTime(types.TypeDecorator):
46
+ """A DateTime type that also accepts ISO8601 strings.
47
+
48
+ Takes into account converting timezone aware datetime objects into
49
+ naive form and back when needed.
50
+
51
+ """
52
+
53
+ impl = DateTime()
54
+ cache_ok = True
55
+
56
+ def __init__(
57
+ self,
58
+ stored_tz: ZoneInfo | None = ZoneInfo("UTC"),
59
+ returned_tz: ZoneInfo = ZoneInfo("UTC"),
60
+ stored_naive_sqlite=True,
61
+ stored_naive_mysql=True,
62
+ stored_naive_postgres=False, # Forces timezone-awareness
63
+ ):
64
+ self._stored_naive_dialect = {
65
+ "sqlite": stored_naive_sqlite,
66
+ "mysql": stored_naive_mysql,
67
+ "postgres": stored_naive_postgres,
68
+ }
69
+ self._stored_tz: ZoneInfo | None = stored_tz # None = Local timezone
70
+ self._returned_tz: ZoneInfo = returned_tz
71
+
72
+ def _stored_naive(self, dialect):
73
+ if dialect.name not in self._stored_naive_dialect:
74
+ raise NotImplementedError(dialect.name)
75
+ return self._stored_naive_dialect.get(dialect.name)
76
+
77
+ def process_bind_param(self, value, dialect):
78
+ if value is None:
79
+ return None
80
+
81
+ if isinstance(value, str):
82
+ try:
83
+ value: datetime = datetime.fromisoformat(value)
84
+ except ValueError as err:
85
+ raise ValueError(f"Unable to parse datetime string: {value}") from err
86
+
87
+ if not isinstance(value, datetime):
88
+ raise ValueError(f"Expected datetime or ISO8601 string, but got {value!r}")
89
+
90
+ if not value.tzinfo:
91
+ raise ValueError(
92
+ f"Provided timestamp {value=} has no tzinfo -"
93
+ " this is problematic and may cause inconsistencies in stored timestamps.\n"
94
+ " Please always work with tz-aware datetimes / attach tzinfo to your datetime objects:"
95
+ " e.g. datetime.now(tz=timezone.utc) or use datetime_obj.astimezone() with no arguments if you need to "
96
+ "attach the local timezone to a local naive timestamp."
97
+ )
98
+
99
+ # Check that we need to convert the timezone to match self._stored_tz timezone:
100
+ if self._stored_naive(dialect):
101
+ # if self._stored_tz is None, we use our local/system timezone.
102
+ stored_tz = self._stored_tz
103
+
104
+ # astimezone converts to the stored timezone (local timezone if None)
105
+ # replace strips the TZ info --> naive datetime object
106
+ value = value.astimezone(tz=stored_tz).replace(tzinfo=None)
107
+
108
+ return value
109
+
110
+ def process_result_value(self, value, dialect):
111
+ if value is None:
112
+ return None
113
+ if not isinstance(value, datetime):
114
+ raise NotImplementedError(f"{value=} not a datetime object")
115
+
116
+ if self._stored_naive(dialect):
117
+ # Here we add back the tzinfo to the naive timestamp
118
+ # from the DB to make it aware again.
119
+ if value.tzinfo is None:
120
+ # we are definitely given a naive timestamp, so handle it.
121
+ # add back the timezone info if stored_tz is set
122
+ if self._stored_tz:
123
+ value = value.replace(tzinfo=self._stored_tz)
124
+ else:
125
+ # if stored as a local time, add back the system timezone info...
126
+ value = value.astimezone()
127
+ else:
128
+ raise ValueError(
129
+ f"stored_naive is True for {dialect.name=}, but the database engine returned "
130
+ "a tz-aware datetime. You need to check the SQLAlchemy model is consistent with the DB schema."
131
+ )
132
+
133
+ # finally, convert the datetime according to the "returned_tz"
134
+ value = value.astimezone(self._returned_tz)
135
+
136
+ # phew...
137
+ return value
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diracx-db
3
- Version: 0.0.1a46
3
+ Version: 0.0.1a48
4
4
  Summary: TODO
5
5
  License: GPL-3.0-only
6
6
  Classifier: Intended Audience :: Science/Research
@@ -3,35 +3,35 @@ diracx/db/__main__.py,sha256=3yaUP1ig-yaPSQM4wy6CtSXXHivQg-hIz2FeBt7joBc,1714
3
3
  diracx/db/exceptions.py,sha256=1nn-SZLG-nQwkxbvHjZqXhE5ouzWj1f3qhSda2B4ZEg,83
4
4
  diracx/db/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
5
5
  diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
6
- diracx/db/os/job_parameters.py,sha256=loAc-bo3u-RMAp_H1g8VRt8T-rCCsXp_d9aCvg5OS-A,1225
6
+ diracx/db/os/job_parameters.py,sha256=3w_CeA2z-cY5pWwXkGu-Fod27FobbUXuwVKK-jN037U,1479
7
7
  diracx/db/os/utils.py,sha256=V4T-taos64SFNcorfIr7mq5l5y88K6TzyCj1YqWk8VI,11562
8
8
  diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
9
9
  diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
10
10
  diracx/db/sql/auth/db.py,sha256=QJtBqMrhOf97UvMG0WpyjsgIRiu19v04FoDzXAyXtT0,8952
11
11
  diracx/db/sql/auth/schema.py,sha256=x2PEbmM_bNPdZUN5BMGMrdSmX8zkDeJ3P9XfhLBGBTs,3173
12
12
  diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
13
- diracx/db/sql/dummy/db.py,sha256=IW4FzG7ERKbhZvC32KL7Rodu2u-zKAf8BryO4VAdJew,1650
13
+ diracx/db/sql/dummy/db.py,sha256=MKSUSJI1BlRgK08tjCfkCkOz02asvJAeBw60pAdiGV8,1212
14
14
  diracx/db/sql/dummy/schema.py,sha256=9zI53pKlzc6qBezsyjkatOQrNZdGCjwgjQ8Iz_pyAXs,789
15
15
  diracx/db/sql/job/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
16
- diracx/db/sql/job/db.py,sha256=TnEc0fckiuMJAZg2v1_Pbwfn7kDPDam6TXp9ySuiddk,11910
17
- diracx/db/sql/job/schema.py,sha256=eFgZshe6NEzOM2qI0HI9Y3abrqDMoQIwa9L0vZugHcU,5431
16
+ diracx/db/sql/job/db.py,sha256=bX-4OMyW4h9tqeTE3OvonxTXlL6j_Qvv9uEtK5SthN8,10120
17
+ diracx/db/sql/job/schema.py,sha256=fJdmiLp6psdAjo_CoBfSAGSYk2NJkSBwvik9tznESD0,5740
18
18
  diracx/db/sql/job_logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
19
19
  diracx/db/sql/job_logging/db.py,sha256=hyklARuEj3R1sSJ8UaObRprmsRx7RjbKAcbfgT9BwRg,5496
20
20
  diracx/db/sql/job_logging/schema.py,sha256=k6uBw-RHAcJ5GEleNpiWoXEJBhCiNG-y4xAgBKHZjjM,2524
21
21
  diracx/db/sql/pilot_agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
22
22
  diracx/db/sql/pilot_agents/db.py,sha256=6CQ0QGV4NhsGKVCygEtE4kmIjT89xJwrIMuYZTslWFE,1231
23
- diracx/db/sql/pilot_agents/schema.py,sha256=KeWnFSpYOTrT3-_rOCFjbjNnPNXKnUZiJVsu4vv5U2U,2149
23
+ diracx/db/sql/pilot_agents/schema.py,sha256=BTFLuiwcxAvAtTvTP9C7DbGtXoM-IHVDG9k7HMx62AA,2211
24
24
  diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
25
25
  diracx/db/sql/sandbox_metadata/db.py,sha256=FtyPx6GAGJAH-lmuw8PQj6_KGHG6t3AC3-E9uWf-JNs,10236
26
26
  diracx/db/sql/sandbox_metadata/schema.py,sha256=V5gV2PHwzTbBz_th9ribLfE7Lqk8YGemDmvqq4jWQJ4,1530
27
27
  diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
28
28
  diracx/db/sql/task_queue/db.py,sha256=2qul1D2tX2uCI92N591WK5xWHakG0pNibzDwKQ7W-I8,6246
29
29
  diracx/db/sql/task_queue/schema.py,sha256=5efAgvNYRkLlaJ2NzRInRfmVa3tyIzQu2l0oRPy4Kzw,3258
30
- diracx/db/sql/utils/__init__.py,sha256=QkvpqBuIAgkAOywAssYzdxSzUQVZlSUumK7mPxotXfM,547
31
- diracx/db/sql/utils/base.py,sha256=HYQuX16mgg9LAMtAEmbTmJFIN0OSMe1Hcb57dtl7LCc,12367
30
+ diracx/db/sql/utils/__init__.py,sha256=XYbv-AJAPl7bb8dETpjc07olmtXQ0h1MFUbLqjAphQE,585
31
+ diracx/db/sql/utils/base.py,sha256=Dn7a-ICuFNTT3w0QtbS63uF1S4Wnn5NK3cu6Pzn4a6A,15641
32
32
  diracx/db/sql/utils/functions.py,sha256=_E4tc9Gti6LuSh7QEyoqPJSvCuByVqvRenOXCzxsulE,4014
33
- diracx/db/sql/utils/types.py,sha256=yU-tXsu6hFGPsr9ba1n3ZjGPnHQI_06lbpkTeDCWJtg,1287
34
- diracx_db-0.0.1a46.dist-info/METADATA,sha256=hu6BtBuYz30tuMO8Z40gyZhSf53QzqTKO45ryw6tTL4,675
35
- diracx_db-0.0.1a46.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
36
- diracx_db-0.0.1a46.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
37
- diracx_db-0.0.1a46.dist-info/RECORD,,
33
+ diracx/db/sql/utils/types.py,sha256=KNZWJfpvHTjfIPg6Nn7zY-rS0q3ybnirHcTcLAYSYbE,5118
34
+ diracx_db-0.0.1a48.dist-info/METADATA,sha256=A-s-kA2Qo5jEmURg_Hi6xAz1H7gGs2cgkq1hyBK7i7s,675
35
+ diracx_db-0.0.1a48.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
36
+ diracx_db-0.0.1a48.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
37
+ diracx_db-0.0.1a48.dist-info/RECORD,,