Skip to content

Commit 322d298

Browse files
Karakatiza666Karakatiza666
committed
[pipeline-manager] Add API endpoint to download dataflow graph
Store program_config and circuit_profile gzipped, but include them unzipped in the support bundle Add compression middleware to pipeline-manager Only collect what is requested as a part of support bundle Add an option to not collect the current profile when downloading support bundle Add support for no-dataflow-graph to fda Add tests for retrieving dataflow graph QOL devcontainer improvements Signed-off-by: Karakatiza666 <bulakh.96@gmail.com>
1 parent 85590c2 commit 322d298

15 files changed

Lines changed: 1037 additions & 73 deletions

File tree

‎.devcontainer/Dockerfile‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@ FROM ubuntu:24.04 AS base
44

55
ARG USERNAME=user
66

7-
RUN useradd -m $USERNAME && usermod -aG sudo $USERNAME
8-
RUN echo "user:pass" | chpasswd
7+
# Rename the default user (uid 1000) from `ubuntu` to USERNAME
8+
RUN usermod -l $USERNAME -d /home/$USERNAME -m ubuntu
9+
ENV HOME /home/$USERNAME
10+
RUN echo "$USERNAME:pass" | chpasswd
911

1012
# Set locale to fix postgres embedded startup
1113
ENV LC_ALL=en_US.UTF-8
@@ -58,7 +60,6 @@ RUN pip3 install --break-system-packages gdown kafka-python-ng
5860
## Switch to non-root user
5961

6062
USER $USERNAME
61-
ENV HOME /home/$USERNAME
6263

6364
## Install rustup and common components
6465

‎.devcontainer/devcontainer.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
],
77
"service": "workspace",
88
"runServices": [
9-
"redpanda"
9+
"postgres-dev",
10+
"db-admin"
1011
],
1112
"workspaceFolder": "/workspaces/feldera",
1213
"shutdownAction": "stopCompose",

‎.devcontainer/docker-compose.devcontainer.yml‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,20 @@ services:
1111
- RUST_VERSION=1.87.0
1212
- RUST_BACKTRACE=1
1313
- REDPANDA_BROKERS=redpanda:9092
14+
volumes:
15+
- /var/run/docker.sock:/var/run/docker.sock
16+
postgres-dev:
17+
image: postgres:15-alpine
18+
restart: always
19+
ports:
20+
- 5435:5432
21+
environment:
22+
POSTGRES_PASSWORD: postgres
23+
POSTGRES_USER: postgres
24+
POSTGRES_DB: postgres
25+
PGDATA: /var/lib/pg/data
26+
db-admin:
27+
image: adminer:5-standalone
28+
restart: always
29+
ports:
30+
- 5436:8080

‎.devcontainer/postCreate.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
# Install Playwright system dependencies
2-
cd ./web-console && bun x playwright install-deps
2+
cd ./web-console && bunx playwright install-deps
33
# Install Playwright browsers
4-
bun x playwright install
4+
bunx playwright install

