# generated by datamodel-codegen:
#   filename:  schema/metadataIngestion/workflow.json
#   timestamp: 2022-04-20T15:52:56+00:00

from __future__ import annotations

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,
    databaseServiceQueryUsagePipeline,
    messagingServiceMetadataPipeline,
)


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 WorkflowConfig(BaseModel):
    class Config:
        extra = Extra.forbid

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


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

    config: Optional[
        Union[
            databaseServiceMetadataPipeline.DatabaseServiceMetadataPipeline,
            databaseServiceQueryUsagePipeline.DatabaseServiceQueryUsagePipeline,
            dashboardServiceMetadataPipeline.DashboardServiceMetadataPipeline,
            messagingServiceMetadataPipeline.MessagingServiceMetadataPipeline,
        ]
    ] = 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: serviceConnection.ServiceConnection = Field(
        ...,
        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',
    )
