Skip to content
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
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,22 @@ def _coerce_unrecognized_expression_to_none(value: Any) -> Any:
]
"""``AssetExpression | None`` that degrades a legacy/unrecognized stored shape to ``None``."""


_maybe_asset_expression_adapter: TypeAdapter[Union[AssetExpression, None]] = TypeAdapter(MaybeAssetExpression)


def validate_asset_expression(value: Any) -> AssetExpression | None:
"""
Validate a raw ``asset_expression`` tree into the typed model, the way a response field does.

A route that rewrites the tree after its response model is built -- redacting the assets the
caller may not read, for instance -- must assign the result of this rather than the raw tree.
FastAPI does not validate the nested models of an already-built response, so a raw tree would
bypass ``MaybeAssetExpression`` and reach the client verbatim, legacy shapes included.
"""
return _maybe_asset_expression_adapter.validate_python(value)


# Common Bulk Data Models
T = TypeVar("T")
K = TypeVar("K")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,15 @@ class PartitionedDagRunCollectionResponse(BaseModel):
class PartitionedDagRunAssetResponse(BaseModel):
"""Asset info within a partitioned Dag run detail."""

asset_id: int
asset_name: str
asset_uri: str
asset_id: int | None
asset_name: str | None
asset_uri: str | None
hidden: bool = False
"""True when the caller may not read this asset: ``asset_id``, ``asset_name`` and ``asset_uri``
are blanked and the UI must not link to the asset. The counts and partition keys stay, so the
run's progress is still readable -- they describe the Dag's own schedule, which the caller can
already read, rather than the asset's identity. Mirrors the redaction
``redact_asset_expression`` applies to the sibling ``asset_expression`` field."""
received: bool
received_count: int
required_count: int
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5531,14 +5531,24 @@ components:
PartitionedDagRunAssetResponse:
properties:
asset_id:
type: integer
anyOf:
- type: integer
- type: 'null'
title: Asset Id
asset_name:
type: string
anyOf:
- type: string
- type: 'null'
title: Asset Name
asset_uri:
type: string
anyOf:
- type: string
- type: 'null'
title: Asset Uri
hidden:
type: boolean
title: Hidden
default: false
received:
type: boolean
title: Received
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from sqlalchemy import delete, func, insert, select, update

