# generated by datamodel-codegen:
#   filename:  metadataIngestion/workflow.json
#   timestamp: 2022-10-03T15:16:21+00:00

from __future__ import annotations

from enum import Enum
from typing import Any, Optional, Union

from pydantic import BaseModel, Extra, Field

from ..entity.services.connections import serviceConnection
from ..entity.services.connections.metadata import openMetadataConnection
from ..type import basic
from . import (
    dashboardServiceMetadataPipeline,
    databaseServiceMetadataPipeline,
    databaseServiceProfilerPipeline,
    databaseServiceQueryLineagePipeline,
    databaseServiceQueryUsagePipeline,
    messagingServiceMetadataPipeline,
    mlmodelServiceMetadataPipeline,
    pipelineServiceMetadataPipeline,
    testSuitePipeline,
)


class ComponentConfig(BaseModel):
    pass

    class Config:
        extra = Extra.allow


class Processor(BaseModel):
    class Config:
        extra = Extra.forbid

    type: str = Field(..., description='Type of processor component ex: pii-processor')
    config: Optional[ComponentConfig] = None


class Stage(BaseModel):
    class Config:
        extra = Extra.forbid

    type: str = Field(..., description='Type of stage component ex: table-usage')
    config: Optional[ComponentConfig] = None


class Sink(BaseModel):
    class Config:
        extra = Extra.forbid

    type: str = Field(..., description='Type of sink component ex: metadata')
    config: Optional[ComponentConfig] = None


class BulkSink(BaseModel):
    class Config:
        extra = Extra.forbid

    type: str = Field(..., description='Type of BulkSink component ex: metadata-usage')
    config: Optional[ComponentConfig] = None


class LogLevels(Enum):
    DEBUG = 'DEBUG'
    INFO = 'INFO'
    WARN = 'WARN'
    ERROR = 'ERROR'


class WorkflowConfig(BaseModel):
    class Config:
        extra = Extra.forbid

    loggerLevel: Optional[LogLevels] = LogLevels.INFO
    openMetadataServerConfig: openMetadataConnection.OpenMetadataConnection
    config: Optional[ComponentConfig] = None


class SourceConfig(BaseModel):
    class Config:
        extra = Extra.forbid

    config: Optional[
        Union[
            databaseServiceMetadataPipeline.DatabaseServiceMetadataPipeline,
            databaseServiceQueryUsagePipeline.DatabaseServiceQueryUsagePipeline,
            databaseServiceQueryLineagePipeline.DatabaseServiceQueryLineagePipeline,
            dashboardServiceMetadataPipeline.DashboardServiceMetadataPipeline,
            messagingServiceMetadataPipeline.MessagingServiceMetadataPipeline,
            databaseServiceProfilerPipeline.DatabaseServiceProfilerPipeline,
            pipelineServiceMetadataPipeline.PipelineServiceMetadataPipeline,
            mlmodelServiceMetadataPipeline.MlModelServiceMetadataPipeline,
            testSuitePipeline.TestSuitePipeline,
        ]
    ] = None


class Source(BaseModel):
    class Config:
        extra = Extra.forbid

    type: str = Field(
        ...,
        description='Type of the source connector ex: mysql, snowflake, tableau etc..',
    )
    serviceName: str = Field(
        ...,
        description='Type of the source connector ex: mysql, snowflake, tableau etc..',
    )
    serviceConnection: Optional[serviceConnection.ServiceConnection] = Field(
        None,
        description='Connection configuration for the source. ex: mysql , tableau connection.',
    )
    sourceConfig: SourceConfig


class OpenMetadataWorkflowConfig(BaseModel):
    class Config:
        extra = Extra.forbid

    source: Source
    processor: Optional[Processor] = None
    sink: Optional[Sink] = None
    stage: Optional[Stage] = None
    bulkSink: Optional[BulkSink] = None
    workflowConfig: WorkflowConfig


class MetadataWorkflow1(BaseModel):
    class Config:
        extra = Extra.forbid

    id: Optional[basic.Uuid] = Field(
        None, description='Unique identifier that identifies this pipeline.'
    )
    name: Optional[basic.EntityName] = Field(
        None, description='Name that identifies this pipeline instance uniquely.'
    )
    openMetadataWorkflowConfig: Optional[OpenMetadataWorkflowConfig] = Field(
        None, description='OpenMetadata Ingestion Workflow Config.'
    )


class MetadataWorkflow(BaseModel):
    class Config:
        extra = Extra.forbid

    __root__: Union[MetadataWorkflow1, Any, Any] = Field(
        ...,
        description='OpenMetadata Ingestion Framework definition.',
        title='MetadataWorkflow',
    )
