# generated by datamodel-codegen:
#   filename:  api/services/ingestionPipelines/testServiceConnection.json
#   timestamp: 2022-10-03T15:16:21+00:00

from __future__ import annotations

from enum import Enum
from typing import Optional, Union

from pydantic import BaseModel, Extra, Field

from ....entity.services import (
    dashboardService,
    databaseService,
    messagingService,
    mlmodelService,
    pipelineService,
)
from ....entity.services.connections.metadata import secretsManagerProvider


class ConnectionType(Enum):
    Database = 'Database'
    Dashboard = 'Dashboard'
    Messaging = 'Messaging'
    Pipeline = 'Pipeline'
    MlModel = 'MlModel'


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

    connection: Optional[
        Union[
            databaseService.DatabaseConnection,
            dashboardService.DashboardConnection,
            messagingService.MessagingConnection,
            pipelineService.PipelineConnection,
            mlmodelService.MlModelConnection,
        ]
    ] = Field(None, description='Connection object.')
    connectionType: Optional[ConnectionType] = Field(
        None, description='Type of service such as Database, Dashboard, Messaging, etc.'
    )
    secretsManagerProvider: Optional[
        secretsManagerProvider.SecretsManagerProvider
    ] = None
    clusterName: Optional[str] = Field(
        'openmetadata',
        description='Cluster name to differentiate OpenMetadata Server instance',
    )
