Fabric Data Engineering MCP Server

/**
 * Pipeline Service
 * Manages Data Factory pipelines, runs, and monitoring
 */

import * as vscode from 'vscode';
import { FabricApiService } from '../core/api/fabric-api-service';

export interface Pipeline {
    id: string;
    displayName: string;
    description?: string;
    workspaceId: string;
    lastRunStatus?: 'Succeeded' | 'Failed' | 'InProgress' | 'Cancelled' | 'Queued';
    lastRunTime?: Date;
    definition?: PipelineDefinition;
}

export interface PipelineDefinition {
    activities: PipelineActivity[];
    parameters?: Record<string, PipelineParameter>;
    variables?: Record<string, PipelineVariable>;
}

export interface PipelineActivity {
    name: string;
    type: string;
    dependsOn?: string[];
    inputs?: ActivityInput[];
    outputs?: ActivityOutput[];
    typeProperties?: Record<string, unknown>;
}

export interface ActivityInput {
    referenceName: string;
    type: string;
}

export interface ActivityOutput {
    referenceName: string;
    type: string;
}

export interface PipelineParameter {
    type: 'String' | 'Int' | 'Float' | 'Bool' | 'Array' | 'Object';
    defaultValue?: unknown;
}

export interface PipelineVariable {
    type: 'String' | 'Bool' | 'Array';
    defaultValue?: unknown;
}

export interface PipelineRun {
    id: string;
    pipelineId: string;
    pipelineName?: string;
    status: 'Succeeded' | 'Failed' | 'InProgress' | 'Cancelled' | 'Queued';
    startTime: Date;
    endTime?: Date;
    durationMs?: number;
    errorMessage?: string;
    parameters?: Record<string, unknown>;
}

export interface ActivityRun {
    activityName: string;
    activityType: string;
    status: 'Succeeded' | 'Failed' | 'InProgress' | 'Cancelled';
    startTime: Date;
    endTime?: Date;
    durationMs?: number;
    errorMessage?: string;
    input?: Record<string, unknown>;
    output?: Record<string, unknown>;
}

export class PipelineService {
    private api: FabricApiService;
    private refreshTimers: Map<string, ReturnType<typeof setTimeout>> = new Map();
    private statusCallbacks: Map<string, (run: PipelineRun) => void> = new Map();

    constructor() {
        this.api = FabricApiService.getInstance();
    }

    async getPipelines(workspaceId: string): Promise<Pipeline[]> {
        try {
            const result = await this.api.get<{ value: Pipeline[] }>(
                `/workspaces/${workspaceId}/dataPipelines`
            );
            return result?.value || [];
        } catch (error) {
            console.error('Failed to get pipelines:', error);
            return [];
        }
    }

    async getPipeline(workspaceId: string, pipelineId: string): Promise<Pipeline | null> {
        try {
            return await this.api.get<Pipeline>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}`
            );
        } catch (error) {
            console.error('Failed to get pipeline:', error);
            return null;
        }
    }

    async createPipeline(workspaceId: string, name: string, definition?: PipelineDefinition): Promise<Pipeline | null> {
        try {
            return await this.api.post<Pipeline>(
                `/workspaces/${workspaceId}/dataPipelines`,
                {
                    displayName: name,
                    definition: definition || { activities: [] }
                }
            );
        } catch (error) {
            console.error('Failed to create pipeline:', error);
            return null;
        }
    }

    async updatePipeline(workspaceId: string, pipelineId: string, updates: Partial<Pipeline>): Promise<Pipeline | null> {
        try {
            return await this.api.patch<Pipeline>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}`,
                updates
            );
        } catch (error) {
            console.error('Failed to update pipeline:', error);
            return null;
        }
    }

    async deletePipeline(workspaceId: string, pipelineId: string): Promise<boolean> {
        try {
            await this.api.delete(`/workspaces/${workspaceId}/dataPipelines/${pipelineId}`);
            return true;
        } catch (error) {
            console.error('Failed to delete pipeline:', error);
            return false;
        }
    }

    async runPipeline(workspaceId: string, pipelineId: string, parameters?: Record<string, unknown>): Promise<PipelineRun | null> {
        try {
            const result = await this.api.post<{ id: string }>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}/jobs/instances`,
                { parameters }
            );
            if (result?.id) {
                return {
                    id: result.id,
                    pipelineId,
                    status: 'InProgress',
                    startTime: new Date()
                };
            }
            return null;
        } catch (error) {
            console.error('Failed to run pipeline:', error);
            return null;
        }
    }

    async cancelPipelineRun(workspaceId: string, pipelineId: string, runId: string): Promise<boolean> {
        try {
            await this.api.post(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}/jobs/instances/${runId}/cancel`,
                {}
            );
            return true;
        } catch (error) {
            console.error('Failed to cancel pipeline run:', error);
            return false;
        }
    }

    async getPipelineRuns(workspaceId: string, pipelineId: string, maxResults: number = 100): Promise<PipelineRun[]> {
        try {
            const result = await this.api.get<{ value: PipelineRun[] }>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}/jobs/instances`,
                { top: maxResults }
            );
            return result?.value || [];
        } catch (error) {
            console.error('Failed to get pipeline runs:', error);
            return [];
        }
    }

    async getPipelineRun(workspaceId: string, pipelineId: string, runId: string): Promise<PipelineRun | null> {
        try {
            return await this.api.get<PipelineRun>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}/jobs/instances/${runId}`
            );
        } catch (error) {
            console.error('Failed to get pipeline run:', error);
            return null;
        }
    }

    async getActivityRuns(workspaceId: string, pipelineId: string, runId: string): Promise<ActivityRun[]> {
        try {
            const result = await this.api.get<{ value: ActivityRun[] }>(
                `/workspaces/${workspaceId}/dataPipelines/${pipelineId}/jobs/instances/${runId}/activities`
            );
            return result?.value || [];
        } catch (error) {
            console.error('Failed to get activity runs:', error);
            return [];
        }
    }

    startStatusMonitoring(workspaceId: string, pipelineId: string, runId: string, callback: (run: PipelineRun) => void): void {
        const key = `${workspaceId}:${pipelineId}:${runId}`;
        this.statusCallbacks.set(key, callback);

        const poll = async () => {
            const run = await this.getPipelineRun(workspaceId, pipelineId, runId);
            if (run) {
                const cb = this.statusCallbacks.get(key);
                cb?.(run);
                
                if (run.status === 'InProgress' || run.status === 'Queued') {
                    const timer = setTimeout(poll, 5000);
                    this.refreshTimers.set(key, timer);
                } else {
                    this.stopStatusMonitoring(workspaceId, pipelineId, runId);
                }
            }
        };

        poll();
    }

    stopStatusMonitoring(workspaceId: string, pipelineId: string, runId: string): void {
        const key = `${workspaceId}:${pipelineId}:${runId}`;
        const timer = this.refreshTimers.get(key);
        if (timer) {
            clearTimeout(timer);
            this.refreshTimers.delete(key);
        }
        this.statusCallbacks.delete(key);
    }

    dispose(): void {
        for (const timer of this.refreshTimers.values()) {
            clearTimeout(timer);
        }
        this.refreshTimers.clear();
        this.statusCallbacks.clear();
    }
}