from airflow.api.common import delete_dag as delete_dag_module
from airflow.api_fastapi.common.asset_expression import redact_asset_expression
from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
from airflow.api_fastapi.common.db.common import SessionDep, apply_filters_to_select, paginated_select
from airflow.api_fastapi.common.db.dags import eager_load_teams, generate_dag_with_latest_run_query
Expand Down Expand Up @@ -58,7 +59,11 @@
)
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
from airflow.api_fastapi.core_api.datamodels.common import BulkBody, BulkResponse
from airflow.api_fastapi.core_api.datamodels.common import (
BulkBody,
BulkResponse,
validate_asset_expression,
)
from airflow.api_fastapi.core_api.datamodels.dags import (
BulkDAGBody,
DAGCollectionResponse,
Expand All @@ -71,6 +76,7 @@
from airflow.api_fastapi.core_api.security import (
EditableDagsFilterDep,
GetUserDep,
ReadableAssetsFilterDep,
ReadableDagsFilterDep,
requires_access_dag,
requires_access_dag_bulk,
Expand Down Expand Up @@ -239,7 +245,11 @@ def get_dag(
dependencies=[Depends(requires_access_dag(method="GET"))],
)
def get_dag_details(
dag_id: str, session: SessionDep, dag_bag: DagBagDep, user: GetUserDep
dag_id: str,
session: SessionDep,
dag_bag: DagBagDep,
user: GetUserDep,
readable_assets_filter: ReadableAssetsFilterDep,
) -> DAGDetailsResponse:
"""Get details of Dag."""
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
Expand Down Expand Up @@ -280,7 +290,16 @@ def get_dag_details(
setattr(dag_model, "active_runs_count", active_runs_count)
setattr(dag_model, "latest_dag_version", latest_dag_version)

return DAGDetailsResponse.model_validate(dag_model)
details = DAGDetailsResponse.model_validate(dag_model)
# The redaction lands on the response, not on `dag_model`: that is a live ORM instance whose
# session commits when the request ends, so assigning the blanked leaves back onto it would
# persist them.
details.asset_expression = validate_asset_expression(
redact_asset_expression(
dag_model.asset_expression, readable_asset_ids=readable_assets_filter.value or set()
)
)
return details


@dags_router.patch(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from sqlalchemy.orm import defaultload

from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
from airflow.api_fastapi.common.asset_expression import redact_asset_expression
from airflow.api_fastapi.common.dagbag import DagBagDep
from airflow.api_fastapi.common.db.common import (
SessionDep,
Expand Down Expand Up @@ -79,6 +80,7 @@
from airflow.api_fastapi.core_api.routes.public.hitl import HITLDetailView
from airflow.api_fastapi.core_api.security import (
GetUserDep,
ReadableAssetsFilterDep,
ReadableDagsFilterDep,
requires_access_dag,
)
Expand Down Expand Up @@ -155,6 +157,7 @@ def get_dags(
],
has_pending_actions: QueryPendingActionsFilter,
readable_dags_filter: ReadableDagsFilterDep,
readable_assets_filter: ReadableAssetsFilterDep,
session: SessionDep,
dag_bag: DagBagDep,
user: GetUserDep,
Expand Down Expand Up @@ -311,6 +314,7 @@ def get_dags(
# Build the dict dynamically from DAGResponse.model_fields so that new fields
# added to DAGResponse are picked up automatically without code changes here.
dag_runs_by_dag_id: dict[str, DAGWithLatestDagRunsResponse] = {}
readable_asset_ids = readable_assets_filter.value or set()
for dag in dags:
dag_data = {
DAG_ALIAS_MAPPING.get(field_name, field_name): getattr(
Expand All @@ -320,7 +324,9 @@ def get_dags(
}
dag_data.update(
{
"asset_expression": dag.asset_expression,
"asset_expression": redact_asset_expression(
dag.asset_expression, readable_asset_ids=readable_asset_ids
),
"latest_dag_runs": [],
"has_unfinished_runs": has_unfinished_runs_by_dag_id[dag.dag_id],
"pending_actions": pending_actions_by_dag_id[dag.dag_id],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from fastapi import Depends, HTTPException, status
from sqlalchemy import and_, select

from airflow.api_fastapi.common.asset_expression import redact_asset_expression
from airflow.api_fastapi.common.db.common import SessionDep, apply_filters_to_select, paginated_select
from airflow.api_fastapi.common.parameters import (
QueryLimit,
Expand All @@ -42,6 +43,7 @@
)
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import (
ReadableAssetsFilterDep,
ReadableDagsFilterDep,
requires_access_asset,
requires_access_dag,
Expand Down Expand Up @@ -251,6 +253,7 @@ def get_partitioned_dag_runs(
limit: QueryLimit,
offset: QueryOffset,
readable_dags_filter: ReadableDagsFilterDep,
readable_assets_filter: ReadableAssetsFilterDep,
dag_id: QueryPartitionedDagRunDagIdFilter,
has_created_dag_run_id: QueryPartitionedDagRunHasCreatedDagRunIdFilter,
) -> PartitionedDagRunCollectionResponse:
Expand Down Expand Up @@ -351,7 +354,11 @@ def get_partitioned_dag_runs(

asset_expressions: dict[str, dict | None] | None = None
if dag_id.value is None:
asset_expressions = {dm.dag_id: dm.asset_expression for dm in dag_models.values()}
readable_asset_ids = readable_assets_filter.value or set()
asset_expressions = {
dm.dag_id: redact_asset_expression(dm.asset_expression, readable_asset_ids=readable_asset_ids)
for dm in dag_models.values()
}

model_data: dict[str, Any] = {
"partitioned_dag_runs": results,
Expand All @@ -370,6 +377,7 @@ def get_pending_partitioned_dag_run(
dag_id: str,
partition_key: str,
session: SessionDep,
readable_assets_filter: ReadableAssetsFilterDep,
partitioned_dag_run_id: int | None = None,
) -> PartitionedDagRunDetailResponse:
"""Return full details for pending PartitionedDagRun."""
Expand Down Expand Up @@ -437,8 +445,13 @@ def get_pending_partitioned_dag_run(
has_rollup_mappers = dag_model is not None and dag_model.has_rollup_mappers
rollup_timetable = load_partitioned_timetable(dag_id, session) if has_rollup_mappers else None

readable_asset_ids = readable_assets_filter.value or set()

assets = []
for asset_row in asset_rows:
# The row itself is still needed in full below: the rollup mapper is keyed by the asset's
# real name and uri, and the received-key lookup by its real id. Only the response is blanked.
is_hidden = asset_row.id not in readable_asset_ids
received_keys = sorted(received_keys_by_asset.get(asset_row.id, set()))
required_keys: list[str] = [partition_key]
is_rollup = (
Expand Down Expand Up @@ -480,9 +493,10 @@ def get_pending_partitioned_dag_run(
received_count = 1 if received_keys else 0
assets.append(
PartitionedDagRunAssetResponse(
asset_id=asset_row.id,
asset_name=asset_row.name,
asset_uri=asset_row.uri,
asset_id=None if is_hidden else asset_row.id,
asset_name=None if is_hidden else asset_row.name,
asset_uri=None if is_hidden else asset_row.uri,
hidden=is_hidden,
received=received_count >= required_count and required_count > 0,
received_count=received_count,
required_count=required_count,
Expand All @@ -496,7 +510,11 @@ def get_pending_partitioned_dag_run(

total_received = sum(a.received_count for a in assets)
total_required = sum(a.required_count for a in assets)
asset_expression = dag_model.asset_expression if dag_model is not None else None
asset_expression = (
redact_asset_expression(dag_model.asset_expression, readable_asset_ids=readable_asset_ids)
if dag_model is not None
else None
)

model_data: dict[str, Any] = {
"id": partitioned_dag_run.id,
Expand Down
32 changes: 29 additions & 3 deletions airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12954,17 +12954,43 @@ export const $OklchColor = {
export const $PartitionedDagRunAssetResponse = {
properties: {
asset_id: {
type: 'integer',
anyOf: [
{
type: 'integer'
},
{
type: 'null'
}
],
title: 'Asset Id'
},
asset_name: {
type: 'string',
anyOf: [
{
type: 'string'
},
{
type: 'null'
}
],
title: 'Asset Name'
},
asset_uri: {
type: 'string',
anyOf: [
{
type: 'string'
},
{
type: 'null'
}
],
title: 'Asset Uri'
},
hidden: {
type: 'boolean',
title: 'Hidden',
default: false
},
received: {
type: 'boolean',
title: 'Received'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3317,9 +3317,10 @@ export type OklchColor = string;
* Asset info within a partitioned Dag run detail.
*/
export type PartitionedDagRunAssetResponse = {
asset_id: number;
asset_name: string;
asset_uri: string;
asset_id: number | null;
asset_name: string | null;
asset_uri: string | null;
hidden?: boolean;
received: boolean;
received_count: number;
required_count: number;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,13 @@
* specific language governing permissions and limitations
* under the License.
*/
import { render, waitFor } from "@testing-library/react";
import "@testing-library/jest-dom/vitest";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import { expect, it, vi } from "vitest";

import { PartitionedDagRunService } from "openapi/requests/services.gen";

import i18n from "src/i18n/config";
import { Wrapper } from "src/utils/Wrapper";

import { AssetProgressCell } from "./AssetProgressCell";
Expand Down Expand Up @@ -68,3 +70,62 @@ it("fetches separate details for pending runs sharing a partition key", async ()
});
});
});

it("labels a hidden asset instead of linking to it, leaving readable assets linked", async () => {
// The shape the pending-detail endpoint serves for an asset the caller may not read.
vi.spyOn(PartitionedDagRunService, "getPendingPartitionedDagRun").mockResolvedValue({
assets: [
{
asset_id: 1,
asset_name: "sales_daily",
asset_uri: "s3://bucket/sales",
hidden: false,
is_rollup: true,
received: false,
received_count: 0,
received_keys: [],
required_count: 1,
required_keys: ["2024-07-01"],
},
{
asset_id: null,
asset_name: null,
asset_uri: null,
hidden: true,
is_rollup: true,
received: false,
received_count: 0,
received_keys: [],
required_count: 1,
required_keys: ["2024-07-01"],
},
],
dag_id: "consumer",
id: 1,
partition_key: "2024-07-01",
total_received: 0,
total_required: 2,
});

render(
<Wrapper>
<AssetProgressCell
dagId="consumer"
partitionedDagRunId={1}
partitionKey="2024-07-01"
totalReceived={0}
totalRequired={2}
/>
</Wrapper>,
);

const trigger = screen.getByRole("button");

// The trigger is disabled while the detail query is in flight, which swallows the click.
await waitFor(() => expect(trigger).not.toBeDisabled());
fireEvent.click(trigger);

await waitFor(() => expect(screen.getByText(i18n.t("common:expression.hiddenAsset"))).toBeInTheDocument());
expect(screen.getByRole("link", { name: "sales_daily" })).toHaveAttribute("href", "/assets/1");
expect(screen.getAllByRole("link")).toHaveLength(1);
});
Loading
Loading