Skip to content

Commit e360a9f

Browse files
serinamarieÅsmund Østvold
authored and
Åsmund Østvold
committed
Add /artifacts/latest/count route (PrefectHQ#9090)
1 parent ad7ef4a commit e360a9f

File tree

4 files changed

+107
-0
lines changed

4 files changed

+107
-0
lines changed

src/prefect/server/api/artifacts.py

+23
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,29 @@ async def count_artifacts(
161161
)
162162

163163

164+
@router.post("/latest/count")
165+
async def count_latest_artifacts(
166+
artifacts: filters.ArtifactCollectionFilter = None,
167+
flow_runs: filters.FlowRunFilter = None,
168+
task_runs: filters.TaskRunFilter = None,
169+
flows: filters.FlowFilter = None,
170+
deployments: filters.DeploymentFilter = None,
171+
db: PrefectDBInterface = Depends(provide_database_interface),
172+
) -> int:
173+
"""
174+
Count artifacts from the database.
175+
"""
176+
async with db.session_context() as session:
177+
return await models.artifacts.count_latest_artifacts(
178+
session=session,
179+
artifact_filter=artifacts,
180+
flow_run_filter=flow_runs,
181+
task_run_filter=task_runs,
182+
flow_filter=flows,
183+
deployment_filter=deployments,
184+
)
185+
186+
164187
@router.patch("/{id}", status_code=204)
165188
async def update_artifact(
166189
artifact: actions.ArtifactUpdate,

src/prefect/server/models/artifacts.py

+34
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,40 @@ async def count_artifacts(
379379
return result.scalar_one()
380380

381381

382+
@inject_db
383+
async def count_latest_artifacts(
384+
session: sa.orm.Session,
385+
db: PrefectDBInterface,
386+
artifact_filter: filters.ArtifactCollectionFilter = None,
387+
flow_run_filter: filters.FlowRunFilter = None,
388+
task_run_filter: filters.TaskRunFilter = None,
389+
deployment_filter: filters.DeploymentFilter = None,
390+
flow_filter: filters.FlowFilter = None,
391+
) -> int:
392+
"""
393+
Counts artifacts.
394+
Args:
395+
session: A database session
396+
artifact_filter: Only select artifacts matching this filter
397+
flow_run_filter: Only select artifacts whose flow runs matching this filter
398+
task_run_filter: Only select artifacts whose task runs matching this filter
399+
"""
400+
query = sa.select(sa.func.count(db.ArtifactCollection.id))
401+
402+
query = await _apply_artifact_filters(
403+
query,
404+
db=db,
405+
artifact_filter=artifact_filter,
406+
flow_run_filter=flow_run_filter,
407+
task_run_filter=task_run_filter,
408+
deployment_filter=deployment_filter,
409+
flow_filter=flow_filter,
410+
)
411+
412+
result = await session.execute(query)
413+
return result.scalar_one()
414+
415+
382416
@inject_db
383417
async def update_artifact(
384418
session: sa.orm.Session,

tests/server/api/test_artifacts.py

+26
Original file line numberDiff line numberDiff line change
@@ -727,6 +727,32 @@ async def test_count_artifacts_by_deployment(self, flow_artifacts, client):
727727
json = response.json()
728728
assert json == 2
729729

730+
async def test_counting_latest_artifacts_by_flow_name(self, flow_artifacts, client):
731+
flow_name = flow_artifacts[0]["name"]
732+
flow_filter = dict(
733+
flows=schemas.filters.FlowFilter(
734+
name=schemas.filters.FlowFilterName(any_=[flow_name])
735+
).dict(json_compatible=True)
736+
)
737+
response = await client.post("/artifacts/latest/count", json=flow_filter)
738+
assert response.status_code == status.HTTP_200_OK
739+
json = response.json()
740+
assert json == 1
741+
742+
async def test_counting_latest_artifacts_by_deployment(
743+
self, flow_artifacts, client
744+
):
745+
deployment_id = flow_artifacts[3]
746+
deployment_filter = dict(
747+
deployments=schemas.filters.DeploymentFilter(
748+
id=schemas.filters.DeploymentFilterId(any_=[deployment_id])
749+
).dict(json_compatible=True)
750+
)
751+
response = await client.post("/artifacts/latest/count", json=deployment_filter)
752+
assert response.status_code == status.HTTP_200_OK
753+
json = response.json()
754+
assert json == 1
755+
730756

731757
class TestUpdateArtifact:
732758
async def test_update_artifact_succeeds(self, artifact, client):

tests/server/models/test_artifacts.py

+24
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,30 @@ async def test_counting_artifacts_by_deployment(
289289
)
290290
assert result == 2
291291

292+
async def test_counting_latest_artifacts_by_flow_name(
293+
self, flow_artifacts, session
294+
):
295+
flow_name = flow_artifacts[0].name
296+
result = await models.artifacts.count_latest_artifacts(
297+
session=session,
298+
flow_filter=schemas.filters.FlowFilter(
299+
name=schemas.filters.FlowFilterName(any_=[flow_name])
300+
),
301+
)
302+
assert result == 1
303+
304+
async def test_counting_latest_artifacts_by_deployment(
305+
self, deployment_artifacts, session
306+
):
307+
deployment_id = deployment_artifacts[0].deployment_id
308+
result = await models.artifacts.count_latest_artifacts(
309+
session=session,
310+
deployment_filter=schemas.filters.DeploymentFilter(
311+
id=schemas.filters.DeploymentFilterId(any_=[deployment_id])
312+
),
313+
)
314+
assert result == 1
315+
292316

293317
class TestUpdateArtifacts:
294318
@pytest.fixture

0 commit comments

Comments
 (0)