Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add /artifacts/latest/count route #9090

Merged
merged 8 commits into from
Apr 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/prefect/server/api/artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,29 @@ async def count_artifacts(
)


@router.post("/latest/count")
async def count_latest_artifacts(
artifacts: filters.ArtifactCollectionFilter = None,
flow_runs: filters.FlowRunFilter = None,
task_runs: filters.TaskRunFilter = None,
flows: filters.FlowFilter = None,
deployments: filters.DeploymentFilter = None,
db: PrefectDBInterface = Depends(provide_database_interface),
) -> int:
"""
Count artifacts from the database.
"""
async with db.session_context() as session:
return await models.artifacts.count_latest_artifacts(
session=session,
artifact_filter=artifacts,
flow_run_filter=flow_runs,
task_run_filter=task_runs,
flow_filter=flows,
deployment_filter=deployments,
)


@router.patch("/{id}", status_code=204)
async def update_artifact(
artifact: actions.ArtifactUpdate,
Expand Down
34 changes: 34 additions & 0 deletions src/prefect/server/models/artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,6 +379,40 @@ async def count_artifacts(
return result.scalar_one()


@inject_db
async def count_latest_artifacts(
session: sa.orm.Session,
db: PrefectDBInterface,
artifact_filter: filters.ArtifactCollectionFilter = None,
flow_run_filter: filters.FlowRunFilter = None,
task_run_filter: filters.TaskRunFilter = None,
deployment_filter: filters.DeploymentFilter = None,
flow_filter: filters.FlowFilter = None,
) -> int:
"""
Counts artifacts.
Args:
session: A database session
artifact_filter: Only select artifacts matching this filter
flow_run_filter: Only select artifacts whose flow runs matching this filter
task_run_filter: Only select artifacts whose task runs matching this filter
"""
query = sa.select(sa.func.count(db.ArtifactCollection.id))

query = await _apply_artifact_filters(
query,
db=db,
artifact_filter=artifact_filter,
flow_run_filter=flow_run_filter,
task_run_filter=task_run_filter,
deployment_filter=deployment_filter,
flow_filter=flow_filter,
)

result = await session.execute(query)
return result.scalar_one()


@inject_db
async def update_artifact(
session: sa.orm.Session,
Expand Down
26 changes: 26 additions & 0 deletions tests/server/api/test_artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -727,6 +727,32 @@ async def test_count_artifacts_by_deployment(self, flow_artifacts, client):
json = response.json()
assert json == 2

async def test_counting_latest_artifacts_by_flow_name(self, flow_artifacts, client):
flow_name = flow_artifacts[0]["name"]
flow_filter = dict(
flows=schemas.filters.FlowFilter(
name=schemas.filters.FlowFilterName(any_=[flow_name])
).dict(json_compatible=True)
)
response = await client.post("/artifacts/latest/count", json=flow_filter)
assert response.status_code == status.HTTP_200_OK
json = response.json()
assert json == 1

async def test_counting_latest_artifacts_by_deployment(
self, flow_artifacts, client
):
deployment_id = flow_artifacts[3]
deployment_filter = dict(
deployments=schemas.filters.DeploymentFilter(
id=schemas.filters.DeploymentFilterId(any_=[deployment_id])
).dict(json_compatible=True)
)
response = await client.post("/artifacts/latest/count", json=deployment_filter)
assert response.status_code == status.HTTP_200_OK
json = response.json()
assert json == 1


class TestUpdateArtifact:
async def test_update_artifact_succeeds(self, artifact, client):
Expand Down
24 changes: 24 additions & 0 deletions tests/server/models/test_artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,30 @@ async def test_counting_artifacts_by_deployment(
)
assert result == 2

async def test_counting_latest_artifacts_by_flow_name(
self, flow_artifacts, session
):
flow_name = flow_artifacts[0].name
result = await models.artifacts.count_latest_artifacts(
session=session,
flow_filter=schemas.filters.FlowFilter(
name=schemas.filters.FlowFilterName(any_=[flow_name])
),
)
assert result == 1

async def test_counting_latest_artifacts_by_deployment(
self, deployment_artifacts, session
):
deployment_id = deployment_artifacts[0].deployment_id
result = await models.artifacts.count_latest_artifacts(
session=session,
deployment_filter=schemas.filters.DeploymentFilter(
id=schemas.filters.DeploymentFilterId(any_=[deployment_id])
),
)
assert result == 1


class TestUpdateArtifacts:
@pytest.fixture
Expand Down