# generated by datamodel-codegen:
#   filename:  entities.yaml
#   timestamp: 2023-02-13T15:29:46+00:00

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import AnyUrl, BaseModel, Field


class Tag(BaseModel):
    name: str


class BooleanFieldStat(BaseModel):
    true_count: int
    false_count: int
    nulls_count: int


class ComplexFieldStat(BaseModel):
    nulls_count: int
    unique_count: int


class NumberFieldStat(BaseModel):
    low_value: float
    high_value: float
    mean_value: Optional[float] = None
    median_value: Optional[float] = None
    nulls_count: int
    unique_count: int


class StringFieldStat(BaseModel):
    max_length: int
    avg_length: float
    nulls_count: int
    unique_count: int


class BinaryFieldStat(BaseModel):
    max_length: int
    avg_length: float
    nulls_count: int
    unique_count: int


class IntegerFieldStat(BaseModel):
    low_value: int
    high_value: int
    mean_value: Optional[int] = None
    median_value: Optional[int] = None
    nulls_count: int
    unique_count: int


class DateTimeFieldStat(BaseModel):
    low_value: datetime
    high_value: datetime
    mean_value: Optional[datetime] = None
    median_value: Optional[datetime] = None
    nulls_count: int
    unique_count: int


class Type(Enum):
    TYPE_STRING = 'TYPE_STRING'
    TYPE_NUMBER = 'TYPE_NUMBER'
    TYPE_INTEGER = 'TYPE_INTEGER'
    TYPE_BOOLEAN = 'TYPE_BOOLEAN'
    TYPE_CHAR = 'TYPE_CHAR'
    TYPE_DATETIME = 'TYPE_DATETIME'
    TYPE_TIME = 'TYPE_TIME'
    TYPE_STRUCT = 'TYPE_STRUCT'
    TYPE_BINARY = 'TYPE_BINARY'
    TYPE_LIST = 'TYPE_LIST'
    TYPE_MAP = 'TYPE_MAP'
    TYPE_UNION = 'TYPE_UNION'
    TYPE_DURATION = 'TYPE_DURATION'
    TYPE_UNKNOWN = 'TYPE_UNKNOWN'


class DataSetFieldType(BaseModel):
    type: Type
    logical_type: Optional[str] = None
    is_nullable: bool


class DataTransformer(BaseModel):
    source_code_url: Optional[str] = None
    sql: Optional[str] = None
    inputs: List[str]
    outputs: List[str]


class DataConsumer(BaseModel):
    inputs: Optional[List[str]] = None


class DataInput(BaseModel):
    outputs: Optional[List[str]] = None


class DataEntityGroup(BaseModel):
    entities_list: List[str]
    group_oddrn: Optional[str] = None


class DataQualityTestExpectation(BaseModel):
    type: Optional[str] = Field(None, example='expect_table_row_count_to_be_between')


class LinkedUrl(BaseModel):
    name: str
    url: str


class JobRunStatus(Enum):
    SUCCESS = 'SUCCESS'
    FAILED = 'FAILED'
    SKIPPED = 'SKIPPED'
    BROKEN = 'BROKEN'
    ABORTED = 'ABORTED'
    RUNNING = 'RUNNING'
    UNKNOWN = 'UNKNOWN'


class QualityRunStatus(Enum):
    SUCCESS = 'SUCCESS'
    FAILED = 'FAILED'
    SKIPPED = 'SKIPPED'
    BROKEN = 'BROKEN'
    ABORTED = 'ABORTED'
    RUNNING = 'RUNNING'
    UNKNOWN = 'UNKNOWN'


class DataTransformerRun(BaseModel):
    transformer_oddrn: str
    start_time: datetime
    end_time: Optional[datetime] = None
    status_reason: Optional[str] = None
    status: JobRunStatus


class DataQualityTestRun(BaseModel):
    data_quality_test_oddrn: str
    start_time: datetime
    end_time: Optional[datetime] = None
    status_reason: Optional[str] = None
    status: QualityRunStatus


class DataEntityType(Enum):
    TABLE = 'TABLE'
    VIEW = 'VIEW'
    FILE = 'FILE'
    FEATURE_GROUP = 'FEATURE_GROUP'
    KAFKA_TOPIC = 'KAFKA_TOPIC'
    JOB = 'JOB'
    JOB_RUN = 'JOB_RUN'
    ML_MODEL = 'ML_MODEL'
    ML_MODEL_TRAINING = 'ML_MODEL_TRAINING'
    DASHBOARD = 'DASHBOARD'
    DAG = 'DAG'
    ML_EXPERIMENT = 'ML_EXPERIMENT'
    GRAPH_NODE = 'GRAPH_NODE'
    UNKNOWN = 'UNKNOWN'
    MICROSERVICE = 'MICROSERVICE'
    API_CALL = 'API_CALL'
    DATABASE_SERVICE = 'DATABASE_SERVICE'
    API_SERVICE = 'API_SERVICE'
    KAFKA_SERVICE = 'KAFKA_SERVICE'