‎crates/adapters/src/server.rs‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1556,6 +1556,7 @@ async fn heap_profile() -> impl Responder {
15561556
match prof_ctl.dump_pprof() {
15571557
Ok(profile) => Ok(HttpResponse::Ok()
15581558
.content_type("application/protobuf")
1559+
.insert_header(header::ContentEncoding::Identity)
15591560
.body(profile)),
15601561
Err(e) => Err(PipelineError::HeapProfilerError {
15611562
error: e.to_string(),
@@ -1575,21 +1576,16 @@ async fn dump_profile(state: WebData<ServerState>) -> Result<HttpResponse, Pipel
15751576
Ok(HttpResponse::Ok()
15761577
.insert_header(header::ContentType("application/zip".parse().unwrap()))
15771578
.insert_header(header::ContentDisposition::attachment("profile.zip"))
1579+
.insert_header(header::ContentEncoding::Identity)
15781580
.body(state.controller()?.async_graph_profile().await?.as_zip()))
15791581
}
15801582

15811583
#[get("/dump_json_profile")]
15821584
async fn dump_json_profile(state: WebData<ServerState>) -> Result<HttpResponse, PipelineError> {
15831585
Ok(HttpResponse::Ok()
1584-
.insert_header(header::ContentType("application/zip".parse().unwrap()))
1585-
.insert_header(header::ContentDisposition::attachment("profile.zip"))
1586-
.body(
1587-
state
1588-
.controller()?
1589-
.async_json_profile()
1590-
.await?
1591-
.as_json_zip(),
1592-
))
1586+
.insert_header(header::ContentType("application/json".parse().unwrap()))
1587+
.insert_header(header::ContentDisposition::attachment("profile.json"))
1588+
.body(state.controller()?.async_json_profile().await?.as_json()))
15931589
}
15941590

15951591
/// Dump the low-level IR of the circuit.

‎crates/fda/src/cli.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,9 @@ pub enum PipelineAction {
536536
/// Skip system configuration collection.
537537
#[arg(long)]
538538
no_system_config: bool,
539+
/// Skip dataflow graph collection.
540+
#[arg(long)]
541+
no_dataflow_graph: bool,
539542
},
540543
/// Enter the ad-hoc SQL shell for a pipeline.
541544
Shell {

‎crates/fda/src/main.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1528,6 +1528,7 @@ async fn pipeline(format: OutputFormat, action: PipelineAction, client: Client)
15281528
no_stats,
15291529
no_pipeline_config,
15301530
no_system_config,
1531+
no_dataflow_graph,
15311532
} => {
15321533
let mut request = client.get_pipeline_support_bundle().pipeline_name(&name);
15331534

@@ -1553,6 +1554,9 @@ async fn pipeline(format: OutputFormat, action: PipelineAction, client: Client)
15531554
if no_system_config {
15541555
request = request.system_config(false);
15551556
}
1557+
if no_dataflow_graph {
1558+
request = request.dataflow_graph(false);
1559+
}
15561560

15571561
let response = request
15581562
.send()

‎crates/ir/test/regen.bash‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,34 @@ fda start sample_c
1111
fda stop sample_c
1212
# end
1313

14-
fda program info sample_a | jq .dataflow > sample_a.json
15-
fda program info sample_b | jq .dataflow > sample_b.json
16-
fda program info sample_b_mod1 | jq .dataflow > sample_b_mod1.json
17-
fda program info sample_b_mod2 | jq .dataflow > sample_b_mod2.json
18-
fda program info sample_b_mod3 | jq .dataflow > sample_b_mod3.json
19-
fda program info sample_c | jq .dataflow > sample_c.json
14+
# Extract dataflow graphs from support bundles
15+
extract_dataflow() {
16+
local pipeline_name=$1
17+
local output_file=$2
18+
19+
# Get support bundle with only dataflow graph enabled (disable other expensive collections)
20+
local bundle_file="${pipeline_name}-support-bundle.zip"
21+
fda support-bundle "$pipeline_name" \
22+
--no-circuit-profile \
23+
--no-heap-profile \
24+
--no-metrics \
25+
--no-logs \
26+
--no-stats \
27+
--no-pipeline-config \
28+
--no-system-config \
29+
--output "$bundle_file"
30+
31+
# Extract dataflow_graph.json from the zip archive (take first match if multiple exist)
32+
local dataflow_file=$(unzip -Z1 "$bundle_file" | grep 'dataflow_graph\.json$' | head -n 1)
33+
unzip -p "$bundle_file" "$dataflow_file" > "$output_file"
34+
35+
# Clean up the temporary bundle file
36+
rm "$bundle_file"
37+
}
38+
39+
extract_dataflow sample_a sample_a.json
40+
extract_dataflow sample_b sample_b.json
41+
extract_dataflow sample_b_mod1 sample_b_mod1.json
42+
extract_dataflow sample_b_mod2 sample_b_mod2.json
43+
extract_dataflow sample_b_mod3 sample_b_mod3.json
44+
extract_dataflow sample_c sample_c.json

‎crates/pipeline-manager/src/api/endpoints/pipeline_interaction.rs‎

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -744,7 +744,7 @@ pub(crate) async fn get_pipeline_circuit_profile(
744744
responses(
745745
(status = OK
746746
, description = "Circuit performance profile in JSON format"
747-
, content_type = "application/zip"
747+
, content_type = "application/json"
748748
, body = Object),
749749
(status = NOT_FOUND
750750
, description = "Pipeline with that name does not exist"
@@ -772,6 +772,9 @@ pub(crate) async fn get_pipeline_circuit_json_profile(
772772
request: HttpRequest,
773773
) -> Result<HttpResponse, ManagerError> {
774774
let pipeline_name = path.into_inner();
775+
776+
// Get the JSON profile from the pipeline and return it directly
777+
// The Compress middleware will automatically compress the response based on Accept-Encoding
775778
state
776779
.runner
777780
.forward_http_request_to_pipeline_by_name(
@@ -786,6 +789,87 @@ pub(crate) async fn get_pipeline_circuit_json_profile(
786789
.await
787790
}
788791

792+
/// Shared helper function to extract dataflow graph from a pipeline's program_info.
793+
///
794+
/// This function encapsulates the common logic for retrieving and validating
795+
/// the dataflow graph from a pipeline's compiled program information.
796+
pub(crate) async fn extract_pipeline_dataflow_graph(
797+
state: &ServerState,
798+
tenant_id: TenantId,
799+
pipeline_name: &str,
800+
) -> Result<feldera_ir::Dataflow, ManagerError> {
801+
// Get pipeline from database
802+
let pipeline = state
803+
.db
804+
.lock()
805+
.await
806+
.get_pipeline(tenant_id, pipeline_name)
807+
.await?;
808+
809+
// Extract dataflow from program_info
810+
if let Some(program_info_value) = &pipeline.program_info {
811+
// Parse program_info JSON to extract dataflow field
812+
match serde_json::from_value::<crate::db::types::program::ProgramInfo>(
813+
program_info_value.clone(),
814+
) {
815+
Ok(program_info) => {
816+
if let Some(dataflow) = program_info.dataflow {
817+
Ok(dataflow)
818+
} else {
819+
Err(ApiError::ProgramInfoMissesDataflow {
820+
pipeline_name: pipeline_name.to_string(),
821+
}
822+
.into())
823+
}
824+
}
825+
Err(e) => Err(ApiError::InvalidProgramInfo {
826+
error: e.to_string(),
827+
}
828+
.into()),
829+
}
830+
} else {
831+
Err(ApiError::ProgramNotCompiled {
832+
pipeline_name: pipeline_name.to_string(),
833+
}
834+
.into())
835+
}
836+
}
837+
838+
/// Get Dataflow Graph
839+
///
840+
/// Retrieve the dataflow graph of a pipeline.
841+
/// The dataflow graph is generated during SQL compilation and shows the structure
842+
/// of the compiled SQL program including the Calcite plan and MIR nodes.
843+
#[utoipa::path(
844+
context_path = "/v0",
845+
security(("JSON web token (JWT) or API key" = [])),
846+
params(
847+
("pipeline_name" = String, Path, description = "Unique pipeline name"),
848+
),
849+
responses(
850+
(status = OK
851+
, description = "Dataflow graph retrieved successfully"
852+
, content_type = "application/json"
853+
, body = Object),
854+
(status = NOT_FOUND
855+
, description = "Pipeline with that name does not exist or dataflow graph is not available"
856+
, body = ErrorResponse
857+
, example = json!(examples::error_unknown_pipeline_name())),
858+
(status = INTERNAL_SERVER_ERROR, body = ErrorResponse),
859+
),
860+
tag = "Metrics & Debugging"
861+
)]
862+
#[get("/pipelines/{pipeline_name}/dataflow_graph")]
863+
pub(crate) async fn get_pipeline_dataflow_graph(
864+
state: WebData<ServerState>,
865+
tenant_id: ReqData<TenantId>,
866+
path: web::Path<String>,
867+
) -> Result<HttpResponse, ManagerError> {
868+
let pipeline_name = path.into_inner();
869+
let dataflow = extract_pipeline_dataflow_graph(&state, *tenant_id, &pipeline_name).await?;
870+
Ok(HttpResponse::Ok().json(dataflow))
871+
}
872+
789873
/// Sync Checkpoints To S3
790874
///
791875
/// Syncs latest checkpoints to the object store configured in pipeline config.

‎crates/pipeline-manager/src/api/endpoints/pipeline_interaction/support_bundle.rs‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -171,10 +171,22 @@ pub(crate) async fn get_pipeline_support_bundle(
171171
state.config.support_data_retention,
172172
)
173173
.await?;
174-
bundles.insert(
175-
0,
176-
SupportBundleData::collect(&state, &client, *tenant_id, &pipeline_name).await?,
177-
);
174+
175+
// Only collect all requested data from the running pipeline if collect=true
176+
if support_bundle_params.collect {
177+
bundles.insert(
178+
0,
179+
SupportBundleData::collect(
180+
&state,
181+
&client,
182+
*tenant_id,
183+
&pipeline_name,
184+
&support_bundle_params,
185+
)
186+
.await?,
187+
);
188+
};
189+
178190
let bundle = SupportBundleZip::create(&pipeline, bundles, &support_bundle_params).await?;
179191

180192
Ok(HttpResponse::Ok()
@@ -186,5 +198,6 @@ pub(crate) async fn get_pipeline_support_bundle(
186198
pipeline_name
187199
),
188200
))
201+
.insert_header(actix_http::ContentEncoding::Identity)
189202
.body(bundle.buffer))
190203
}

0 commit comments

Comments
 (0)