Missed a spot - add org id to build_uri (#2674)
This commit is contained in:
@@ -72,7 +72,9 @@ class ArtifactManager:
|
||||
path: str | None = None,
|
||||
) -> str:
|
||||
artifact_id = generate_artifact_id()
|
||||
uri = app.STORAGE.build_uri(artifact_id, step, artifact_type)
|
||||
uri = app.STORAGE.build_uri(
|
||||
organization_id=step.organization_id, artifact_id=artifact_id, step=step, artifact_type=artifact_type
|
||||
)
|
||||
return await self._create_artifact(
|
||||
aio_task_primary_key=step.task_id,
|
||||
artifact_id=artifact_id,
|
||||
|
||||
@@ -39,7 +39,7 @@ FILE_EXTENTSION_MAP: dict[ArtifactType, str] = {
|
||||
|
||||
class BaseStorage(ABC):
|
||||
@abstractmethod
|
||||
def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -28,9 +28,9 @@ class LocalStorage(BaseStorage):
|
||||
def __init__(self, artifact_path: str = settings.ARTIFACT_STORAGE_PATH) -> None:
|
||||
self.artifact_path = artifact_path
|
||||
|
||||
def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
file_ext = FILE_EXTENTSION_MAP[artifact_type]
|
||||
return f"file://{self.artifact_path}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
|
||||
return f"file://{self.artifact_path}/{organization_id}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
|
||||
|
||||
async def retrieve_global_workflows(self) -> list[str]:
|
||||
file_path = Path(f"{self.artifact_path}/{settings.ENV}/global_workflows.txt")
|
||||
|
||||
@@ -35,9 +35,9 @@ class S3Storage(BaseStorage):
|
||||
self.async_client = AsyncAWSClient()
|
||||
self.bucket = bucket or settings.AWS_S3_BUCKET_ARTIFACTS
|
||||
|
||||
def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str:
|
||||
file_ext = FILE_EXTENTSION_MAP[artifact_type]
|
||||
return f"s3://{self.bucket}/{self._PATH_VERSION}/{settings.ENV}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
|
||||
return f"s3://{self.bucket}/{self._PATH_VERSION}/{settings.ENV}/{organization_id}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
|
||||
|
||||
async def retrieve_global_workflows(self) -> list[str]:
|
||||
uri = f"s3://{self.bucket}/{settings.ENV}/global_workflows.txt"
|
||||
|
||||
@@ -31,10 +31,15 @@ def local_storage() -> LocalStorage:
|
||||
class TestLocalStorageBuildURIs:
|
||||
def test_build_uri(self, local_storage: LocalStorage) -> None:
|
||||
step = create_fake_step(TEST_STEP_ID)
|
||||
uri = local_storage.build_uri("artifact123", step, ArtifactType.LLM_PROMPT)
|
||||
uri = local_storage.build_uri(
|
||||
organization_id=TEST_ORGANIZATION_ID,
|
||||
artifact_id="artifact123",
|
||||
step=step,
|
||||
artifact_type=ArtifactType.LLM_PROMPT,
|
||||
)
|
||||
assert (
|
||||
uri
|
||||
== f"file://{local_storage.artifact_path}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt"
|
||||
== f"file://{local_storage.artifact_path}/{TEST_ORGANIZATION_ID}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt"
|
||||
)
|
||||
|
||||
def test_build_log_uri(self, local_storage: LocalStorage) -> None:
|
||||
|
||||
@@ -31,10 +31,15 @@ def s3_storage() -> S3Storage:
|
||||
class TestS3StorageBuildURIs:
|
||||
def test_build_uri(self, s3_storage: S3Storage) -> None:
|
||||
step = create_fake_step(TEST_STEP_ID)
|
||||
uri = s3_storage.build_uri("artifact123", step, ArtifactType.LLM_PROMPT)
|
||||
uri = s3_storage.build_uri(
|
||||
organization_id=TEST_ORGANIZATION_ID,
|
||||
artifact_id="artifact123",
|
||||
step=step,
|
||||
artifact_type=ArtifactType.LLM_PROMPT,
|
||||
)
|
||||
assert (
|
||||
uri
|
||||
== f"s3://{TEST_BUCKET}/v1/{settings.ENV}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt"
|
||||
== f"s3://{TEST_BUCKET}/v1/{settings.ENV}/{TEST_ORGANIZATION_ID}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt"
|
||||
)
|
||||
|
||||
def test_build_log_uri(self, s3_storage: S3Storage) -> None:
|
||||
|
||||
Reference in New Issue
Block a user