class MetadataExtension(BaseModel):
    schema_url: AnyUrl = Field(
        ...,
        description='The JSON Pointer (https://tools.ietf.org/html/rfc6901) URL to the corresponding version of the schema definition for this extension',
        example='https://raw.githubusercontent.com/opendatadiscovery/opendatadiscovery-specification/main/specification/extensions/glue.json#/definitions/GlueDataSetExtension',
    )
    metadata: Dict[str, Any]


class DataSource(BaseModel):
    oddrn: str
    name: str
    description: Optional[str] = None


class DataSourceList(BaseModel):
    items: List[DataSource]


class IngestionAlertType(Enum):
    DATA_SOURCE_CONNECTION_ERROR = 'DATA_SOURCE_CONNECTION_ERROR'
    DATA_SOURCE_GENERIC_ERROR = 'DATA_SOURCE_GENERIC_ERROR'


class IngestionAlert(BaseModel):
    description: str
    type: IngestionAlertType
    related_data_source_oddrn: str
    related_data_entity_oddrn: Optional[str] = None


class IngestionAlertList(BaseModel):
    items: List[IngestionAlert]


class CompactDataEntity(BaseModel):
    oddrn: str
    type: DataEntityType


class CompactDataEntityList(BaseModel):
    items: List[CompactDataEntity]


class BaseObject(BaseModel):
    oddrn: str = Field(..., example='//aws/glue/{account_id}/{database}/{tablename}')
    name: str
    version: Optional[str] = None
    description: Optional[str] = None
    owner: Optional[str] = Field(None, example='//aws/iam/{account_id}/user/name')
    metadata: Optional[List[MetadataExtension]] = None
    tags: Optional[List[Tag]] = None


class DataSetFieldStat(BaseModel):
    tags: Optional[List[Tag]] = None
    complex_stats: Optional[ComplexFieldStat] = None
    boolean_stats: Optional[BooleanFieldStat] = None
    integer_stats: Optional[IntegerFieldStat] = None
    number_stats: Optional[NumberFieldStat] = None
    string_stats: Optional[StringFieldStat] = None
    binary_stats: Optional[BinaryFieldStat] = None
    datetime_stats: Optional[DateTimeFieldStat] = None


class DataSetField(BaseObject):
    parent_field_oddrn: Optional[str] = None
    type: DataSetFieldType
    is_primary_key: Optional[bool] = None
    is_sort_key: Optional[bool] = None
    is_key: Optional[bool] = None
    is_value: Optional[bool] = None
    default_value: Optional[str] = None
    stats: Optional[DataSetFieldStat] = None


class DataQualityTest(BaseModel):
    suite_name: str
    dataset_list: List[str]
    expectation: DataQualityTestExpectation
    suite_url: Optional[str] = None
    linked_url_list: Optional[List[LinkedUrl]] = None


class DataSet(BaseModel):
    parent_oddrn: Optional[str] = None
    rows_number: Optional[int] = None
    field_list: List[DataSetField]


class DataSetStatistics(BaseModel):
    dataset_oddrn: str
    fields: Dict[str, DataSetFieldStat] = Field(
        ...,
        description='key is a field oddrn. For example: "//aws/glue/{account_id}/{database}/{tablename}/columns/{column_name}"\n',
    )


class DataEntity(BaseObject):
    updated_at: Optional[datetime] = None
    created_at: Optional[datetime] = None
    type: DataEntityType
    dataset: Optional[DataSet] = None
    data_transformer: Optional[DataTransformer] = None
    data_transformer_run: Optional[DataTransformerRun] = None
    data_quality_test: Optional[DataQualityTest] = None
    data_quality_test_run: Optional[DataQualityTestRun] = None
    data_input: Optional[DataInput] = None
    data_consumer: Optional[DataConsumer] = None
    data_entity_group: Optional[DataEntityGroup] = None


class DataEntityList(BaseModel):
    data_source_oddrn: str = Field(..., example='//aws/glue/123456789010/')
    items: Optional[List[DataEntity]] = None


class DatasetStatisticsList(BaseModel):
    items: List[DataSetStatistics]
