# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: pql.proto
# plugin: python-betterproto
from dataclasses import dataclass
from typing import Dict, List

import betterproto


class ObjectType(betterproto.Enum):
    """Top level object type used in Clause, as well as privilege"""

    UNDEFINED = 0
    ROLE = 1
    ENGINE = 2
    CONNECTION = 3
    DATASET = 4
    MODEL = 5
    SCHEDULE = 6
    ENVIRONMENT = 7
    DEPLOYMENT = 8


class ParseContextDialect(betterproto.Enum):
    DEFAULT_ANSI = 0
    MYSQL = 1
    POSTGRESQL = 2
    SNOWFLAKE = 3
    REDSHIFT = 4
    BIGQUERY = 5
    DATABRICKS = 6


class ParseStatementClauseType(betterproto.Enum):
    UNDEFINED = 0
    SELECT = 1
    JOIN = 2
    INSERT = 3
    WITH = 4
    UNION = 5
    UNION_ALL = 6
    INTERSECT = 7
    ORDER_BY = 8
    CREATE_TABLE = 9
    CREATE_ROLE = 10
    CREATE_ENGINE = 11
    CREATE_CONNECTION = 12
    CREATE_DATASET = 13
    CREATE_MODEL = 14
    RETRAIN_MODEL = 15
    FORK_MODEL = 16
    ALTER_MODEL = 17
    CREATE_SCHEDULE = 18
    DROP_ROLE = 20
    DROP_ENGINE = 21
    DROP_CONNECTION = 22
    DROP_DATASET = 23
    DROP_MODEL = 24
    DROP_SCHEDULE = 25
    GRANT = 50
    REVOKE = 51
    PREDICT = 100
    EVALUATE = 101
    VISUALIZE = 102
    FORECAST = 103
    EXPORT_MODEL = 110
    SHOW_ROLES = 201
    SHOW_GRANTS = 202
    SHOW_ENGINES = 203
    SHOW_CONNECTIONS = 204
    SHOW_DATASETS = 205
    SHOW_MODELS = 206
    SHOW_METRICS = 207
    SHOW_DEPLOYMENTS = 208
    USE_CONNECTION = 250
    USE_ENGINE = 251
    USE_MODEL = 252
    USE_ENVIRONMENT = 253
    CREATE_ENVIRONMENT = 301
    DROP_ENVIRONMENT = 302
    ALTER_ENVIRONMENT = 303
    CREATE_DEPLOYMENT = 321
    DROP_DEPLOYMENT = 322
    ALTER_DEPLOYMENT = 323
    MOVE_DEPLOYMENT = 324
    DESCRIBE_DEPLOYMENT = 330
    APPROVE_DEPLOYMENT = 341
    REJECT_DEPLOYMENT = 342
    REVERT_DEPLOYMENT = 343


class CreateConnectionClauseConnectionType(betterproto.Enum):
    UNDEFINED = 0
    ADLS = 1
    GCS = 2
    S3 = 3
    MYSQL = 4
    POSTGRESQL = 5
    SNOWFLAKE = 6
    REDSHIFT = 7
    BIGQUERY = 8
    DATABRICKS = 9


class ModelClauseActionType(betterproto.Enum):
    CREATE = 0
    RETRAIN = 1
    FORK = 2
    ALTER = 3


class FeatureFeatureType(betterproto.Enum):
    UNDEFINED = 0
    BINARY = 1
    NUMBER = 2
    CATEGORY = 3
    SET = 4
    BAG = 5
    SEQUENCE = 6
    TEXT = 7
    TIMESERIES = 8
    AUDIO = 9
    IMAGE = 10
    DATE = 11
    H3 = 12
    VECTOR = 13


class CreateEngineClauseEngineType(betterproto.Enum):
    GENERAL = 0
    SERVING = 1


class PredictClausePredictType(betterproto.Enum):
    UNDEFINED = 0
    PREDICT = 1
    EVALUATE = 2


