WIP: Move closer to supporting sandboxes

Author: chrisburr

src/diracx/core/models.py

@@ -112,3 +112,18 @@ class UserInfo(BaseModel):
     preferred_username: str
     dirac_group: str
     vo: str
+
+
+class SandboxChecksum(StrEnum):
+    SHA256 = "sha256"
+
+
+class SandboxFormat(StrEnum):
+    TAR_GZ = "tar.gz"
+
+
+class SandboxInfo(BaseModel):
+    checksum_algorithm: SandboxChecksum
+    checksum: str = Field(pattern=r"^[0-f]{64}$")
+    size: int = Field(ge=1)
+    format: SandboxFormat

src/diracx/core/s3.py

@@ -0,0 +1,80 @@
+"""Utilities for interacting with S3-compatible storage."""
+from __future__ import annotations
+
+import base64
+from typing import TypedDict
+
+from botocore.errorfactory import ClientError
+
+PRESIGNED_URL_TIMEOUT = 5 * 60
+
+
+class S3PresignedPostInfo(TypedDict):
+    url: str
+    fields: dict[str, str]
+
+
+def hack_get_s3_client():
+    # TODO: Use async
+    import boto3
+    from botocore.config import Config
+
+    s3_cred = {
+        "endpoint": "http://christohersmbp4.localdomain:32000",
+        "access_key_id": "console",
+        "secret_access_key": "console123",
+    }
+    bucket_name = "sandboxes"
+    my_config = Config(signature_version="v4")
+    s3 = boto3.client(
+        "s3",
+        endpoint_url=s3_cred["endpoint"],
+        aws_access_key_id=s3_cred["access_key_id"],
+        aws_secret_access_key=s3_cred["secret_access_key"],
+        config=my_config,
+    )
+    try:
+        s3.create_bucket(Bucket=bucket_name)
+    except Exception:
+        pass
+    return s3, bucket_name
+
+
+def s3_object_exists(s3_client, bucket_name, key) -> bool:
+    """Check if an object exists in an S3 bucket."""
+    try:
+        s3_client.head_object(Bucket=bucket_name, Key=key)
+    except ClientError as e:
+        if e.response["Error"]["Code"] != "404":
+            raise
+        return False
+    else:
+        return True
+
+
+def generate_presigned_upload(
+    s3_client, bucket_name, key, checksum_algorithm, checksum, size
+) -> S3PresignedPostInfo:
+    """Generate a presigned URL and fields for uploading a file to S3
+
+    The signature is restricted to only accept data with the given checksum and size.
+    """
+    fields = {
+        "x-amz-checksum-algorithm": checksum_algorithm,
+        f"x-amz-checksum-{checksum_algorithm}": b16_to_b64(checksum),
+    }
+    conditions = [["content-length-range", size, size]] + [
+        {k: v} for k, v in fields.items()
+    ]
+    return s3_client.generate_presigned_post(
+        Bucket=bucket_name,
+        Key=key,
+        Fields=fields,
+        Conditions=conditions,
+        ExpiresIn=PRESIGNED_URL_TIMEOUT,
+    )
+
+
+def b16_to_b64(hex_string: str) -> str:
+    """Convert hexadecimal encoded data to base64 encoded data"""
+    return base64.b64encode(base64.b16decode(hex_string.upper())).decode()

src/diracx/db/sql/sandbox_metadata/db.py

@@ -1,96 +1,86 @@
-""" SandboxMetadataDB frontend
-"""
-
 from __future__ import annotations
 
-import datetime
-
 import sqlalchemy
 
-from diracx.db.sql.utils import BaseSQLDB
+from diracx.core.models import SandboxInfo, UserInfo
+from diracx.db.sql.utils import BaseSQLDB, utcnow
 
 from .schema import Base as SandboxMetadataDBBase
 from .schema import sb_Owners, sb_SandBoxes
 
 # In legacy DIRAC the SEName column was used to support multiple different
 # storage backends. This is no longer the case, so we hardcode the value to
 # S3 to represent the new DiracX system.
-SE_NAME = "S3"
+SE_NAME = "ProductionSandboxSE"
+PFN_PREFIX = "/S3/"
 
 
 class SandboxMetadataDB(BaseSQLDB):
     metadata = SandboxMetadataDBBase.metadata
 
-    async def _get_put_owner(self, owner: str, owner_group: str) -> int:
-        """adds a new owner/ownerGroup pairs, while returning their ID if already existing
-
-        Args:
-            owner (str): user name
-            owner_group (str): group of the owner
-        """
+    async def upsert_owner(self, user: UserInfo) -> int:
+        """Get the id of the owner from the database"""
+        # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
         stmt = sqlalchemy.select(sb_Owners.OwnerID).where(
-            sb_Owners.Owner == owner, sb_Owners.OwnerGroup == owner_group
+            sb_Owners.Owner == user.preferred_username,
+            sb_Owners.OwnerGroup == user.dirac_group,
+            # TODO: Add VO
         )
         result = await self.conn.execute(stmt)
         if owner_id := result.scalar_one_or_none():
             return owner_id
 
