Update a Source Schema

from typing import Any

from validio_sdk.client import Client


async def update_source_jtd_schema(source_id: str):
    client = Client()
    await update_bigquery_source_jtd_schema(client, source_id)

    print(f"Updated schema for source with id {source_id}")


async def get_jtd_schema(
    client: Client,
    credential_id: str,
    project: str,
    dataset: str,
    table: str,
) -> dict[str, Any]:
    query = """
    query GcpBigQueryInferSchema($input: GcpBigQueryInferSchemaInput!) {
      gcpBigQueryInferSchema(input: $input)
    }
    """
    variable_values = {
        "input": {
            "credentialId": credential_id,
            "project": project,
            "dataset": dataset,
            "table": table,
        }
    }

    result = await client.execute(
        query,
        variable_values=variable_values,
    )

    return result["gcpBigQueryInferSchema"]


async def verify_jtd_schema(
    client: Client,
    source_id: str,
    jtd_schema: dict[str, Any],
) -> None:
    mutation = """
    mutation SourceSchemaVerify($id: SourceId!, $input: JsonTypeDefinition!) {
      sourceSchemaVerify(id: $id, input: $input) {
        segmentationConflicts {
          fields
          segmentation { id }
        }
        validatorConflicts {
          fields
          validator { id }
        }
        windowConflicts {
          fields
          window { id }
        }
      }
    }
    """

    result = await client.execute(
        mutation,
        variable_values={
            "id": source_id,
            "input": jtd_schema,
        },
    )

    for conflict_type in [
        "segmentationConflicts",
        "validatorConflicts",
        "windowConflicts",
    ]:
        assert not result["sourceSchemaVerify"][
            conflict_type
        ], f"New schema is invalid for the current setup: {result}"


async def update_bigquery_source_jtd_schema(
    client: Client,
    source_id: str,
) -> None:
    source_query = """
    query Source($id: SourceId!) {
      source(id: $id) {
        credential {
          id
        }

        ... on GcpBigQuerySource {
          config {
            project
            dataset
            table
          }
        }
      }
    }
    """

    source = await client.execute(source_query, variable_values={"id": source_id})
    source_config = source["source"]["config"]

    jtd_schema = await get_jtd_schema(
        client,
        source["source"]["credential"]["id"],
        source_config["project"],
        source_config["dataset"],
        source_config["table"],
    )

    await verify_jtd_schema(client, source_id, jtd_schema)

    update_mutation = """
    mutation SourceUpdate($input: SourceUpdateInput!) {
      sourceUpdate(input: $input) {
        id
      }
    }
    """

    variable_values = {
        "input": {
            "id": source_id,
            "jtdSchema": jtd_schema,
        }
    }

    await client.execute(
        update_mutation,
        variable_values=variable_values,
    )