class PredictTargetPredictPropertyType(betterproto.Enum):
    EMPTY = 0
    EXPLANATION = 1
    CONFIDENCE = 2
    PROBABILITIES = 3
    METADATA = 4


class InputInputType(betterproto.Enum):
    UNDEFINED = 0
    IDENTIFIER = 1
    BINARY = 2
    NUMBER = 3
    STRING = 4
    ARRAY = 5
    SAMPLE_CHOICE = 6
    SAMPLE_GRID = 7
    RANGE_INT = 8
    RANGE_FLOAT = 9
    SPLIT = 10


class InputScaleType(betterproto.Enum):
    AUTO = 0
    LOG = 1
    LINEAR = 2


class SplitSpecSplitType(betterproto.Enum):
    UNDEFINED = 0
    RANDOM = 1
    FIXED = 2
    STRATIFY = 3


class VisualizeClauseVisualizeType(betterproto.Enum):
    UNDEFINED = 0
    ROC_CURVES = 1
    CONFUSION_MATRIX = 2
    RESIDUALS = 3
    LEARNING_CURVES = 4


class PrivilegePrivilegeType(betterproto.Enum):
    NONE = 0
    READ = 1
    USAGE = 2
    WRITE = 5
    OWN = 15


class TransitionDeploymentClauseTransitionType(betterproto.Enum):
    NONE = 0
    APPROVE = 1
    REJECT = 2
    REVERT = 3


@dataclass
class ParseContext(betterproto.Message):
    source_dialect: "ParseContextDialect" = betterproto.enum_field(1)
    target_dialect: "ParseContextDialect" = betterproto.enum_field(2)
    rewrite: bool = betterproto.bool_field(3)
    default_connection_name: str = betterproto.string_field(4)
    connection_info: Dict[str, "ConnectionInfo"] = betterproto.map_field(
        5, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
    )
    dataset_mappings: List["DatasetMapping"] = betterproto.message_field(6)
    target_mappings: List["TargetMapping"] = betterproto.message_field(7)
    default_limit: int = betterproto.int32_field(8)


@dataclass
class ConnectionInfo(betterproto.Message):
    name: str = betterproto.string_field(1)
    type: str = betterproto.string_field(2)
    config: str = betterproto.string_field(3)
    uuid: str = betterproto.string_field(4)


