Support for latest Airflow pipelines (TaskFlow & Classic)#25
Open
Resmung0 wants to merge 4 commits into
Open
Conversation
- Improved Airflow DAG detection logic to include TaskFlow decorators. - Implemented Airflow CLI integration using `airflow dags show` for accurate structure extraction. - Added a robust regex-based fallback parser for static analysis when the CLI is unavailable. - Implemented task code snippet extraction for both TaskFlow and classic operators. - Added comprehensive unit tests for AirflowParser. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
Contributor
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
Reviewer's GuideExtends the Airflow parser to support Airflow 2.x/3.x TaskFlow and classic DAGs using either CLI-based DOT extraction or a regex-based static fallback, and adds tests for detection, regex parsing, and DOT parsing. Sequence diagram for AirflowParser parse flow with CLI and regex fallbacksequenceDiagram
participant Caller
participant AirflowParser
participant fs as FileSystem
participant AirflowCmd as airflow_or_uv
Caller->>AirflowParser: parse(content, filePath)
AirflowParser->>AirflowParser: getAirflowCmd(cwd)
alt cached command in airflowCmdCache
AirflowParser-->>AirflowParser: return cached CommandInfo
else no cached command
AirflowParser->>fs: check venv airflow executables
alt venv airflow found
AirflowParser-->>AirflowParser: set CommandInfo (local venv)
else venv airflow not found
AirflowParser->>AirflowCmd: exec uv --version
AirflowParser->>AirflowCmd: exec uv run airflow version
alt uv airflow available
AirflowParser-->>AirflowParser: set CommandInfo (uv run airflow)
else uv airflow not available
AirflowParser->>AirflowCmd: exec airflow version
alt global airflow available
AirflowParser-->>AirflowParser: set CommandInfo (global airflow)
else no airflow
AirflowParser-->>AirflowParser: CommandInfo null
end
end
end
AirflowParser->>AirflowParser: cache CommandInfo by cwd
end
alt CommandInfo not null
AirflowParser->>AirflowParser: extractDagId(content)
alt dagId found
AirflowParser->>AirflowCmd: exec "airflow dags show dagId" (DOT)
AirflowCmd-->>AirflowParser: DOT stdout
AirflowParser->>AirflowParser: parseDot(dot, filePath)
AirflowParser-->>AirflowParser: PipelineData
AirflowParser->>AirflowParser: enrich nodes with extractTaskSnippet and extractOperatorSnippet
AirflowParser-->>Caller: PipelineData (CLI based)
else no dagId
AirflowParser->>AirflowParser: parseWithRegex(content, filePath)
AirflowParser-->>Caller: PipelineData (regex based)
end
else no airflow command available
AirflowParser->>AirflowParser: parseWithRegex(content, filePath)
AirflowParser-->>Caller: PipelineData (regex based)
end
Updated class diagram for AirflowParser and related typesclassDiagram
class IParser {
<<interface>>
+name: string
+canParse(fileName: string, content: string): boolean
+parse(content: string, filePath: string): Promise~PipelineData~
}
class AirflowParser {
+name: string
-airflowCmdCache: Map~string, AirflowCmdEntry~
+canParse(fileName: string, content: string): boolean
+parse(content: string, filePath: string): Promise~PipelineData~
-extractDagId(content: string): string
-parseWithCLI(dagId: string, filePath: string, cmdInfo: CommandInfo, cwd: string): Promise~PipelineData~
-parseDot(dot: string, filePath: string): PipelineData
-parseWithRegex(content: string, filePath: string): PipelineData
-extractTaskSnippet(content: string, taskId: string): any[]
-extractOperatorSnippet(content: string, taskId: string): any[]
-extractOperatorSnippetByVar(content: string, varName: string): any[]
-getAirflowCmd(cwd: string): Promise~AirflowCmdEntry~
}
class CommandInfo {
+command: string
+args: string[]
}
class AirflowCmdEntry {
+commandInfo: CommandInfo
+isLocal: boolean
}
class PipelineData {
+filePath: string
+framework: string
+nodes: PipelineNode[]
+edges: PipelineEdge[]
}
class PipelineNode {
+id: string
+label: string
+type: string
+data: any
}
class PipelineEdge {
+id: string
+source: string
+target: string
}
IParser <|.. AirflowParser
AirflowParser ..> CommandInfo
AirflowParser ..> AirflowCmdEntry
AirflowParser ..> PipelineData
PipelineData "1" *-- "*" PipelineNode
PipelineData "1" *-- "*" PipelineEdge
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 9 issues, and left some high level feedback:
- The CLI invocation in
parseWithCLIbuilds a single shell command string with the DAG ID and command name interpolated, which can break on spaces and special characters and is harder to secure; consider usingexecFile/spawnwith an args array (reusingCommandInfo.args) instead ofexecwith a joined string. - The snippet extractors (
extractTaskSnippet,extractOperatorSnippet,extractOperatorSnippetByVar) hardcodepath: 'dag.py'instead of using the actualfilePath, which can lead to misleading metadata when parsing files with different names; passingfilePaththrough and using it in the snippet objects would improve accuracy. - The snippet extraction helpers currently return
any[], which loses type information and makes downstream use more error-prone; defining a smallCodeSnippetinterface and updating these methods and thecodeDepsusage to use that type would improve clarity and type safety.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The CLI invocation in `parseWithCLI` builds a single shell command string with the DAG ID and command name interpolated, which can break on spaces and special characters and is harder to secure; consider using `execFile`/`spawn` with an args array (reusing `CommandInfo.args`) instead of `exec` with a joined string.
- The snippet extractors (`extractTaskSnippet`, `extractOperatorSnippet`, `extractOperatorSnippetByVar`) hardcode `path: 'dag.py'` instead of using the actual `filePath`, which can lead to misleading metadata when parsing files with different names; passing `filePath` through and using it in the snippet objects would improve accuracy.
- The snippet extraction helpers currently return `any[]`, which loses type information and makes downstream use more error-prone; defining a small `CodeSnippet` interface and updating these methods and the `codeDeps` usage to use that type would improve clarity and type safety.
## Individual Comments
### Comment 1
<location path="src/extension/parsers/data-processing/AirflowParser.ts" line_range="68-69" />
<code_context>
+ return null;
+ }
+
+ private async parseWithCLI(dagId: string, filePath: string, cmdInfo: CommandInfo, cwd: string): Promise<PipelineData> {
+ const fullCmd = [cmdInfo.command, ...cmdInfo.args, 'dags', 'show', dagId].join(' ');
+ const { stdout } = await execPromise(fullCmd, { cwd });
+
</code_context>
<issue_to_address>
**🚨 issue (security):** Building the CLI command as a single joined string risks issues with dag IDs containing spaces/special characters and opens the door to shell-injection-like problems.
Interpolating `dagId` into a single string command means spaces or shell metacharacters in the ID will be interpreted by the shell, which is brittle and potentially exploitable. Prefer `execFile`/`spawn` with an argument array so `dagId` is passed as a separate argument, or ensure robust quoting/escaping so it cannot modify the command line.
</issue_to_address>
### Comment 2
<location path="src/extension/parsers/data-processing/AirflowParser.ts" line_range="86-95" />
<code_context>
+ while ((match = operatorRegex.exec(content)) !== null) {
+ const varName = match[1];
+ const taskId = match[2] || varName;
+ nodes.push({
+ id: taskId,
+ label: taskId,
+ type: 'default',
+ data: {
+ framework: this.name,
+ variableName: varName,
+ codeDeps: this.extractOperatorSnippet(content, taskId) || this.extractOperatorSnippetByVar(content, varName)
+ }
+ });
</code_context>
<issue_to_address>
**issue (bug_risk):** The fallback to `extractOperatorSnippetByVar` never runs because an empty array from `extractOperatorSnippet` is still truthy.
Because `extractOperatorSnippet` always returns an array, the `||` never falls through to `extractOperatorSnippetByVar`, even when the first result is `[]`. If you want to fall back when no snippet is found by `taskId`, check the array length instead, e.g.:
```ts
const byTaskId = this.extractOperatorSnippet(content, taskId);
const codeDeps = byTaskId.length > 0
? byTaskId
: this.extractOperatorSnippetByVar(content, varName);
```
This makes the var-name lookup a real fallback when the task-based lookup finds nothing.
</issue_to_address>
### Comment 3
<location path="src/extension/parsers/data-processing/AirflowParser.ts" line_range="219-220" />
<code_context>
+ }
+ }
+
+ return [{
+ path: 'dag.py',
+ snippet: lines.slice(startIdx, endIdx + 1).join('\n')
+ }];
</code_context>
<issue_to_address>
**suggestion:** Hardcoding the snippet path to 'dag.py' may be misleading when parsing other files and makes cross-file usage less flexible.
Because `parse` already knows the concrete `filePath`, always returning `'dag.py'` here can mislead consumers when the DAG is defined in another file. Instead, thread `filePath` (or its basename) into `extractTaskSnippet`/`extractOperatorSnippet*` and use that value for `path` so snippet metadata reflects the actual source file and can be reused across files.
Suggested implementation:
```typescript
private extractTaskSnippet(content: string, filePath: string, taskId: string): any[] {
```
```typescript
return [{
path: filePath,
snippet: lines.slice(startIdx, endIdx + 1).join('\n')
}];
```
```typescript
private extractOperatorSnippet(content: string, filePath: string, taskId: string): any[] {
```
```typescript
return [{
path: filePath,
snippet: lines[lineIdx].trim()
```
To fully implement this change, you will also need to:
1. Update all call sites of `extractTaskSnippet` to pass `filePath` as the second argument, e.g. `this.extractTaskSnippet(content, filePath, taskId)`.
2. Update all call sites of `extractOperatorSnippet` to pass `filePath` as the second argument, e.g. `this.extractOperatorSnippet(content, filePath, taskId)`.
3. If you prefer to store only the basename in `path`, you can derive it at the call site using `import path from 'path';` and pass `path.basename(filePath)` instead of the full path.
</issue_to_address>
### Comment 4
<location path="tests/AirflowParser.test.ts" line_range="41-42" />
<code_context>
+ });
+ });
+
+ describe('parseWithRegex', () => {
+ it('should extract tasks and edges from TaskFlow DAG', () => {
+ const content = `
+@dag(dag_id="test_dag")
</code_context>
<issue_to_address>
**suggestion (testing):** Missing tests for the top-level `parse` function, especially CLI and fallback behavior
Please add tests that cover the orchestration logic in `parse`, including:
- Airflow CLI available and successful (mock `getAirflowCmd` / `child_process.exec`), asserting `parse` uses DOT output rather than regex.
- Airflow CLI available but returning DOT without `digraph`, ensuring `parseWithCLI` throws and `parse` falls back to regex.
- No Airflow CLI available, verifying `parse` falls back directly to regex.
- No `dag_id` extracted, confirming `parse` skips CLI and uses regex.
This will ensure the integration paths and fallbacks are exercised, not just the lower-level helpers.
Suggested implementation:
```typescript
expect(parser.canParse('script.py', content)).toBe(false);
});
});
describe('parse', () => {
const filePath = 'dags/test_dag.py';
const dagContent = `
from airflow.decorators import dag, task
@dag(dag_id="test_dag")
def my_dag():
@task
def my_task():
return "hello"
my_task()
`;
const nonDagContent = `
def my_dag():
return "not an airflow dag"
`;
afterEach(() => {
jest.restoreAllMocks();
});
it('should use Airflow CLI when available and successful', async () => {
const cliResult = { nodes: ['task1'], edges: [] };
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockResolvedValue(cliResult);
const parseWithRegexSpy = jest.spyOn(parser as any, 'parseWithRegex');
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithCLISpy).toHaveBeenCalledWith(filePath, dagContent, 'test_dag');
expect(parseWithRegexSpy).not.toHaveBeenCalled();
expect(result).toBe(cliResult);
});
it('should fall back to regex when Airflow CLI returns invalid DOT', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockImplementation(async () => {
// Simulate parseWithCLI failing because DOT output is missing "digraph"
throw new Error('Invalid DOT output from Airflow CLI');
});
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, dagContent);
expect(result).toBe(regexResult);
});
it('should fall back directly to regex when Airflow CLI is not available', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
// Simulate "no CLI available" by having parseWithCLI always throw
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockImplementation(async () => {
throw new Error('Airflow CLI not available');
});
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, dagContent);
expect(result).toBe(regexResult);
});
it('should skip CLI and use regex when no dag_id is extracted', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
const parseWithCLISpy = jest.spyOn(parser as any, 'parseWithCLI');
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, nonDagContent);
expect(parseWithCLISpy).not.toHaveBeenCalled();
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, nonDagContent);
expect(result).toBe(regexResult);
});
});
```
These tests assume:
1. `parser` has methods `parse`, `parseWithCLI`, and `parseWithRegex` and that `parse` has the signature `parse(filePath: string, content: string)`.
2. `parse` internally extracts `dag_id` (e.g. `"test_dag"`) from the content and passes it as the third argument to `parseWithCLI(filePath, content, dagId)`.
If any of these assumptions differ in your codebase, adjust:
- The arguments in `toHaveBeenCalledWith` expectations to match the actual `parseWithCLI` / `parseWithRegex` signatures.
- The DAG content / `dag_id` string (`"test_dag"`) to align with your real extraction logic.
If `parseWithCLI` / `parseWithRegex` are not instance methods on `parser` but are imported helpers, replace `jest.spyOn(parser as any, 'parseWithCLI')` with `jest.spyOn` on the actual module export, and update the expectations accordingly.
</issue_to_address>
### Comment 5
<location path="tests/AirflowParser.test.ts" line_range="42-51" />
<code_context>
+ });
+
+ describe('parseWithRegex', () => {
+ it('should extract tasks and edges from TaskFlow DAG', () => {
+ const content = `
+@dag(dag_id="test_dag")
+def my_dag():
+ @task
+ def task_a():
+ return "a"
+
+ @task
+ def task_b():
+ return "b"
+
+ task_a() >> task_b()
+`;
+ const result = (parser as any).parseWithRegex(content, 'dag.py');
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' }));
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' }));
</code_context>
<issue_to_address>
**suggestion (testing):** No tests cover `<<` or `set_upstream` dependency definitions in the regex parser
`parseWithRegex` already supports `<<` and `set_upstream`, but the tests only cover `>>` and `set_downstream`. Please add coverage for:
- `task_b() << task_a()` and assert the edge is `source: 'task_a', target: 'task_b'`.
- `task1.set_upstream(task2)` and assert the edge is `source: 'task2', target: 'task1'`.
This will verify edge direction is handled correctly for all supported forms and guard against regressions.
</issue_to_address>
### Comment 6
<location path="tests/AirflowParser.test.ts" line_range="57-60" />
<code_context>
+ expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' }));
+ });
+
+ it('should extract tasks from classic Operators', () => {
+ const content = `
+task1 = BashOperator(task_id="bash_task", bash_command="echo 1")
+task2 = PythonOperator(task_id="python_task", python_callable=my_func)
+task1 >> task2
+`;
+ const result = (parser as any).parseWithRegex(content, 'dag.py');
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'bash_task' }));
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'python_task' }));
+ expect(result.edges).toContainEqual(expect.objectContaining({ source: 'bash_task', target: 'python_task' }));
+ });
+
</code_context>
<issue_to_address>
**suggestion (testing):** Missing tests for operator regex when `task_id` is omitted and when edges use variable names instead of task IDs
`parseWithRegex` also supports operators where `task_id` is omitted (e.g. `task1 = BashOperator(...)`, using `task1` as the id) and dependency resolution via either `variableName` or `id`. The new test only covers explicit `task_id`s. Please add tests that:
- Cover operators without `task_id` (e.g. `task1 = BashOperator()`, `task2 = BashOperator()`, `task1 >> task2`) and assert edges are created between the inferred ids.
- Cover dependencies where the variable name differs from `task_id`, verifying that `getTaskId` resolves the correct node id for both bitshift and method-based dependencies.
This will exercise the variableName vs id resolution logic and guard against regressions in dependency parsing.
```suggestion
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' }));
expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' }));
});
it('should extract tasks from classic Operators', () => {
const content = `
task1 = BashOperator(task_id="bash_task", bash_command="echo 1")
task2 = PythonOperator(task_id="python_task", python_callable=my_func)
task1 >> task2
`;
const result = (parser as any).parseWithRegex(content, 'dag.py');
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'bash_task' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'python_task' }));
expect(result.edges).toContainEqual(expect.objectContaining({ source: 'bash_task', target: 'python_task' }));
});
it('should infer task ids from variable names when task_id is omitted', () => {
const content = `
task1 = BashOperator()
task2 = BashOperator()
task1 >> task2
`;
const result = (parser as any).parseWithRegex(content, 'dag.py');
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task1' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task2' }));
expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task1', target: 'task2' }));
});
it('should resolve dependencies using both variable names and task_ids', () => {
const content = `
bash = BashOperator(task_id="bash_task")
python = PythonOperator(task_id="python_task")
# bitshift dependencies mixing variable names and task_ids
bash >> "python_task"
"bash_task" >> python
# method-based dependencies mixing variable names and task_ids
bash.set_downstream("python_task")
python.set_upstream("bash_task")
`;
const result = (parser as any).parseWithRegex(content, 'dag.py');
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'bash_task' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'python_task' }));
expect(result.edges).toContainEqual(
expect.objectContaining({ source: 'bash_task', target: 'python_task' }),
);
});
```
</issue_to_address>
### Comment 7
<location path="tests/AirflowParser.test.ts" line_range="56-59" />
<code_context>
+
+ task_a() >> task_b()
+`;
+ const result = (parser as any).parseWithRegex(content, 'dag.py');
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' }));
+ expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' }));
+ expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' }));
+ });
+
</code_context>
<issue_to_address>
**suggestion (testing):** Code snippet extraction for tasks and operators is not asserted in tests
The parser now populates `data.codeDeps` via `extractTaskSnippet` / `extractOperatorSnippet` (and in the CLI path limits to the first match), but the current tests don’t check this at all.
Please extend the TaskFlow and classic-operator regex tests to assert that `data.codeDeps` is set on the relevant nodes and that at least one snippet includes the expected `def task_a` or `BashOperator(task_id="bash_task"` text.
Optionally, add a focused unit test for one of the snippet helpers (e.g. `(parser as any).extractTaskSnippet`) to verify multi-line function bodies, the returned `path`, and snippet contents. This will help guard the new "Code Snippets" feature against future regex/indentation changes.
Suggested implementation:
```typescript
const result = (parser as any).parseWithRegex(content, 'dag.py');
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' }));
expect(result.edges).toContainEqual(
expect.objectContaining({ source: 'task_a', target: 'task_b' }),
);
const taskANode = result.nodes.find((n: any) => n.id === 'task_a');
expect(taskANode).toBeDefined();
expect(taskANode?.data?.codeDeps?.length).toBeGreaterThan(0);
expect(taskANode?.data?.codeDeps?.[0]).toEqual(
expect.objectContaining({
path: 'dag.py',
snippet: expect.stringContaining('def task_a'),
}),
);
});
```
I could only see the TaskFlow regex test in the snippet, so the following additional changes are required elsewhere in `tests/AirflowParser.test.ts`:
1. **Classic operator regex test (`BashOperator`)**:
- Locate the test that exercises the classic operator regex, likely where content includes `BashOperator(task_id="bash_task"` or similar.
- After the existing assertions on `result.nodes` / `result.edges`, add expectations similar to:
```ts
const bashNode = result.nodes.find((n: any) => n.id === 'bash_task');
expect(bashNode).toBeDefined();
expect(bashNode?.data?.codeDeps?.length).toBeGreaterThan(0);
expect(bashNode?.data?.codeDeps?.[0]).toEqual(
expect.objectContaining({
path: 'dag.py', // or the file name passed into parseWithRegex
snippet: expect.stringContaining('BashOperator(task_id="bash_task"'),
}),
);
```
- Ensure you use the actual `path` argument you pass to `parseWithRegex` in that test (it might not be `'dag.py'`).
2. **Focused unit test for `extractTaskSnippet` (optional but recommended)**:
- Add a new `describe` block under `describe('AirflowParser', () => { ... })`, e.g.:
```ts
describe('extractTaskSnippet', () => {
it('extracts multi-line TaskFlow function bodies with path and snippet', () => {
const content = `
from airflow.decorators import task
@task
def task_a():
x = 1
y = 2
return x + y
`;
const parserAny = parser as any;
const snippet = parserAny.extractTaskSnippet(content, 'dag.py', 'task_a');
expect(snippet).toEqual(
expect.objectContaining({
path: 'dag.py',
snippet: expect.stringContaining('def task_a'),
}),
);
expect(snippet.snippet).toContain('x = 1');
expect(snippet.snippet).toContain('return x + y');
});
});
```
- If `extractTaskSnippet` has a different signature (e.g. returns an array or includes additional metadata), adjust the expectations accordingly but still assert:
- The returned `path` matches the passed file name.
- The snippet includes `def task_a` and multiple lines from the function body.
These additional changes will ensure that both TaskFlow and classic operator parsing paths assert `data.codeDeps` population and that the snippet extraction helper is covered by focused tests.
</issue_to_address>
### Comment 8
<location path="tests/AirflowParser.test.ts" line_range="85-86" />
<code_context>
+ });
+ });
+
+ describe('parseDot', () => {
+ it('should parse Airflow DOT output', () => {
+ const dot = `
+digraph test_dag {
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a test around DOT validation in the CLI flow (e.g., missing `digraph`)
Currently we only test the happy-path DOT input. Since `parseWithCLI` performs an extra `stdout.includes('digraph')` validation and throws otherwise, please add a test where the mocked CLI returns output without `digraph`. Assert that:
- `parseWithCLI` throws, and
- `parse` falls back to the regex-based parsing (e.g., by spying on `parseWithRegex`).
This will verify the fallback behavior for malformed/non-DOT CLI output.
Suggested implementation:
```typescript
expect(result.nodes).toHaveLength(2);
expect(result.nodes[0]).toEqual(expect.objectContaining({ id: 'task_1', label: 'task_1' }));
expect(result.edges).toHaveLength(1);
expect(result.edges[0]).toEqual(expect.objectContaining({ source: 'task_1', target: 'task_2' }));
});
it('should throw in parseWithCLI when CLI output is not DOT (missing digraph)', async () => {
const stdout = `
some random output
without the expected keyword
`;
const runCliSpy = jest
.spyOn(parser as any, 'runCli')
.mockResolvedValueOnce({ stdout });
await expect((parser as any).parseWithCLI('dag.py')).rejects.toThrow();
expect(runCliSpy).toHaveBeenCalled();
});
it('should fall back to regex parsing when CLI output validation fails', async () => {
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockRejectedValueOnce(new Error('CLI output is not DOT'));
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockReturnValueOnce({
nodes: [{ id: 't1', label: 't1' }],
edges: [],
});
const result = await (parser as any).parse('content', 'dag.py');
expect(parseWithCLISpy).toHaveBeenCalled();
expect(parseWithRegexSpy).toHaveBeenCalledWith('content', 'dag.py');
expect(result).toEqual({
nodes: [{ id: 't1', label: 't1' }],
edges: [],
});
});
```
1. The above assumes that `AirflowParser` (or whatever `parser` is) exposes a method named `runCli` that `parseWithCLI` uses internally to invoke the Airflow CLI and capture `{ stdout }`. If the actual helper has a different name or is accessed differently, replace `'runCli'` in the `jest.spyOn` call with the correct method name or adapt the mocking to the existing pattern in this test file (e.g., spying on a `cliExecutor`, `exec` wrapper, or a module-level function).
2. Ensure `parseWithCLI` is an `async` function that returns a promise and performs the `stdout.includes('digraph')` validation; the first new test relies on it rejecting when the validation fails.
3. If `parse` is not already using `parseWithCLI` first and then `parseWithRegex` on failure, make sure its implementation matches this behavior so the fallback test passes: `parse` should call `parseWithCLI`, `catch` any error, then call and return `parseWithRegex`.
</issue_to_address>
### Comment 9
<location path="src/extension/parsers/data-processing/AirflowParser.ts" line_range="25" />
<code_context>
+ return hasAirflowImport && hasDagDefinition;
}
async parse(content: string, filePath: string): Promise<PipelineData> {
- // Placeholder implementation
+ const cwd = path.dirname(filePath);
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring `parse`, `getAirflowCmd`, and the snippet helpers into smaller focused functions to flatten control flow and remove duplicated logic while preserving behavior.
The added functionality is solid, but there are a few spots where you can reduce complexity and flatten control flow without changing behavior.
---
### 1. `parse` – separate CLI path from fallback
`parse` currently handles: cwd detection, airflow resolution, dagId extraction, CLI invocation, snippet enrichment, and regex fallback in one method.
You already have `parseWithCLI`; a thin wrapper for the whole “CLI + enrichment” path makes `parse` easier to follow and test.
```ts
async parse(content: string, filePath: string): Promise<PipelineData> {
const cwd = path.dirname(filePath);
const airflowInfo = await this.getAirflowCmd(cwd);
const cliResult = airflowInfo
? await this.tryParseWithCLI(content, filePath, airflowInfo.commandInfo, cwd)
: null;
return cliResult ?? this.parseWithRegex(content, filePath);
}
private async tryParseWithCLI(
content: string,
filePath: string,
cmdInfo: CommandInfo,
cwd: string
): Promise<PipelineData | null> {
const dagId = this.extractDagId(content);
if (!dagId) return null;
try {
const data = await this.parseWithCLI(dagId, filePath, cmdInfo, cwd);
data.nodes = data.nodes.map(node => ({
...node,
data: {
...node.data,
codeDeps: [
...this.extractTaskSnippet(content, node.id),
...this.extractOperatorSnippet(content, node.id),
].slice(0, 1),
},
}));
return data;
} catch (error) {
console.error("Airflow CLI parsing failed, falling back to regex:", error);
return null;
}
}
```
This keeps `parse` as “orchestrator” only, and pushes the nested CLI/try-catch + enrichment into a dedicated helper.
---
### 2. `getAirflowCmd` – split discovery concerns
`getAirflowCmd` is doing: venv search, `uv` detection, global `airflow` detection, and caching. Splitting the discovery into small helpers makes the control flow clearer while keeping caching + behavior intact.
```ts
private async getAirflowCmd(cwd: string): Promise<{ commandInfo: CommandInfo; isLocal: boolean } | null> {
if (this.airflowCmdCache.has(cwd)) {
return this.airflowCmdCache.get(cwd)!;
}
let result = this.findVenvAirflow(cwd)
?? await this.findUvAirflow(cwd)
?? await this.findGlobalAirflow();
this.airflowCmdCache.set(cwd, result);
return result;
}
private findVenvAirflow(cwd: string): { commandInfo: CommandInfo; isLocal: boolean } | null {
const isWindows = process.platform === 'win32';
const venvDirs = ['.venv', 'venv', 'env'];
for (const dir of venvDirs) {
const venvAirflow = path.join(
cwd,
dir,
isWindows ? 'Scripts' : 'bin',
isWindows ? 'airflow.exe' : 'airflow',
);
if (fs.existsSync(venvAirflow)) {
return { commandInfo: { command: venvAirflow, args: [] }, isLocal: true };
}
}
return null;
}
private async findUvAirflow(cwd: string): Promise<{ commandInfo: CommandInfo; isLocal: boolean } | null> {
try {
await execPromise('uv --version');
await execPromise('uv run airflow version', { cwd });
return { commandInfo: { command: 'uv', args: ['run', 'airflow'] }, isLocal: false };
} catch {
return null;
}
}
private async findGlobalAirflow(): Promise<{ commandInfo: CommandInfo; isLocal: boolean } | null> {
try {
await execPromise('airflow version');
return { commandInfo: { command: 'airflow', args: [] }, isLocal: false };
} catch {
return null;
}
}
```
This keeps functionality and caching but removes nested `try` blocks and local state juggling.
---
### 3. Snippet helpers – introduce shared utilities
You have three helpers that all split content into lines and do line-based searches; `extractTaskSnippet` adds nested indentation logic. Extracting small, generic helpers removes duplication and makes `extractTaskSnippet` easier to reason about.
```ts
private findLineIndex(lines: string[], predicate: (line: string, idx: number) => boolean): number {
return lines.findIndex(predicate);
}
private getIndentedBlock(lines: string[], headerIdx: number): { start: number; end: number } {
const defIdx = lines.findIndex((line, i) => i > headerIdx && line.includes('def ') && !line.trim().startsWith('#'));
if (defIdx === -1) return { start: headerIdx, end: headerIdx };
const startIndent = lines[defIdx].search(/\S/);
let endIdx = defIdx;
for (let i = defIdx + 1; i < lines.length; i++) {
const trimmed = lines[i].trim();
if (trimmed === '' || trimmed.startsWith('#')) {
endIdx = i;
continue;
}
const currentIndent = lines[i].search(/\S/);
if (currentIndent <= startIndent && currentIndent !== -1) break;
endIdx = i;
}
return { start: headerIdx, end: endIdx };
}
```
Then `extractTaskSnippet` and the operator helpers become simpler:
```ts
private extractTaskSnippet(content: string, taskId: string): any[] {
const lines = content.split('\n');
const startIdx = this.findLineIndex(
lines,
(line, idx) =>
line.includes('@task') &&
(lines[idx + 1]?.includes(`def ${taskId}`) || lines[idx + 2]?.includes(`def ${taskId}`)),
);
if (startIdx === -1) return [];
const { start, end } = this.getIndentedBlock(lines, startIdx);
return [{ path: 'dag.py', snippet: lines.slice(start, end + 1).join('\n') }];
}
private extractOperatorSnippet(content: string, taskId: string): any[] {
const lines = content.split('\n');
const lineIdx = this.findLineIndex(
lines,
line =>
line.includes(`task_id=["']${taskId}["']`) ||
line.includes(`task_id = ["']${taskId}["']`),
);
if (lineIdx === -1) return [];
return [{ path: 'dag.py', snippet: lines[lineIdx].trim() }];
}
private extractOperatorSnippetByVar(content: string, varName: string): any[] {
const lines = content.split('\n');
const lineIdx = this.findLineIndex(
lines,
line => line.includes(`${varName} =`) && line.includes('Operator'),
);
if (lineIdx === -1) return [];
return [{ path: 'dag.py', snippet: lines[lineIdx].trim() }];
}
```
This keeps behavior but centralizes the line-search and indentation logic, reducing nested loops and duplicated conditions.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+41
to
+42
| describe('parseWithRegex', () => { | ||
| it('should extract tasks and edges from TaskFlow DAG', () => { |
There was a problem hiding this comment.
suggestion (testing): Missing tests for the top-level parse function, especially CLI and fallback behavior
Please add tests that cover the orchestration logic in parse, including:
- Airflow CLI available and successful (mock
getAirflowCmd/child_process.exec), assertingparseuses DOT output rather than regex. - Airflow CLI available but returning DOT without
digraph, ensuringparseWithCLIthrows andparsefalls back to regex. - No Airflow CLI available, verifying
parsefalls back directly to regex. - No
dag_idextracted, confirmingparseskips CLI and uses regex.
This will ensure the integration paths and fallbacks are exercised, not just the lower-level helpers.
Suggested implementation:
expect(parser.canParse('script.py', content)).toBe(false);
});
});
describe('parse', () => {
const filePath = 'dags/test_dag.py';
const dagContent = `
from airflow.decorators import dag, task
@dag(dag_id="test_dag")
def my_dag():
@task
def my_task():
return "hello"
my_task()
`;
const nonDagContent = `
def my_dag():
return "not an airflow dag"
`;
afterEach(() => {
jest.restoreAllMocks();
});
it('should use Airflow CLI when available and successful', async () => {
const cliResult = { nodes: ['task1'], edges: [] };
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockResolvedValue(cliResult);
const parseWithRegexSpy = jest.spyOn(parser as any, 'parseWithRegex');
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithCLISpy).toHaveBeenCalledWith(filePath, dagContent, 'test_dag');
expect(parseWithRegexSpy).not.toHaveBeenCalled();
expect(result).toBe(cliResult);
});
it('should fall back to regex when Airflow CLI returns invalid DOT', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockImplementation(async () => {
// Simulate parseWithCLI failing because DOT output is missing "digraph"
throw new Error('Invalid DOT output from Airflow CLI');
});
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, dagContent);
expect(result).toBe(regexResult);
});
it('should fall back directly to regex when Airflow CLI is not available', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
// Simulate "no CLI available" by having parseWithCLI always throw
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockImplementation(async () => {
throw new Error('Airflow CLI not available');
});
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, dagContent);
expect(parseWithCLISpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, dagContent);
expect(result).toBe(regexResult);
});
it('should skip CLI and use regex when no dag_id is extracted', async () => {
const regexResult = { nodes: ['task_from_regex'], edges: [] };
const parseWithCLISpy = jest.spyOn(parser as any, 'parseWithCLI');
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockResolvedValue(regexResult);
const result = await parser.parse(filePath, nonDagContent);
expect(parseWithCLISpy).not.toHaveBeenCalled();
expect(parseWithRegexSpy).toHaveBeenCalledTimes(1);
expect(parseWithRegexSpy).toHaveBeenCalledWith(filePath, nonDagContent);
expect(result).toBe(regexResult);
});
});These tests assume:
parserhas methodsparse,parseWithCLI, andparseWithRegexand thatparsehas the signatureparse(filePath: string, content: string).parseinternally extractsdag_id(e.g."test_dag") from the content and passes it as the third argument toparseWithCLI(filePath, content, dagId).
If any of these assumptions differ in your codebase, adjust:
- The arguments in
toHaveBeenCalledWithexpectations to match the actualparseWithCLI/parseWithRegexsignatures. - The DAG content /
dag_idstring ("test_dag") to align with your real extraction logic.
If parseWithCLI / parseWithRegex are not instance methods on parser but are imported helpers, replace jest.spyOn(parser as any, 'parseWithCLI') with jest.spyOn on the actual module export, and update the expectations accordingly.
Comment on lines
+42
to
+51
| it('should extract tasks and edges from TaskFlow DAG', () => { | ||
| const content = ` | ||
| @dag(dag_id="test_dag") | ||
| def my_dag(): | ||
| @task | ||
| def task_a(): | ||
| return "a" | ||
|
|
||
| @task | ||
| def task_b(): |
There was a problem hiding this comment.
suggestion (testing): No tests cover << or set_upstream dependency definitions in the regex parser
parseWithRegex already supports << and set_upstream, but the tests only cover >> and set_downstream. Please add coverage for:
task_b() << task_a()and assert the edge issource: 'task_a', target: 'task_b'.task1.set_upstream(task2)and assert the edge issource: 'task2', target: 'task1'.
This will verify edge direction is handled correctly for all supported forms and guard against regressions.
Comment on lines
+57
to
+60
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' })); | ||
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' })); | ||
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' })); | ||
| }); |
There was a problem hiding this comment.
suggestion (testing): Missing tests for operator regex when task_id is omitted and when edges use variable names instead of task IDs
parseWithRegex also supports operators where task_id is omitted (e.g. task1 = BashOperator(...), using task1 as the id) and dependency resolution via either variableName or id. The new test only covers explicit task_ids. Please add tests that:
- Cover operators without
task_id(e.g.task1 = BashOperator(),task2 = BashOperator(),task1 >> task2) and assert edges are created between the inferred ids. - Cover dependencies where the variable name differs from
task_id, verifying thatgetTaskIdresolves the correct node id for both bitshift and method-based dependencies.
This will exercise the variableName vs id resolution logic and guard against regressions in dependency parsing.
Suggested change
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' })); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' })); | |
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' })); | |
| }); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' })); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' })); | |
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' })); | |
| }); | |
| it('should extract tasks from classic Operators', () => { | |
| const content = ` | |
| task1 = BashOperator(task_id="bash_task", bash_command="echo 1") | |
| task2 = PythonOperator(task_id="python_task", python_callable=my_func) | |
| task1 >> task2 | |
| `; | |
| const result = (parser as any).parseWithRegex(content, 'dag.py'); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'bash_task' })); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'python_task' })); | |
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'bash_task', target: 'python_task' })); | |
| }); | |
| it('should infer task ids from variable names when task_id is omitted', () => { | |
| const content = ` | |
| task1 = BashOperator() | |
| task2 = BashOperator() | |
| task1 >> task2 | |
| `; | |
| const result = (parser as any).parseWithRegex(content, 'dag.py'); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task1' })); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task2' })); | |
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task1', target: 'task2' })); | |
| }); | |
| it('should resolve dependencies using both variable names and task_ids', () => { | |
| const content = ` | |
| bash = BashOperator(task_id="bash_task") | |
| python = PythonOperator(task_id="python_task") | |
| # bitshift dependencies mixing variable names and task_ids | |
| bash >> "python_task" | |
| "bash_task" >> python | |
| # method-based dependencies mixing variable names and task_ids | |
| bash.set_downstream("python_task") | |
| python.set_upstream("bash_task") | |
| `; | |
| const result = (parser as any).parseWithRegex(content, 'dag.py'); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'bash_task' })); | |
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'python_task' })); | |
| expect(result.edges).toContainEqual( | |
| expect.objectContaining({ source: 'bash_task', target: 'python_task' }), | |
| ); | |
| }); |
Comment on lines
+56
to
+59
| const result = (parser as any).parseWithRegex(content, 'dag.py'); | ||
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' })); | ||
| expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' })); | ||
| expect(result.edges).toContainEqual(expect.objectContaining({ source: 'task_a', target: 'task_b' })); |
There was a problem hiding this comment.
suggestion (testing): Code snippet extraction for tasks and operators is not asserted in tests
The parser now populates data.codeDeps via extractTaskSnippet / extractOperatorSnippet (and in the CLI path limits to the first match), but the current tests don’t check this at all.
Please extend the TaskFlow and classic-operator regex tests to assert that data.codeDeps is set on the relevant nodes and that at least one snippet includes the expected def task_a or BashOperator(task_id="bash_task" text.
Optionally, add a focused unit test for one of the snippet helpers (e.g. (parser as any).extractTaskSnippet) to verify multi-line function bodies, the returned path, and snippet contents. This will help guard the new "Code Snippets" feature against future regex/indentation changes.
Suggested implementation:
const result = (parser as any).parseWithRegex(content, 'dag.py');
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_a' }));
expect(result.nodes).toContainEqual(expect.objectContaining({ id: 'task_b' }));
expect(result.edges).toContainEqual(
expect.objectContaining({ source: 'task_a', target: 'task_b' }),
);
const taskANode = result.nodes.find((n: any) => n.id === 'task_a');
expect(taskANode).toBeDefined();
expect(taskANode?.data?.codeDeps?.length).toBeGreaterThan(0);
expect(taskANode?.data?.codeDeps?.[0]).toEqual(
expect.objectContaining({
path: 'dag.py',
snippet: expect.stringContaining('def task_a'),
}),
);
});I could only see the TaskFlow regex test in the snippet, so the following additional changes are required elsewhere in tests/AirflowParser.test.ts:
-
Classic operator regex test (
BashOperator):- Locate the test that exercises the classic operator regex, likely where content includes
BashOperator(task_id="bash_task"or similar. - After the existing assertions on
result.nodes/result.edges, add expectations similar to:const bashNode = result.nodes.find((n: any) => n.id === 'bash_task'); expect(bashNode).toBeDefined(); expect(bashNode?.data?.codeDeps?.length).toBeGreaterThan(0); expect(bashNode?.data?.codeDeps?.[0]).toEqual( expect.objectContaining({ path: 'dag.py', // or the file name passed into parseWithRegex snippet: expect.stringContaining('BashOperator(task_id="bash_task"'), }), );
- Ensure you use the actual
pathargument you pass toparseWithRegexin that test (it might not be'dag.py').
- Locate the test that exercises the classic operator regex, likely where content includes
-
Focused unit test for
extractTaskSnippet(optional but recommended):- Add a new
describeblock underdescribe('AirflowParser', () => { ... }), e.g.:describe('extractTaskSnippet', () => { it('extracts multi-line TaskFlow function bodies with path and snippet', () => { const content = ` from airflow.decorators import task @task def task_a(): x = 1 y = 2 return x + y `; const parserAny = parser as any; const snippet = parserAny.extractTaskSnippet(content, 'dag.py', 'task_a'); expect(snippet).toEqual( expect.objectContaining({ path: 'dag.py', snippet: expect.stringContaining('def task_a'), }), ); expect(snippet.snippet).toContain('x = 1'); expect(snippet.snippet).toContain('return x + y'); }); });
- If
extractTaskSnippethas a different signature (e.g. returns an array or includes additional metadata), adjust the expectations accordingly but still assert:- The returned
pathmatches the passed file name. - The snippet includes
def task_aand multiple lines from the function body.
- The returned
- Add a new
These additional changes will ensure that both TaskFlow and classic operator parsing paths assert data.codeDeps population and that the snippet extraction helper is covered by focused tests.
Comment on lines
+85
to
+86
| describe('parseDot', () => { | ||
| it('should parse Airflow DOT output', () => { |
There was a problem hiding this comment.
suggestion (testing): Consider adding a test around DOT validation in the CLI flow (e.g., missing digraph)
Currently we only test the happy-path DOT input. Since parseWithCLI performs an extra stdout.includes('digraph') validation and throws otherwise, please add a test where the mocked CLI returns output without digraph. Assert that:
parseWithCLIthrows, andparsefalls back to the regex-based parsing (e.g., by spying onparseWithRegex).
This will verify the fallback behavior for malformed/non-DOT CLI output.
Suggested implementation:
expect(result.nodes).toHaveLength(2);
expect(result.nodes[0]).toEqual(expect.objectContaining({ id: 'task_1', label: 'task_1' }));
expect(result.edges).toHaveLength(1);
expect(result.edges[0]).toEqual(expect.objectContaining({ source: 'task_1', target: 'task_2' }));
});
it('should throw in parseWithCLI when CLI output is not DOT (missing digraph)', async () => {
const stdout = `
some random output
without the expected keyword
`;
const runCliSpy = jest
.spyOn(parser as any, 'runCli')
.mockResolvedValueOnce({ stdout });
await expect((parser as any).parseWithCLI('dag.py')).rejects.toThrow();
expect(runCliSpy).toHaveBeenCalled();
});
it('should fall back to regex parsing when CLI output validation fails', async () => {
const parseWithCLISpy = jest
.spyOn(parser as any, 'parseWithCLI')
.mockRejectedValueOnce(new Error('CLI output is not DOT'));
const parseWithRegexSpy = jest
.spyOn(parser as any, 'parseWithRegex')
.mockReturnValueOnce({
nodes: [{ id: 't1', label: 't1' }],
edges: [],
});
const result = await (parser as any).parse('content', 'dag.py');
expect(parseWithCLISpy).toHaveBeenCalled();
expect(parseWithRegexSpy).toHaveBeenCalledWith('content', 'dag.py');
expect(result).toEqual({
nodes: [{ id: 't1', label: 't1' }],
edges: [],
});
});- The above assumes that
AirflowParser(or whateverparseris) exposes a method namedrunClithatparseWithCLIuses internally to invoke the Airflow CLI and capture{ stdout }. If the actual helper has a different name or is accessed differently, replace'runCli'in thejest.spyOncall with the correct method name or adapt the mocking to the existing pattern in this test file (e.g., spying on acliExecutor,execwrapper, or a module-level function). - Ensure
parseWithCLIis anasyncfunction that returns a promise and performs thestdout.includes('digraph')validation; the first new test relies on it rejecting when the validation fails. - If
parseis not already usingparseWithCLIfirst and thenparseWithRegexon failure, make sure its implementation matches this behavior so the fallback test passes:parseshould callparseWithCLI,catchany error, then call and returnparseWithRegex.
- Use `execFile` for secure CLI invocation. - Implement proper source file path tracking for code snippets. - Add `CodeSnippet` interface for improved type safety. - Fix fallback logic for operator snippet extraction. - Refactor parser for better modularity and testability. - Expand test suite to cover CLI orchestration and edge cases. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
- Improved Airflow DAG detection logic to include TaskFlow decorators. - Implemented secure Airflow CLI integration using `execFile` and `airflow dags show`. - Added a robust regex-based fallback parser for static analysis. - Implemented task code snippet extraction with proper file path tracking. - Added comprehensive unit tests for AirflowParser covering orchestration and edge cases. - Introduced `CodeSnippet` interface for better type safety. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
… robustness refinements - Switched to `execFile` for secure Airflow CLI interaction. - Modularized parser logic and improved snippet extraction with proper file path tracking. - Added comprehensive orchestration and fallback tests. - Verified support for TaskFlow API, classic operators, and varied dependency definition styles (bitshift, method-based). - Integrated `CodeSnippet` interface for enhanced type safety. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces support for Apache Airflow version 2.x and 3.x.
Key features:
@dagand@taskdecorators from the TaskFlow API in addition to classicDAG()constructor calls.airflow dags showto retrieve the exact DAG structure in DOT format, ensuring high-fidelity visualization.>>,<<) and method-based dependency definitions (set_downstream,set_upstream).PR created automatically by Jules for task 15059195706509956576 started by @Resmung0
Summary by Sourcery
Add robust Airflow DAG parser that supports modern TaskFlow and classic DAG definitions, including optional CLI-assisted graph extraction and a regex-based fallback.
New Features:
Enhancements:
Tests: