"""Models for the TableDiff test case"""

from typing import List, Optional, Union

from pydantic import BaseModel

from metadata.generated.schema.entity.data.table import (
    Column,
    Table,
    TableProfilerConfig,
)
from metadata.generated.schema.entity.services.databaseService import (
    DatabaseConnection,
    DatabaseServiceType,
)
from metadata.ingestion.models.custom_pydantic import CustomSecretStr


class TableParameter(BaseModel):
    serviceUrl: Union[str, dict]
    path: str
    columns: List[Column]
    database_service_type: DatabaseServiceType
    privateKey: Optional[CustomSecretStr]
    passPhrase: Optional[CustomSecretStr]


class TableDiffRuntimeParameters(BaseModel):
    table1: TableParameter
    table2: TableParameter
    keyColumns: List[str]
    extraColumns: List[str]
    whereClause: Optional[str]
    table_profile_config: Optional[TableProfilerConfig]


class TableCustomSQLQueryRuntimeParameters(BaseModel):
    conn_config: DatabaseConnection
    entity: Table