diracx-db 0.0.1a50__py3-none-any.whl → 0.0.2__py3-none-any.whl

diracx/db/__main__.py

@@ -39,6 +39,7 @@ async def init_sql():
                 if db._db_url.startswith("sqlite"):
                     await conn.exec_driver_sql("PRAGMA foreign_keys=ON")
                 await conn.run_sync(db.metadata.create_all)
+                await db.post_create(conn)
 
 
 async def init_os():

diracx/db/sql/auth/db.py

@@ -1,16 +1,21 @@
 from __future__ import annotations
 
+import logging
 import secrets
+from datetime import UTC, datetime
+from itertools import pairwise
 
-from sqlalchemy import insert, select, update
+from dateutil.rrule import MONTHLY, rrule
+from sqlalchemy import insert, select, text, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
+from sqlalchemy.ext.asyncio import AsyncConnection
 from uuid_utils import UUID, uuid7
 
 from diracx.core.exceptions import (
     AuthorizationError,
     TokenNotFoundError,
 )
-from diracx.db.sql.utils import BaseSQLDB, hash, substract_date
+from diracx.db.sql.utils import BaseSQLDB, hash, substract_date, uuid7_from_datetime
 
 from .schema import (
     AuthorizationFlows,
@@ -25,10 +30,72 @@ from .schema import Base as AuthDBBase
 USER_CODE_ALPHABET = "BCDFGHJKLMNPQRSTVWXZ"
 MAX_RETRY = 5
 
+logger = logging.getLogger(__name__)
+
 
 class AuthDB(BaseSQLDB):
     metadata = AuthDBBase.metadata
 
+    @classmethod
+    async def post_create(cls, conn: AsyncConnection) -> None:
+        """Create partitions if it is a MySQL DB and it does not have
+        it yet and the table does not have any data yet.
+        We do this as a post_create step as sqlalchemy does not support
+        partition so well.
+        """
+        if conn.dialect.name == "mysql":
+            check_partition_query = text(
+                "SELECT PARTITION_NAME FROM information_schema.partitions "
+                "WHERE TABLE_NAME = 'RefreshTokens' AND PARTITION_NAME is not NULL"
+            )
+            partition_names = (await conn.execute(check_partition_query)).all()
+
+            if not partition_names:
+                # Create a monthly partition from today until 2 years
+                # The partition are named p_<year>_<month>
+                start_date = datetime.now(tz=UTC).replace(
+                    day=1, hour=0, minute=0, second=0, microsecond=0
+                )
+                end_date = start_date.replace(year=start_date.year + 2)
+
+                dates = [
+                    dt for dt in rrule(MONTHLY, dtstart=start_date, until=end_date)
+                ]
+
+                partition_list = []
+                for name, limit in pairwise(dates):
+                    partition_list.append(
+                        f"PARTITION p_{name.year}_{name.month} "
+                        f"VALUES LESS THAN ('{str(uuid7_from_datetime(limit, randomize=False)).replace('-', '')}')"
+                    )
+                partition_list.append("PARTITION p_future VALUES LESS THAN (MAXVALUE)")
+
+                alter_query = text(
+                    f"ALTER TABLE RefreshTokens PARTITION BY RANGE COLUMNS (JTI) ({','.join(partition_list)})"
+                )
+
+                check_table_empty_query = text("SELECT * FROM RefreshTokens LIMIT 1")
+                refresh_table_content = (
+                    await conn.execute(check_table_empty_query)
+                ).all()
+                if refresh_table_content:
+                    logger.warning(
+                        "RefreshTokens table not empty. Run the following query yourself"
+                    )
+                    logger.warning(alter_query)
+                    return
+
+                await conn.execute(alter_query)
+
+                partition_names = (
+                    await conn.execute(
+                        check_partition_query, {"table_name": "RefreshTokens"}
+                    )
+                ).all()
+                assert partition_names, (
+                    f"There should be partitions now {partition_names}"
+                )
+
     async def device_flow_validate_user_code(
         self, user_code: str, max_validity: int
     ) -> str:

diracx/db/sql/utils/base.py

@@ -155,6 +155,11 @@ class BaseSQLDB(metaclass=ABCMeta):
                     raise
         return db_urls
 
+    @classmethod
+    async def post_create(cls, conn: AsyncConnection) -> None:
+        """Execute actions after the schema has been created."""
+        return
+
     @classmethod
     def transaction(cls) -> Self:
         raise NotImplementedError("This should never be called")

{diracx_db-0.0.1a50.dist-info → diracx_db-0.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-db
-Version: 0.0.1a50
+Version: 0.0.2
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -12,8 +12,11 @@ Requires-Python: >=3.11
 Requires-Dist: diracx-core
 Requires-Dist: opensearch-py[async]
 Requires-Dist: pydantic>=2.10
+Requires-Dist: python-dateutil
 Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Requires-Dist: uuid-utils
 Provides-Extra: testing
 Requires-Dist: diracx-testing; extra == 'testing'
 Requires-Dist: freezegun; extra == 'testing'
+Provides-Extra: types
+Requires-Dist: types-python-dateutil; extra == 'types'

{diracx_db-0.0.1a50.dist-info → diracx_db-0.0.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 diracx/db/__init__.py,sha256=2oeUeVwZq53bo_ZOflEYZsBn7tcR5Tzb2AIu0TAWELM,109
-diracx/db/__main__.py,sha256=3yaUP1ig-yaPSQM4wy6CtSXXHivQg-hIz2FeBt7joBc,1714
+diracx/db/__main__.py,sha256=6YlmpiU1cLLHjKLy1DfdEOQUyvSla-MbJsJ7aQwAOVs,1757
 diracx/db/exceptions.py,sha256=1nn-SZLG-nQwkxbvHjZqXhE5ouzWj1f3qhSda2B4ZEg,83
 diracx/db/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
@@ -7,7 +7,7 @@ diracx/db/os/job_parameters.py,sha256=3w_CeA2z-cY5pWwXkGu-Fod27FobbUXuwVKK-jN037
 diracx/db/os/utils.py,sha256=V4T-taos64SFNcorfIr7mq5l5y88K6TzyCj1YqWk8VI,11562
 diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/auth/db.py,sha256=QJtBqMrhOf97UvMG0WpyjsgIRiu19v04FoDzXAyXtT0,8952
+diracx/db/sql/auth/db.py,sha256=sGXKFLy1iCdrkg1952_F-c7nRSIx4F1JuDloIV6AXN8,11814
 diracx/db/sql/auth/schema.py,sha256=9fUV7taDPnoAcoiwRAmQraOmF2Ytoizjs2TFvN7zsVs,3132
 diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/dummy/db.py,sha256=MKSUSJI1BlRgK08tjCfkCkOz02asvJAeBw60pAdiGV8,1212
@@ -28,10 +28,10 @@ diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 diracx/db/sql/task_queue/db.py,sha256=2qul1D2tX2uCI92N591WK5xWHakG0pNibzDwKQ7W-I8,6246
 diracx/db/sql/task_queue/schema.py,sha256=5efAgvNYRkLlaJ2NzRInRfmVa3tyIzQu2l0oRPy4Kzw,3258
 diracx/db/sql/utils/__init__.py,sha256=k1DI4Idlqv36pXn2BhQysb947Peio9DnYaePslkTpUQ,685
-diracx/db/sql/utils/base.py,sha256=ZqqBYSKsRz2kLezOB7CSbxPFQVcozmaxNbrbXuGCKPs,16904
+diracx/db/sql/utils/base.py,sha256=cXwAzWtYtEkukjHSBGaGNtXSozHfXHpZ9p9eXdQTpDY,17065
 diracx/db/sql/utils/functions.py,sha256=_E4tc9Gti6LuSh7QEyoqPJSvCuByVqvRenOXCzxsulE,4014
 diracx/db/sql/utils/types.py,sha256=KNZWJfpvHTjfIPg6Nn7zY-rS0q3ybnirHcTcLAYSYbE,5118
-diracx_db-0.0.1a50.dist-info/METADATA,sha256=tF7YN-0vEiwoSOdb40_lPaZ5n191Tn5NaY_z3ifD2g4,675
-diracx_db-0.0.1a50.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-diracx_db-0.0.1a50.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
-diracx_db-0.0.1a50.dist-info/RECORD,,
+diracx_db-0.0.2.dist-info/METADATA,sha256=n8QuMXRrA6-iLKw14y29Jj2WxBzQZ2plW5qMngBnEqs,780
+diracx_db-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+diracx_db-0.0.2.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
+diracx_db-0.0.2.dist-info/RECORD,,