@dataclass
class DatasetMapping(betterproto.Message):
    source: "DatasetRef" = betterproto.message_field(1)
    dialect: "ParseContextDialect" = betterproto.enum_field(2)
    object_name: str = betterproto.string_field(3)
    query_sql: str = betterproto.string_field(4)
    column_name_to_connection_name: Dict[str, str] = betterproto.map_field(
        5, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
    dataset_id: int = betterproto.uint64_field(6)
    data_format: str = betterproto.string_field(7)


@dataclass
class TargetMapping(betterproto.Message):
    target: "TargetRef" = betterproto.message_field(1)
    model: "ModelRef" = betterproto.message_field(2)
    feature_names: List[str] = betterproto.string_field(3)
    column_name: str = betterproto.string_field(4)
    model_id: int = betterproto.uint64_field(5)


@dataclass
class ParseStatement(betterproto.Message):
    clause_type: "ParseStatementClauseType" = betterproto.enum_field(1)
    parsed_sql: str = betterproto.string_field(2)
    sources: List["DatasetRef"] = betterproto.message_field(3)
    aliases: Dict[str, int] = betterproto.map_field(
        4, betterproto.TYPE_STRING, betterproto.TYPE_INT32
    )
    dependencies: List["NativeSqlDependency"] = betterproto.message_field(5)
    sink: "DatasetRef" = betterproto.message_field(6)
    # oneof clause {
    native_sql: "NativeSqlClause" = betterproto.message_field(10)
    create_object: "CreateObjectClause" = betterproto.message_field(11)
    drop_object: "DropObjectClause" = betterproto.message_field(12)
    grant: "PrivilegeClause" = betterproto.message_field(13)
    revoke: "PrivilegeClause" = betterproto.message_field(14)
    predict: "PredictClause" = betterproto.message_field(15)
    show_object: "ShowObjectClause" = betterproto.message_field(16)
    export_model: "ExportModelClause" = betterproto.message_field(17)
    visualize: "VisualizeClause" = betterproto.message_field(18)
    use_object: "UseObjectClause" = betterproto.message_field(19)
    alter_object: "AlterObjectClause" = betterproto.message_field(20)
    transition_deployment: "TransitionDeploymentClause" = betterproto.message_field(21)
    forecast: "ForecastClause" = betterproto.message_field(22)


@dataclass
class NativeSqlClause(betterproto.Message):
    sql: str = betterproto.string_field(1)
    connection_names: List[str] = betterproto.string_field(2)
    column_names: List["TargetRef"] = betterproto.message_field(3)


@dataclass
class NativeSqlDependency(betterproto.Message):
    alias: str = betterproto.string_field(1)
    clause_type: "ParseStatementClauseType" = betterproto.enum_field(2)
    # oneof clause {
    native_sql: "NativeSqlClause" = betterproto.message_field(3)
    predict: "PredictClause" = betterproto.message_field(4)
    show_object: "ShowObjectClause" = betterproto.message_field(5)


@dataclass
class DropObjectClause(betterproto.Message):
    object_type: "ObjectType" = betterproto.enum_field(1)
    if_exists: bool = betterproto.bool_field(2)
    object_ref: "ObjectRef" = betterproto.message_field(3)


@dataclass
class ObjectRef(betterproto.Message):
    connection_name: str = betterproto.string_field(2)
    name: str = betterproto.string_field(3)


@dataclass
class AlterObjectClause(betterproto.Message):
    object_type: "ObjectType" = betterproto.enum_field(1)
    if_exists: bool = betterproto.bool_field(2)
    # oneof object_name {
    alter_model: "ModelClause" = betterproto.message_field(10)
    alter_deployment: "AlterDeploymentClause" = betterproto.message_field(11)
    move_deployment: "MoveDeploymentClause" = betterproto.message_field(12)


@dataclass
class CreateObjectClause(betterproto.Message):
    object_type: "ObjectType" = betterproto.enum_field(1)
    replace: bool = betterproto.bool_field(2)
    if_not_exists: bool = betterproto.bool_field(3)
    # oneof object_name {
    create_role: "CreateRoleClause" = betterproto.message_field(4)
    create_engine: "CreateEngineClause" = betterproto.message_field(5)
    create_connection: "CreateConnectionClause" = betterproto.message_field(6)
    create_dataset: "CreateDatasetClause" = betterproto.message_field(7)
    create_model: "ModelClause" = betterproto.message_field(10)
    create_schedule: "CreateScheduleClause" = betterproto.message_field(9)
    create_environment: "CreateEnvironmentClause" = betterproto.message_field(11)
    create_deployment: "CreateDeploymentClause" = betterproto.message_field(12)


@dataclass
class CreateConnectionClause(betterproto.Message):
    type: "CreateConnectionClauseConnectionType" = betterproto.enum_field(1)
    name: str = betterproto.string_field(2)
    # oneof properties { // Index from 10 matching connection type
    adls_properties: "AdlsProperties" = betterproto.message_field(11)
    gcs_properties: "GcsProperties" = betterproto.message_field(12)
    s3_properties: "S3Properties" = betterproto.message_field(13)
    mysql_properties: "DatabaseProperties" = betterproto.message_field(14)
    postgres_properties: "DatabaseProperties" = betterproto.message_field(15)
    snowflake_properties: "SnowflakeProperties" = betterproto.message_field(16)
    redshift_properties: "DatabaseProperties" = betterproto.message_field(17)
    big_query_properties: "BigQueryProperties" = betterproto.message_field(18)
    databricks_properties: "DatabricksProperties" = betterproto.message_field(19)


@dataclass
class AdlsProperties(betterproto.Message):
    account_key: str = betterproto.string_field(1)
    connection_string: str = betterproto.string_field(2)
    protocol: str = betterproto.string_field(3)
    url: str = betterproto.string_field(4)


@dataclass
class GcsProperties(betterproto.Message):
    token: str = betterproto.string_field(1)
    url: str = betterproto.string_field(2)


@dataclass
class S3Properties(betterproto.Message):
    aws_access_key_id: str = betterproto.string_field(1)
    aws_secret_access_key: str = betterproto.string_field(2)
    aws_role_arn: str = betterproto.string_field(3)
    url: str = betterproto.string_field(4)


@dataclass
class DatabaseProperties(betterproto.Message):
    address: str = betterproto.string_field(1)
    database: str = betterproto.string_field(2)
    username: str = betterproto.string_field(3)
    password: str = betterproto.string_field(4)


@dataclass
class SnowflakeProperties(betterproto.Message):
    account: str = betterproto.string_field(1)
    warehouse: str = betterproto.string_field(2)
    database: str = betterproto.string_field(3)
    schema: str = betterproto.string_field(4)
    username: str = betterproto.string_field(5)
    password: str = betterproto.string_field(6)


@dataclass
class BigQueryProperties(betterproto.Message):
    token: str = betterproto.string_field(1)
    project: str = betterproto.string_field(2)
    dataset: str = betterproto.string_field(3)


@dataclass
class DatabricksProperties(betterproto.Message):
    access_token: str = betterproto.string_field(1)
    server_host_name: str = betterproto.string_field(2)
    http_path: str = betterproto.string_field(3)
    schema_name: str = betterproto.string_field(4)


@dataclass
class CreateDatasetClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    # oneof from {
    object_names: List[str] = betterproto.string_field(2)
    path: str = betterproto.string_field(3)
    format_properties: Dict[str, str] = betterproto.map_field(
        4, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
    query: "NativeSqlClause" = betterproto.message_field(5)


@dataclass
class DatasetRef(betterproto.Message):
    connection_name: str = betterproto.string_field(1)
    name: str = betterproto.string_field(2)


@dataclass
class ModelClause(betterproto.Message):
    action_type: "ModelClauseActionType" = betterproto.enum_field(1)
    parent_model: "ModelRef" = betterproto.message_field(2)
    name: str = betterproto.string_field(3)
    string_config: str = betterproto.string_field(4)
    model_config: "ModelConfig" = betterproto.message_field(5)
    # oneof from {
    source: "DatasetRef" = betterproto.message_field(6)
    query: "NativeSqlClause" = betterproto.message_field(7)


@dataclass
class CreateScheduleClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    start_date_time: str = betterproto.string_field(2)
    cron_expression: str = betterproto.string_field(3)
    clause_type: "ParseStatementClauseType" = betterproto.enum_field(4)
    # oneof as {
    native_sql: "NativeSqlClause" = betterproto.message_field(5)
    predict: "PredictClause" = betterproto.message_field(6)


@dataclass
class ModelConfig(betterproto.Message):
    input_features: List["Feature"] = betterproto.message_field(1)
    output_features: List["Feature"] = betterproto.message_field(2)
    preprocessing: List["Input"] = betterproto.message_field(3)
    defaults: List["Input"] = betterproto.message_field(7)
    combiner: List["Input"] = betterproto.message_field(4)
    trainer: List["Input"] = betterproto.message_field(5)
    hyperopt: List["Input"] = betterproto.message_field(6)


@dataclass
class Feature(betterproto.Message):
    name: str = betterproto.string_field(1)
    type: "FeatureFeatureType" = betterproto.enum_field(2)
    preprocessing: List["Input"] = betterproto.message_field(3)
    encoder: List["Input"] = betterproto.message_field(4)
    decoder: List["Input"] = betterproto.message_field(5)


@dataclass
class CreateRoleClause(betterproto.Message):
    name: str = betterproto.string_field(1)


@dataclass
class CreateEngineClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    engine_type: "CreateEngineClauseEngineType" = betterproto.enum_field(2)
    template_name: str = betterproto.string_field(3)
    auto_suspend: int = betterproto.uint32_field(4)
    auto_resume: bool = betterproto.bool_field(5)
    compute_properties: Dict[str, str] = betterproto.map_field(
        6, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
    scaling_properties: Dict[str, str] = betterproto.map_field(
        7, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
    logging_properties: Dict[str, str] = betterproto.map_field(
        8, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )


@dataclass
class PredictClause(betterproto.Message):
    predict_type: "PredictClausePredictType" = betterproto.enum_field(1)
    targets: List["PredictTarget"] = betterproto.message_field(2)
    inputs: List["PredictInput"] = betterproto.message_field(3)


@dataclass
class ForecastClause(betterproto.Message):
    horizon: int = betterproto.uint32_field(1)
    targets: List["PredictTarget"] = betterproto.message_field(2)
    inputs: List["PredictInput"] = betterproto.message_field(3)


@dataclass
class PredictInput(betterproto.Message):
    # oneof {
    constant_input: "Input" = betterproto.message_field(1)
    native_sql: "NativeSqlClause" = betterproto.message_field(2)


@dataclass
class PredictTarget(betterproto.Message):
    target: "TargetRef" = betterproto.message_field(1)
    alias: str = betterproto.string_field(2)
    properties: List["PredictTargetPredictPropertyType"] = betterproto.enum_field(3)
    model: "ModelRef" = betterproto.message_field(4)
    deployment: "DeploymentRef" = betterproto.message_field(7)
    categories: List[str] = betterproto.string_field(5)
    format_properties: Dict[str, str] = betterproto.map_field(
        6, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )


@dataclass
class TargetRef(betterproto.Message):
    connection_name: str = betterproto.string_field(1)
    dataset_name: str = betterproto.string_field(2)
    name: str = betterproto.string_field(3)


@dataclass
class ModelRef(betterproto.Message):
    connection_name: str = betterproto.string_field(1)
    name: str = betterproto.string_field(2)
    version: int = betterproto.int32_field(3)
    uuid: str = betterproto.string_field(4)
    run_id: str = betterproto.string_field(5)


@dataclass
class DeploymentRef(betterproto.Message):
    connection_name: str = betterproto.string_field(1)
    name: str = betterproto.string_field(2)
    version: int = betterproto.int32_field(3)
    uuid: str = betterproto.string_field(4)


@dataclass
class Input(betterproto.Message):
    name: str = betterproto.string_field(1)
    type: "InputInputType" = betterproto.enum_field(2)
    values: List[str] = betterproto.string_field(3)
    # oneof input_props {
    split_props: "SplitSpec" = betterproto.message_field(4)
    # }
    scale_type: "InputScaleType" = betterproto.enum_field(10)


@dataclass
class SplitSpec(betterproto.Message):
    type: "SplitSpecSplitType" = betterproto.enum_field(1)
    column: str = betterproto.string_field(2)
    probabilities: List[float] = betterproto.double_field(3)


@dataclass
class ExportModelClause(betterproto.Message):
    model: "ModelRef" = betterproto.message_field(1)
    path: str = betterproto.string_field(2)
    format_properties: Dict[str, str] = betterproto.map_field(
        4, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )


@dataclass
class VisualizeClause(betterproto.Message):
    type: "VisualizeClauseVisualizeType" = betterproto.enum_field(1)
    format_properties: Dict[str, str] = betterproto.map_field(
        2, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
    metrics: List[str] = betterproto.string_field(3)
    targets: List["TargetRef"] = betterproto.message_field(4)
    models: List["ModelRef"] = betterproto.message_field(5)
    query: "NativeSqlClause" = betterproto.message_field(6)


@dataclass
class Privilege(betterproto.Message):
    privilege_types: List["PrivilegePrivilegeType"] = betterproto.enum_field(1)
    object_type: "ObjectType" = betterproto.enum_field(2)
    object_refs: List["ObjectRef"] = betterproto.message_field(3)


@dataclass
class RoleOrUser(betterproto.Message):
    user: str = betterproto.string_field(1)
    role: str = betterproto.string_field(2)


@dataclass
class PrivilegeClause(betterproto.Message):
    """Great privilege or roles to (user or role)"""

    privilege: "Privilege" = betterproto.message_field(1)
    roles: List[str] = betterproto.string_field(2)
    roles_reserved: List[bool] = betterproto.bool_field(3)
    target: "RoleOrUser" = betterproto.message_field(4)


@dataclass
class ShowObjectClause(betterproto.Message):
    like: str = betterproto.string_field(2)
    # Nested properties for special show properties  oneof show_props {
    show_grants: "ShowGrantsClause" = betterproto.message_field(3)
    show_models: "ShowModelsClause" = betterproto.message_field(4)
    show_metrics: "ShowMetricsClause" = betterproto.message_field(5)
    show_deployments: "ShowDeploymentsClause" = betterproto.message_field(6)


@dataclass
class ShowGrantsClause(betterproto.Message):
    object_type: "ObjectType" = betterproto.enum_field(1)
    object_ref: "ObjectRef" = betterproto.message_field(2)
    to: "RoleOrUser" = betterproto.message_field(3)
    of: "RoleOrUser" = betterproto.message_field(4)


@dataclass
class ShowModelsClause(betterproto.Message):
    source: "DatasetRef" = betterproto.message_field(1)


@dataclass
class ShowMetricsClause(betterproto.Message):
    metrics: List[str] = betterproto.string_field(1)
    targets: List["TargetRef"] = betterproto.message_field(2)
    models: List["ModelRef"] = betterproto.message_field(3)


@dataclass
class ShowDeploymentsClause(betterproto.Message):
    engine_name: str = betterproto.string_field(1)
    model: "ModelRef" = betterproto.message_field(2)


@dataclass
class UseObjectClause(betterproto.Message):
    object_type: "ObjectType" = betterproto.enum_field(1)
    # oneof use_props {
    connection_name: str = betterproto.string_field(2)
    engine_name: str = betterproto.string_field(3)
    model: "ModelRef" = betterproto.message_field(4)
    environment_name: str = betterproto.string_field(5)


@dataclass
class CreateEnvironmentClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    connection_name: str = betterproto.string_field(2)
    url: str = betterproto.string_field(3)
    depends_on: str = betterproto.string_field(4)
    where_condition: str = betterproto.string_field(5)
    approvals: List["RoleOrUser"] = betterproto.message_field(6)
    default_properties: Dict[str, str] = betterproto.map_field(
        7, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )


@dataclass
class CreateDeploymentClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    engine_name: str = betterproto.string_field(2)
    deployment_models: List["DeploymentModel"] = betterproto.message_field(4)
    comment: str = betterproto.string_field(5)


@dataclass
class AlterDeploymentClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    deployment_models: List["DeploymentModel"] = betterproto.message_field(2)


@dataclass
class MoveDeploymentClause(betterproto.Message):
    name: str = betterproto.string_field(1)
    engine_name: str = betterproto.string_field(2)
    comment: str = betterproto.string_field(3)


@dataclass
class TransitionDeploymentClause(betterproto.Message):
    action: "TransitionDeploymentClauseTransitionType" = betterproto.enum_field(1)
    deployment: "DeploymentRef" = betterproto.message_field(2)
    comment: str = betterproto.string_field(3)


@dataclass
class DeploymentModel(betterproto.Message):
    model: "ModelRef" = betterproto.message_field(1)
    shadow: bool = betterproto.bool_field(2)
    weight: int = betterproto.uint32_field(3)
