diracx-db 0.0.1a44__tar.gz → 0.0.1a46__tar.gz

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.
Files changed (57) hide show
  1. diracx_db-0.0.1a46/.gitignore +98 -0
  2. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/PKG-INFO +3 -3
  3. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/pyproject.toml +9 -9
  4. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/__main__.py +0 -1
  5. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/os/job_parameters.py +2 -2
  6. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job/db.py +1 -1
  7. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job_logging/db.py +0 -1
  8. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job_logging/schema.py +1 -1
  9. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/pilot_agents/db.py +0 -1
  10. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/sandbox_metadata/db.py +94 -5
  11. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/task_queue/db.py +2 -2
  12. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/utils/base.py +1 -1
  13. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/utils/functions.py +31 -0
  14. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/opensearch/test_index_template.py +4 -2
  15. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/pilot_agents/test_pilot_agents_db.py +0 -1
  16. diracx_db-0.0.1a46/tests/test_freeze_time.py +153 -0
  17. diracx_db-0.0.1a44/setup.cfg +0 -4
  18. diracx_db-0.0.1a44/src/diracx_db.egg-info/PKG-INFO +0 -19
  19. diracx_db-0.0.1a44/src/diracx_db.egg-info/SOURCES.txt +0 -53
  20. diracx_db-0.0.1a44/src/diracx_db.egg-info/dependency_links.txt +0 -1
  21. diracx_db-0.0.1a44/src/diracx_db.egg-info/entry_points.txt +0 -10
  22. diracx_db-0.0.1a44/src/diracx_db.egg-info/requires.txt +0 -8
  23. diracx_db-0.0.1a44/src/diracx_db.egg-info/top_level.txt +0 -1
  24. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/README.md +0 -0
  25. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/__init__.py +0 -0
  26. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/exceptions.py +0 -0
  27. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/os/__init__.py +0 -0
  28. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/os/utils.py +0 -0
  29. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/py.typed +0 -0
  30. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/__init__.py +0 -0
  31. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/auth/__init__.py +0 -0
  32. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/auth/db.py +0 -0
  33. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/auth/schema.py +0 -0
  34. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/dummy/__init__.py +0 -0
  35. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/dummy/db.py +0 -0
  36. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/dummy/schema.py +0 -0
  37. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job/__init__.py +0 -0
  38. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job/schema.py +0 -0
  39. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/job_logging/__init__.py +0 -0
  40. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/pilot_agents/__init__.py +0 -0
  41. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/pilot_agents/schema.py +0 -0
  42. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/sandbox_metadata/__init__.py +0 -0
  43. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/sandbox_metadata/schema.py +0 -0
  44. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/task_queue/__init__.py +0 -0
  45. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/task_queue/schema.py +0 -0
  46. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/utils/__init__.py +0 -0
  47. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/src/diracx/db/sql/utils/types.py +0 -0
  48. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/auth/test_authorization_flow.py +0 -0
  49. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/auth/test_device_flow.py +0 -0
  50. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/auth/test_refresh_token.py +0 -0
  51. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/jobs/test_job_db.py +0 -0
  52. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/jobs/test_job_logging_db.py +0 -0
  53. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/jobs/test_sandbox_metadata.py +0 -0
  54. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/opensearch/test_connection.py +0 -0
  55. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/opensearch/test_search.py +0 -0
  56. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/pilot_agents/__init__.py +0 -0
  57. {diracx_db-0.0.1a44 → diracx_db-0.0.1a46}/tests/test_dummy_db.py +0 -0
