venv/lib/python3.11/site-packages/wandb/apis/public/integrations.py

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Iterable

from pydantic import ValidationError
from typing_extensions import override
from wandb_gql import gql
from wandb_graphql.language.ast import Document

from wandb.apis.paginator import Paginator

if TYPE_CHECKING:
    from wandb.apis.paginator import _Client
    from wandb.automations import Integration, SlackIntegration, WebhookIntegration
    from wandb.automations._generated import (
        GenericWebhookIntegrationConnectionFields,
        IntegrationConnectionFields,
        SlackIntegrationConnectionFields,
    )


class Integrations(Paginator["Integration"]):
    last_response: IntegrationConnectionFields | None
    _query: Document

    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
        from wandb.automations._generated import INTEGRATIONS_BY_ENTITY_GQL

        super().__init__(client, variables, per_page=per_page)
        # All integrations for entity
        self._query = gql(INTEGRATIONS_BY_ENTITY_GQL)

    @property
    def more(self) -> bool:
        """Whether there are more Integrations to fetch."""
        if self.last_response is None:
            return True
        return self.last_response.page_info.has_next_page

    @property
    def cursor(self) -> str | None:
        """The start cursor to use for the next page."""
        if self.last_response is None:
            return None
        return self.last_response.page_info.end_cursor

    @override
    def _update_response(self) -> None:
        """Fetch and parse the response data for the current page."""
        from wandb.automations._generated import IntegrationConnectionFields

        data: dict[str, Any] = self.client.execute(
            self._query, variable_values=self.variables
        )
        try:
            page_data = data["entity"]["integrations"]
            self.last_response = IntegrationConnectionFields.model_validate(page_data)
        except (LookupError, AttributeError, ValidationError) as e:
            raise ValueError("Unexpected response data") from e

    def convert_objects(self) -> Iterable[Integration]:
        """Parse the page data into a list of integrations."""
        from wandb.automations.integrations import _IntegrationEdge

        page = self.last_response
        return [_IntegrationEdge.model_validate(edge).node for edge in page.edges]


class WebhookIntegrations(Paginator["WebhookIntegration"]):
    last_response: GenericWebhookIntegrationConnectionFields | None
    _query: Document

    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
        from wandb.automations._generated import (
            GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL,
        )

        super().__init__(client, variables, per_page=per_page)
        # Webhook integrations for entity
        self._query = gql(GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL)

    @property
    def more(self) -> bool:
        """Whether there are more webhook integrations to fetch."""
        if self.last_response is None:
            return True
        return self.last_response.page_info.has_next_page

    @property
    def cursor(self) -> str | None:
        """The start cursor to use for the next page."""
        if self.last_response is None:
            return None
        return self.last_response.page_info.end_cursor

    @override
    def _update_response(self) -> None:
        """Fetch and parse the response data for the current page."""
        from wandb.automations._generated import (
            GenericWebhookIntegrationConnectionFields,
        )

        data: dict[str, Any] = self.client.execute(
            self._query, variable_values=self.variables
        )
        try:
            page_data = data["entity"]["integrations"]
            self.last_response = (
                GenericWebhookIntegrationConnectionFields.model_validate(page_data)
            )
        except (LookupError, AttributeError, ValidationError) as e:
            raise ValueError("Unexpected response data") from e

    def convert_objects(self) -> Iterable[WebhookIntegration]:
        """Parse the page data into a list of webhook integrations."""
        from wandb.automations import WebhookIntegration

        typename = "GenericWebhookIntegration"
        return [
            # Filter on typename__ needed because the GQL response still
            # includes all integration types
            WebhookIntegration.model_validate(node)
            for edge in self.last_response.edges
            if (node := edge.node) and (node.typename__ == typename)
        ]


class SlackIntegrations(Paginator["SlackIntegration"]):
    last_response: SlackIntegrationConnectionFields | None
    _query: Document

    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
        from wandb.automations._generated import SLACK_INTEGRATIONS_BY_ENTITY_GQL

        super().__init__(client, variables, per_page=per_page)
        # Slack integrations for entity
        self._query = gql(SLACK_INTEGRATIONS_BY_ENTITY_GQL)

    @property
    def more(self) -> bool:
        """Whether there are more Slack integrations to fetch."""
        if self.last_response is None:
            return True
        return self.last_response.page_info.has_next_page

    @property
    def cursor(self) -> str | None:
        """The start cursor to use for the next page."""
        if self.last_response is None:
            return None
        return self.last_response.page_info.end_cursor

    @override
    def _update_response(self) -> None:
        """Fetch and parse the response data for the current page."""
        from wandb.automations._generated import SlackIntegrationConnectionFields

        data: dict[str, Any] = self.client.execute(
            self._query, variable_values=self.variables
        )
        try:
            page_data = data["entity"]["integrations"]
            self.last_response = SlackIntegrationConnectionFields.model_validate(
                page_data
            )
        except (LookupError, AttributeError, ValidationError) as e:
            raise ValueError("Unexpected response data") from e

    def convert_objects(self) -> Iterable[SlackIntegration]:
        """Parse the page data into a list of Slack integrations."""
        from wandb.automations import SlackIntegration

        typename = "SlackIntegration"
        return [
            # Filter on typename__ needed because the GQL response still
            # includes all integration types
            SlackIntegration.model_validate(node)
            for edge in self.last_response.edges
            if (node := edge.node) and (node.typename__ == typename)
        ]