-        stmt = sqlalchemy.insert(sb_Owners).values(Owner=owner, OwnerGroup=owner_group)
+        stmt = sqlalchemy.insert(sb_Owners).values(
+            Owner=user.preferred_username,
+            OwnerGroup=user.dirac_group,
+        )
         result = await self.conn.execute(stmt)
         return result.lastrowid
 
-    async def insert(
-        self, owner: str, owner_group: str, sb_SE: str, se_PFN: str, size: int = 0
-    ) -> int:
-        """inserts a new sandbox in SandboxMetadataDB
-        this is "equivalent" of DIRAC registerAndGetSandbox
-
-        Args:
-            owner (str): user name_
-            owner_group (str): groupd of the owner
-            sb_SE (str): _description_
-            sb_PFN (str): _description_
-            size (int, optional): _description_. Defaults to 0.
-        """
-        owner_id = await self._get_put_owner(owner, owner_group)
+    @staticmethod
+    def get_pfn(bucket_name: str, user: UserInfo, sandbox_info: SandboxInfo) -> str:
+        """Get the sandbox's user namespaced and content addressed PFN"""
+        parts = [
+            "S3",
+            bucket_name,
+            user.vo,
+            user.dirac_group,
+            user.preferred_username,
+            f"{sandbox_info.checksum_algorithm}:{sandbox_info.checksum}.{sandbox_info.format}",
+        ]
+        return "/".join(parts)
+
+    async def insert_sandbox(self, user: UserInfo, pfn: str, size: int):
+        """Add a new sandbox in SandboxMetadataDB"""
+        # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
+        owner_id = await self.upsert_owner(user)
         stmt = sqlalchemy.insert(sb_SandBoxes).values(
-            OwnerId=owner_id, SEName=sb_SE, SEPFN=se_PFN, Bytes=size
+            OwnerId=owner_id, SEName=SE_NAME, SEPFN=pfn, Bytes=size
         )
         try:
             result = await self.conn.execute(stmt)
-            return result.lastrowid
         except sqlalchemy.exc.IntegrityError:
-            # it is a duplicate, try to retrieve SBiD
-            stmt: sqlalchemy.Executable = sqlalchemy.select(sb_SandBoxes.SBId).where(  # type: ignore[no-redef]
-                sb_SandBoxes.SEPFN == se_PFN,
-                sb_SandBoxes.SEName == sb_SE,
-                sb_SandBoxes.OwnerId == owner_id,
-            )
-            result = await self.conn.execute(stmt)
-            sb_ID = result.scalar_one()
-            stmt: sqlalchemy.Executable = (  # type: ignore[no-redef]
-                sqlalchemy.update(sb_SandBoxes)
-                .where(sb_SandBoxes.SBId == sb_ID)
-                .values(LastAccessTime=datetime.datetime.utcnow())
-            )
-            await self.conn.execute(stmt)
-            return sb_ID
-
-    async def exists_and_assigned(self, name: str) -> bool:
-        """Checks if a sandbox exists and has been assigned
+            await self.update_sandbox_last_access_time(pfn)
+        else:
+            assert result.rowcount == 1
+
+    async def update_sandbox_last_access_time(self, pfn: str) -> None:
+        stmt = (
+            sqlalchemy.update(sb_SandBoxes)
+            .where(sb_SandBoxes.SEName == SE_NAME, sb_SandBoxes.SEPFN == pfn)
+            .values(LastAccessTime=utcnow())
+        )
+        result = await self.conn.execute(stmt)
+        assert result.rowcount == 1
 
-        As sandboxes are registered in the DB before uploading to the storage
-        backend we can't on their existence in the database to determine if
-        they have been uploaded. Instead we check if the sandbox has been
-        assigned to a job. If it has then we know it has been uploaded and we
-        can avoid communicating with the storage backend.
-        """
+    async def sandbox_is_assigned(self, pfn: str) -> bool:
+        """Checks if a sandbox exists and has been assigned."""
         stmt: sqlalchemy.Executable = sqlalchemy.select(sb_SandBoxes.Assigned).where(
-            sb_SandBoxes.SEName == SE_NAME,
-            sb_SandBoxes.SEPFN == name,
+            sb_SandBoxes.SEName == SE_NAME, sb_SandBoxes.SEPFN == pfn
         )
         result = await self.conn.execute(stmt)
-        return result.scalar_one()
+        is_assigned = result.scalar_one()
+        return is_assigned
 
     async def delete(self, sandbox_ids: list[int]) -> bool:
         stmt: sqlalchemy.Executable = sqlalchemy.delete(sb_SandBoxes).where(