@@ -0,0 +1,98 @@
1
+ # Python
2
+ *.py[cod]
3
+
4
+ # Conda
5
+ .conda
6
+
7
+ # C extensions
8
+ *.so
9
+ var
10
+ sdist
11
+ lib
12
+ lib64
13
+
14
+ # Packages
15
+ *.egg
16
+ *.egg-info
17
+ dist
18
+ build
19
+ eggs
20
+ parts
21
+ bin
22
+ develop-eggs
23
+ .installed.cfg
24
+ *.whl
25
+
26
+ # Translations
27
+ *.mo
28
+
29
+ # Mr Developer
30
+ .mr.developer.cfg
31
+
32
+ # Installer logs
33
+ pip-log.txt
34
+
35
+ # Unit test / coverage reports
36
+ .coverage
37
+ .tox
38
+ .ruff_cache
39
+ .mypy_cache
40
+
41
+ # Eclipse
42
+ .project
43
+ .pydevproject
44
+ .pyproject
45
+ .settings
46
+ .metadata
47
+
48
+ #VSCode
49
+ .vscode
50
+ .env
51
+
52
+ # Vim
53
+ .*.sw[a-z]
54
+ *.un~
55
+ Session.vim
56
+ *~
57
+
58
+ # Intellij
59
+ .idea/
60
+ LHCbDIRAC.iml
61
+
62
+ # MaxOSX files
63
+ .DS_Store
64
+
65
+ # test stuff
66
+ .pytest_cache
67
+ .cache
68
+ __pycache__
69
+ pytests.xml
70
+ nosetests.xml
71
+ coverage.xml
72
+ Local_*
73
+ .hypothesis
74
+ *.gz
75
+ htmlcov/
76
+ *.xml.temp
77
+
78
+ #remove in case we want a specific LHCb one
79
+ .pylintrc
80
+
81
+ # docs
82
+ # this is auto generated
83
+ docs/source/CodeDocumentation/
84
+ docs/source/AdministratorGuide/Configuration/ExampleConfig.rst
85
+ docs/source/AdministratorGuide/CommandReference
86
+ docs/source/UserGuide/CommandReference
87
+ docs/_build
88
+ docs/source/_build
89
+
90
+
91
+ # CMT junk
92
+ /*/*/x86_64-*-*-*/
93
+ */*/cmt/Makefile
94
+
95
+ # pixi environments
96
+ .pixi
97
+ pixi.lock
98
+ *.egg-info
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: diracx-db
3
- Version: 0.0.1a44
3
+ Version: 0.0.1a46
4
4
  Summary: TODO
5
5
  License: GPL-3.0-only
6
6
  Classifier: Intended Audience :: Science/Research
@@ -9,11 +9,11 @@ Classifier: Programming Language :: Python :: 3
9
9
  Classifier: Topic :: Scientific/Engineering
10
10
  Classifier: Topic :: System :: Distributed Computing
11
11
  Requires-Python: >=3.11
12
- Description-Content-Type: text/markdown
13
12
  Requires-Dist: diracx-core
14
13
  Requires-Dist: opensearch-py[async]
15
14
  Requires-Dist: pydantic>=2.10
16
15
  Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
17
16
  Requires-Dist: uuid-utils
18
17
  Provides-Extra: testing
19
- Requires-Dist: diracx-testing; extra == "testing"
18
+ Requires-Dist: diracx-testing; extra == 'testing'
19
+ Requires-Dist: freezegun; extra == 'testing'
@@ -22,9 +22,7 @@ dependencies = [
22
22
  dynamic = ["version"]
23
23
 
24
24
  [project.optional-dependencies]
25
- testing = [
26
- "diracx-testing",
27
- ]
25
+ testing = ["diracx-testing", "freezegun"]
28
26
 
29
27
  [project.entry-points."diracx.dbs.sql"]
30
28
  AuthDB = "diracx.db.sql:AuthDB"
@@ -37,16 +35,18 @@ TaskQueueDB = "diracx.db.sql:TaskQueueDB"
37
35
  [project.entry-points."diracx.dbs.os"]
38
36
  JobParametersDB = "diracx.db.os:JobParametersDB"
39
37
 
40
- [tool.setuptools.packages.find]
41
- where = ["src"]
42
-
43
38
  [build-system]
44
- requires = ["setuptools>=61", "wheel", "setuptools_scm>=8"]
45
- build-backend = "setuptools.build_meta"
39
+ requires = ["hatchling", "hatch-vcs"]
40
+ build-backend = "hatchling.build"
41
+
42
+ [tool.hatch.version]
43
+ source = "vcs"
46
44
 
47
- [tool.setuptools_scm]
45
+ [tool.hatch.version.raw-options]
48
46
  root = ".."
49
47
 
48
+ [tool.hatch.build.targets.wheel]
49
+ packages = ["src/diracx"]
50
50
 
51
51
  [tool.pytest.ini_options]
52
52
  testpaths = ["tests"]
@@ -31,7 +31,6 @@ async def init_sql():
31
31
  from diracx.db.sql.utils import BaseSQLDB
32
32
 
33
33
  for db_name, db_url in BaseSQLDB.available_urls().items():
34
-
35
34
  logger.info("Initialising %s", db_name)
36
35
  db = BaseSQLDB.available_implementations(db_name)[0](db_url)
37
36
  async with db.engine_context():
@@ -1,6 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
- from DIRAC.Core.Utilities import TimeUtilities
3
+ from datetime import UTC, datetime
4
4
 
5
5
  from diracx.db.os.utils import BaseOSDB
6
6
 
@@ -31,7 +31,7 @@ class JobParametersDB(BaseOSDB):
31
31
  def upsert(self, vo, doc_id, document):
32
32
  document = {
33
33
  "JobID": doc_id,
34
- "timestamp": TimeUtilities.toEpochMilliSeconds(),
34
+ "timestamp": int(datetime.now(tz=UTC).timestamp() * 1000),
35
35
  **document,
36
36
  }
37
37
  return super().upsert(vo, doc_id, document)
@@ -50,7 +50,7 @@ class JobDB(BaseSQLDB):
50
50
  }
51
51
 
52
52
  # TODO: this is copied from the DIRAC JobDB
53
- # but is overwriten in LHCbDIRAC, so we need
53
+ # but is overwritten in LHCbDIRAC, so we need
54
54
  # to find a way to make it dynamic
55
55
  jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
56
56
 
@@ -89,7 +89,6 @@ class JobLoggingDB(BaseSQLDB):
89
89
  status_time,
90
90
  status_source,
91
91
  ) in rows:
92
-
93
92
  values[job_id].append(
94
93
  [
95
94
  status,
@@ -44,7 +44,7 @@ class MagicEpochDateTime(TypeDecorator):
44
44
  """
45
45
  if value is None:
46
46
  return None
47
- # Carefully convert from Decimal to datetime to avoid loosing precision
47
+ # Carefully convert from Decimal to datetime to avoid losing precision
48
48
  value += self.MAGIC_EPOC_NUMBER
49
49
  value_int = int(value)
50
50
  result = datetime.fromtimestamp(value_int, tz=UTC)
@@ -20,7 +20,6 @@ class PilotAgentsDB(BaseSQLDB):
20
20
  grid_type: str = "DIRAC",
21
21
  pilot_stamps: dict | None = None,
22
22
  ) -> None:
23
-
24
23
  if pilot_stamps is None:
25
24
  pilot_stamps = {}
26
25
 
@@ -1,8 +1,25 @@
1
1
  from __future__ import annotations
2
2
 
3
- from typing import Any
3
+ import logging
4
+ from contextlib import asynccontextmanager
5
+ from functools import partial
6
+ from typing import Any, AsyncGenerator
4
7
 
5
- from sqlalchemy import Executable, delete, insert, literal, select, update
8
+ from sqlalchemy import (
9
+ BigInteger,
10
+ Column,
11
+ Executable,
12
+ MetaData,
13
+ Table,
14
+ and_,
15
+ delete,
16
+ exists,
17
+ insert,
18
+ literal,
19
+ or_,
20
+ select,
21
+ update,
22
+ )
6
23
  from sqlalchemy.exc import IntegrityError, NoResultFound
7
24
 
8
25
  from diracx.core.exceptions import (
@@ -12,15 +29,25 @@ from diracx.core.exceptions import (
12
29
  )
13
30
  from diracx.core.models import SandboxInfo, SandboxType, UserInfo
14
31
  from diracx.db.sql.utils.base import BaseSQLDB
15
- from diracx.db.sql.utils.functions import utcnow
32
+ from diracx.db.sql.utils.functions import days_since, utcnow
16
33
 
17
34
  from .schema import Base as SandboxMetadataDBBase
18
35
  from .schema import SandBoxes, SBEntityMapping, SBOwners
19
36
 
37
+ logger = logging.getLogger(__name__)
38
+
20
39
 
21
40
  class SandboxMetadataDB(BaseSQLDB):
22
41
  metadata = SandboxMetadataDBBase.metadata
23
42
 
43
+ # Temporary table to store the sandboxes to delete, see `select_and_delete_expired`
44
+ _temp_table = Table(
45
+ "sb_to_delete",
46
+ MetaData(),
47
+ Column("SBId", BigInteger, primary_key=True),
48
+ prefixes=["TEMPORARY"],
49
+ )
50
+
24
51
  async def get_owner_id(self, user: UserInfo) -> int | None:
25
52
  """Get the id of the owner from the database."""
26
53
  stmt = select(SBOwners.OwnerID).where(
@@ -85,7 +112,13 @@ class SandboxMetadataDB(BaseSQLDB):
85
112
  .values(LastAccessTime=utcnow())
86
113
  )
87
114
  result = await self.conn.execute(stmt)
88
- assert result.rowcount == 1
115
+ if result.rowcount == 0:
116
+ # If the update didn't affect any row, the sandbox doesn't exist
117
+ raise SandboxNotFoundError(pfn, se_name)
118
+ elif result.rowcount != 1:
119
+ raise NotImplementedError(
120
+ "More than one sandbox was updated. This should not happen."
121
+ )
89
122
 
90
123
  async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool | None:
91
124
  """Checks if a sandbox exists and has been assigned."""
@@ -128,7 +161,7 @@ class SandboxMetadataDB(BaseSQLDB):
128
161
  sb_type: SandboxType,
129
162
  se_name: str,
130
163
  ) -> None:
131
- """Mapp sandbox and jobs."""
164
+ """Map sandbox and jobs."""
132
165
  for job_id in jobs_ids:
133
166
  # Define the entity id as 'Entity:entity_id' due to the DB definition:
134
167
  entity_id = self.jobid_to_entity_id(job_id)
@@ -187,3 +220,59 @@ class SandboxMetadataDB(BaseSQLDB):
187
220
  .values(Assigned=False)
188
221
  )
189
222
  await self.conn.execute(unassign_stmt)
223
+
224
+ @asynccontextmanager
225
+ async def delete_unused_sandboxes(
226
+ self, *, limit: int | None = None
227
+ ) -> AsyncGenerator[AsyncGenerator[str, None], None]:
228
+ """Get the sandbox PFNs to delete.
229
+
230
+ The result of this function can be used as an async context manager
231
+ to yield the PFNs of the sandboxes to delete. The context manager
232
+ will automatically remove the sandboxes from the database upon exit.
233
+
234
+ Args:
235
+ limit: If not None, the maximum number of sandboxes to delete.
236
+
237
+ """
238
+ conditions = [
239
+ # If it has assigned to a job but is no longer mapped it can be removed
240
+ and_(
241
+ SandBoxes.Assigned,
242
+ ~exists().where(SBEntityMapping.SBId == SandBoxes.SBId),
243
+ ),
244
+ # If the sandbox is still unassigned after 15 days, remove it
245
+ and_(~SandBoxes.Assigned, days_since(SandBoxes.LastAccessTime) >= 15),
246
+ ]
247
+ # Sandboxes which are not on S3 will be handled by legacy DIRAC
248
+ condition = and_(SandBoxes.SEPFN.like("/S3/%"), or_(*conditions))
249
+
250
+ # Copy the in-flight rows to a temporary table
251
+ await self.conn.run_sync(partial(self._temp_table.create, checkfirst=True))
252
+ select_stmt = select(SandBoxes.SBId).where(condition)
253
+ if limit:
254
+ select_stmt = select_stmt.limit(limit)
255
+ insert_stmt = insert(self._temp_table).from_select(["SBId"], select_stmt)
256
+ await self.conn.execute(insert_stmt)
257
+
258
+ try:
259
+ # Select the sandbox PFNs from the temporary table and yield them
260
+ select_stmt = select(SandBoxes.SEPFN).join(
261
+ self._temp_table, self._temp_table.c.SBId == SandBoxes.SBId
262
+ )
263
+
264
+ async def yield_pfns() -> AsyncGenerator[str, None]:
265
+ async for row in await self.conn.stream(select_stmt):
266
+ yield row.SEPFN
267
+
268
+ yield yield_pfns()
269
+
270
+ # Delete the sandboxes from the main table
271
+ delete_stmt = delete(SandBoxes).where(
272
+ SandBoxes.SBId.in_(select(self._temp_table.c.SBId))
273
+ )
274
+ result = await self.conn.execute(delete_stmt)
275
+ logger.info("Deleted %d expired/unassigned sandboxes", result.rowcount)
276
+
277
+ finally:
278
+ await self.conn.run_sync(partial(self._temp_table.drop, checkfirst=True))
@@ -56,7 +56,7 @@ class TaskQueueDB(BaseSQLDB):
56
56
  )
57
57
  rows = await self.conn.execute(stmt)
58
58
  # Get owners in this group and the amount of times they appear
59
- # TODO: I guess the rows are already a list of tupes
59
+ # TODO: I guess the rows are already a list of tuples
60
60
  # maybe refactor
61
61
  return {r[0]: r[1] for r in rows if r}
62
62
 
@@ -108,7 +108,7 @@ class TaskQueueDB(BaseSQLDB):
108
108
  tq_ids: list[int],
109
109
  priority: float,
110
110
  ):
111
- """Set the priority for a user/userGroup combo given a splitted share."""
111
+ """Set the priority for a user/userGroup combo given a split share."""
112
112
  update_stmt = (
113
113
  update(TaskQueues)
114
114
  .where(TaskQueues.TQId.in_(tq_ids))
@@ -51,7 +51,7 @@ class BaseSQLDB(metaclass=ABCMeta):
51
51
  of the same database. To list the available implementations use
52
52
  `BaseSQLDB.available_implementations(db_name)`. The first entry in this list
53
53
  will be the preferred implementation and it can be initialized by calling
54
- it's `__init__` function with a URL perviously obtained from
54
+ it's `__init__` function with a URL previously obtained from
55
55
  `BaseSQLDB.available_urls`.
56
56
 
57
57
  To control the lifetime of the SQLAlchemy engine used for connecting to the
@@ -103,6 +103,37 @@ def sqlite_date_trunc(element, compiler, **kw):
103
103
  )
104
104
 
105
105
 
106
+ class days_since(expression.FunctionElement): # noqa: N801
107
+ """Sqlalchemy function to get the number of days since a given date.
108
+
109
+ Primarily used to be able to query for a specific resolution of a date e.g.
110
+
111
+ select * from table where days_since(date_column) = 0
112
+ select * from table where days_since(date_column) = 1
113
+ """
114
+
115
+ type = DateTime()
116
+ inherit_cache = False
117
+
118
+ def __init__(self, *args, **kwargs) -> None:
119
+ super().__init__(*args, **kwargs)
120
+
121
+
122
+ @compiles(days_since, "postgresql")
123
+ def pg_days_since(element, compiler, **kw):
124
+ return f"EXTRACT(DAY FROM (now() - {compiler.process(element.clauses)}))"
125
+
126
+
127
+ @compiles(days_since, "mysql")
128
+ def mysql_days_since(element, compiler, **kw):
129
+ return f"DATEDIFF(NOW(), {compiler.process(element.clauses)})"
130
+
131
+
132
+ @compiles(days_since, "sqlite")
133
+ def sqlite_days_since(element, compiler, **kw):
134
+ return f"julianday('now') - julianday({compiler.process(element.clauses)})"
135
+
136
+
106
137
  def substract_date(**kwargs: float) -> datetime:
107
138
  return datetime.now(tz=timezone.utc) - timedelta(**kwargs)
108
139
 
@@ -37,13 +37,15 @@ async def _get_test_index_mappings(dummy_opensearch_db: DummyOSDB):
37
37
 
38
38
  # At this point the index should not exist yet
39
39
  with pytest.raises(opensearchpy.exceptions.NotFoundError):
40
- await dummy_opensearch_db.client.indices.get_mapping(index_name)
40
+ await dummy_opensearch_db.client.indices.get_mapping(index=index_name)
41
41
 
42
42
  # Insert document which will automatically create the index based on the template
43
43
  await dummy_opensearch_db.upsert(vo, document_id, DUMMY_DOCUMENT)
44
44
 
45
45
  # Ensure the result looks as expected and return the mappings
46
- index_mapping = await dummy_opensearch_db.client.indices.get_mapping(index_name)
46
+ index_mapping = await dummy_opensearch_db.client.indices.get_mapping(
47
+ index=index_name
48
+ )
47
49
  assert list(index_mapping) == [index_name]
48
50
  assert list(index_mapping[index_name]) == ["mappings"]
49
51
  return index_mapping[index_name]["mappings"]
@@ -15,7 +15,6 @@ async def pilot_agents_db(tmp_path) -> PilotAgentsDB:
15
15
 
16
16
 
17
17
  async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
18
-
19
18
  async with pilot_agents_db as pilot_agents_db:
20
19
  # Add a pilot reference
21
20
  refs = [f"ref_{i}" for i in range(10)]
@@ -0,0 +1,153 @@
1
+ from __future__ import annotations
2
+
3
+ import re
4
+ from datetime import UTC, datetime, timedelta, timezone
5
+
6
+ import freezegun
7
+ import pytest
8
+ import sqlalchemy
9
+ from sqlalchemy.ext.asyncio import create_async_engine
10
+
11
+ from diracx.testing.time import julian_date, mock_sqlite_time
12
+
13
+ RE_SQLITE_TIME = re.compile(r"(\d{4})-(\d{2})-(\d{2})(?: (\d{2}):(\d{2}):(\d{2}))?")
14
+
15
+
16
+ @pytest.mark.parametrize(
17
+ "dt, expected",
18
+ [
19
+ (datetime(2000, 1, 1, 18, 0, 0, tzinfo=UTC), 2451545.25),
20
+ (datetime(2000, 1, 1, 6, 0, 0, tzinfo=UTC), 2451544.75),
21
+ ],
22
+ )
23
+ def test_julian_date(dt, expected):
24
+ """Test the julian_date function with known values."""
25
+ assert julian_date(dt) == expected
26
+
27
+
28
+ def test_julian_date_non_utc():
29
+ """Test the julian_date function with a non-UTC timezone."""
30
+ non_utc = timezone(timedelta(hours=-5))
31
+ dt = datetime(2000, 1, 1, 18, 0, 0, tzinfo=non_utc)
32
+ jd = julian_date(dt)
33
+ # dt in UTC is 2000-01-01 23:00:00, so fractional day is (23-12)/24 = 11/24
34
+ expected = 2451545 + 11 / 24
35
+ assert abs(jd - expected) < 1e-6
36
+
37
+
38
+ @pytest.mark.parametrize("with_mock", [True, False])
39
+ async def test_freeze_sqlite_datetime(with_mock):
40
+ """Test the SQLite DATETIME() function with freezegun."""
41
+ engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True, echo=True)
42
+ if with_mock:
43
+ sqlalchemy.event.listen(engine.sync_engine, "connect", mock_sqlite_time)
44
+
45
+ async with engine.begin() as conn:
46
+ # DATETIME()
47
+ result = await conn.execute(sqlalchemy.text("SELECT DATETIME()"))
48
+ value = datetime.strptime(result.scalar_one(), "%Y-%m-%d %H:%M:%S").replace(
49
+ tzinfo=UTC
50
+ )
51
+ actual = datetime.now(UTC)
52
+ assert abs(value - actual) < timedelta(seconds=1)
53
+
54
+ if with_mock:
55
+ with freezegun.freeze_time("2000-10-01 12:00:00"):
56
+ result = await conn.execute(sqlalchemy.text("SELECT DATETIME()"))
57
+ value = datetime.strptime(
58
+ result.scalar_one(), "%Y-%m-%d %H:%M:%S"
59
+ ).replace(tzinfo=UTC)
60
+ actual = datetime(2000, 10, 1, 12, 0, 0, tzinfo=UTC)
61
+ assert abs(value - actual) < timedelta(seconds=1)
62
+
63
+ # DATETIME('now')
64
+ result = await conn.execute(sqlalchemy.text("SELECT DATETIME('now')"))
65
+ value = datetime.strptime(result.scalar_one(), "%Y-%m-%d %H:%M:%S").replace(
66
+ tzinfo=UTC
67
+ )
68
+ actual = datetime.now(UTC)
69
+ assert abs(value - actual) < timedelta(seconds=1)
70
+
71
+ if with_mock:
72
+ with freezegun.freeze_time("2000-10-01 12:00:00"):
73
+ result = await conn.execute(sqlalchemy.text("SELECT DATETIME('now')"))
74
+ value = datetime.strptime(
75
+ result.scalar_one(), "%Y-%m-%d %H:%M:%S"
76
+ ).replace(tzinfo=UTC)
77
+ actual = datetime(2000, 10, 1, 12, 0, 0, tzinfo=UTC)
78
+ assert abs(value - actual) < timedelta(seconds=1)
79
+
80
+
81
+ @pytest.mark.parametrize("with_mock", [True, False])
82
+ async def test_freeze_sqlite_julianday(with_mock):
83
+ """Test the SQLite JULIANDAY() function with freezegun."""
84
+ engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True, echo=True)
85
+ if with_mock:
86
+ sqlalchemy.event.listen(engine.sync_engine, "connect", mock_sqlite_time)
87
+
88
+ async with engine.begin() as conn:
89
+ # JULIANDAY()
90
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY()"))
91
+ actual = julian_date(datetime.now(UTC))
92
+ assert abs(result.scalar_one() - actual) < 1e-3
93
+
94
+ if with_mock:
95
+ with freezegun.freeze_time("2000-10-01 12:00:00"):
96
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY()"))
97
+ actual = julian_date(datetime(2000, 10, 1, 12, 0, 0, tzinfo=UTC))
98
+ assert abs(result.scalar_one() - actual) < 1e-3
99
+
100
+ # JULIANDAY('now')
101
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('now')"))
102
+ actual = julian_date(datetime.now(UTC))
103
+ assert abs(result.scalar_one() - actual) < 1e-3
104
+
105
+ if with_mock:
106
+ with freezegun.freeze_time("2000-10-01 12:00:00"):
107
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('now')"))
108
+ actual = julian_date(datetime(2000, 10, 1, 12, 0, 0, tzinfo=UTC))
109
+ assert abs(result.scalar_one() - actual) < 1e-3
110
+
111
+ # JULIANDAY('YYYY-MM-DD')
112
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('1996-05-30')"))
113
+ assert result.scalar_one() == 2450233.5
114
+
115
+ if with_mock:
116
+ with freezegun.freeze_time("2020-12-13 12:34:56"):
117
+ result = await conn.execute(
118
+ sqlalchemy.text("SELECT JULIANDAY('1996-05-30')")
119
+ )
120
+ assert result.scalar_one() == 2450233.5
121
+
122
+ # JULIANDAY('YYYY-MM-DD HH:MM:SS')
123
+ result = await conn.execute(
124
+ sqlalchemy.text("SELECT JULIANDAY('2000-10-01 12:00:00')")
125
+ )
126
+ assert result.scalar_one() == 2451819.0
127
+
128
+ if with_mock:
129
+ with freezegun.freeze_time("2020-12-13 12:34:56"):
130
+ result = await conn.execute(
131
+ sqlalchemy.text("SELECT JULIANDAY('2000-10-01 12:00:00')")
132
+ )
133
+ assert result.scalar_one() == 2451819.0
134
+
135
+ # JULIANDAY('1356')
136
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('1356')"))
137
+ assert result.scalar_one() == 1356.0
138
+
139
+ if with_mock:
140
+ with freezegun.freeze_time("2020-12-13 12:34:56"):
141
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('1356')"))
142
+ assert result.scalar_one() == 1356.0
143
+
144
+ # JULIANDAY('1356.12')
145
+ result = await conn.execute(sqlalchemy.text("SELECT JULIANDAY('1356.12')"))
146
+ assert result.scalar_one() == 1356.12
147
+
148
+ if with_mock:
149
+ with freezegun.freeze_time("2020-12-13 12:34:56"):
150
+ result = await conn.execute(
151
+ sqlalchemy.text("SELECT JULIANDAY('1356.12')")
152
+ )
153
+ assert result.scalar_one() == 1356.12
@@ -1,4 +0,0 @@
1
- [egg_info]
2
- tag_build =
3
- tag_date = 0
4
-
@@ -1,19 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: diracx-db
3
- Version: 0.0.1a44
4
- Summary: TODO
5
- License: GPL-3.0-only
6
- Classifier: Intended Audience :: Science/Research
7
- Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
8
- Classifier: Programming Language :: Python :: 3
9
- Classifier: Topic :: Scientific/Engineering
10
- Classifier: Topic :: System :: Distributed Computing
11
- Requires-Python: >=3.11
12
- Description-Content-Type: text/markdown
13
- Requires-Dist: diracx-core
14
- Requires-Dist: opensearch-py[async]
15
- Requires-Dist: pydantic>=2.10
16
- Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
17
- Requires-Dist: uuid-utils
18
- Provides-Extra: testing
19
- Requires-Dist: diracx-testing; extra == "testing"
@@ -1,53 +0,0 @@
1
- README.md
2
- pyproject.toml
3
- src/diracx/db/__init__.py
4
- src/diracx/db/__main__.py
5
- src/diracx/db/exceptions.py
6
- src/diracx/db/py.typed
7
- src/diracx/db/os/__init__.py
8
- src/diracx/db/os/job_parameters.py
9
- src/diracx/db/os/utils.py
10
- src/diracx/db/sql/__init__.py
11
- src/diracx/db/sql/auth/__init__.py
12
- src/diracx/db/sql/auth/db.py
13
- src/diracx/db/sql/auth/schema.py
14
- src/diracx/db/sql/dummy/__init__.py
15
- src/diracx/db/sql/dummy/db.py
16
- src/diracx/db/sql/dummy/schema.py
17
- src/diracx/db/sql/job/__init__.py
18
- src/diracx/db/sql/job/db.py
19
- src/diracx/db/sql/job/schema.py
20
- src/diracx/db/sql/job_logging/__init__.py
21
- src/diracx/db/sql/job_logging/db.py
22
- src/diracx/db/sql/job_logging/schema.py
23
- src/diracx/db/sql/pilot_agents/__init__.py
24
- src/diracx/db/sql/pilot_agents/db.py
25
- src/diracx/db/sql/pilot_agents/schema.py
26
- src/diracx/db/sql/sandbox_metadata/__init__.py
27
- src/diracx/db/sql/sandbox_metadata/db.py
28
- src/diracx/db/sql/sandbox_metadata/schema.py
29
- src/diracx/db/sql/task_queue/__init__.py
30
- src/diracx/db/sql/task_queue/db.py
31
- src/diracx/db/sql/task_queue/schema.py
32
- src/diracx/db/sql/utils/__init__.py
33
- src/diracx/db/sql/utils/base.py
34
- src/diracx/db/sql/utils/functions.py
35
- src/diracx/db/sql/utils/types.py
36
- src/diracx_db.egg-info/PKG-INFO
37
- src/diracx_db.egg-info/SOURCES.txt
38
- src/diracx_db.egg-info/dependency_links.txt
39
- src/diracx_db.egg-info/entry_points.txt
40
- src/diracx_db.egg-info/requires.txt
41
- src/diracx_db.egg-info/top_level.txt
42
- tests/test_dummy_db.py
43
- tests/auth/test_authorization_flow.py
44
- tests/auth/test_device_flow.py
45
- tests/auth/test_refresh_token.py
46
- tests/jobs/test_job_db.py
47
- tests/jobs/test_job_logging_db.py
48
- tests/jobs/test_sandbox_metadata.py
49
- tests/opensearch/test_connection.py
50
- tests/opensearch/test_index_template.py
51
- tests/opensearch/test_search.py
52
- tests/pilot_agents/__init__.py
53
- tests/pilot_agents/test_pilot_agents_db.py
@@ -1,10 +0,0 @@
1
- [diracx.dbs.os]
2
- JobParametersDB = diracx.db.os:JobParametersDB
3
-
4
- [diracx.dbs.sql]
5
- AuthDB = diracx.db.sql:AuthDB
6
- JobDB = diracx.db.sql:JobDB
7
- JobLoggingDB = diracx.db.sql:JobLoggingDB
8
- PilotAgentsDB = diracx.db.sql:PilotAgentsDB
9
- SandboxMetadataDB = diracx.db.sql:SandboxMetadataDB
10
- TaskQueueDB = diracx.db.sql:TaskQueueDB
@@ -1,8 +0,0 @@
1
- diracx-core
2
- opensearch-py[async]
3
- pydantic>=2.10
4
- sqlalchemy[aiomysql,aiosqlite]>=2
5
- uuid-utils
6
-
7
- [testing]
8
- diracx-testing
@@ -1 +0,0 @@
1
- diracx
File without changes