Add graphql support for taxonomy

.gitignore

@@ -34,3 +34,9 @@ restructure.sh
 .agents
 .claude
 .cursor
+
+# Impala subgraph SDL used by scripts/validate_graphql.py and the offline
+# schema test. Never committed: it carries internal feature-flag/permission
+# names, tenant-specific policy, and staff email addresses from deprecation
+# directives. Keep it local, or point TTD_GRAPHQL_SCHEMA_PATH elsewhere.
+impala.graphql

examples/graphql_example.py

@@ -0,0 +1,112 @@
+"""Example: third-party data taxonomy GraphQL operations via
+ttd-data-python's DataClient.
+
+    TTD_AUTH_TOKEN=... python examples/graphql_example.py
+
+Reads only. The segment upsert mutation is at the bottom, guarded by
+GRAPHQL_EXAMPLE_ALLOW_UPSERT so running this file cannot alter a taxonomy.
+"""
+
+import json
+import os
+
+from ttd_data import DataClient
+
+PROVIDER_ID = os.getenv("GRAPHQL_EXAMPLE_PROVIDER_ID", "eltoro")
+
+token = os.getenv("TTD_AUTH_TOKEN")
+if not token:
+    raise SystemExit("Set TTD_AUTH_TOKEN to a platform token.")
+
+client = DataClient(graphql_http_headers={"TTD-Auth": token})
+
+
+def show(label: str, response: dict) -> None:
+    """Print a response, surfacing GraphQL errors rather than hiding them in
+    a wall of JSON."""
+    print(f"\n{'=' * 60}\n  {label}\n{'=' * 60}")
+    if response.get("errors"):
+        print("GraphQL errors:")
+        print(json.dumps(response["errors"], indent=2))
+        return
+    print(json.dumps(response.get("data"), indent=2))
+
+
+# ---------------------------------------------------------------------------
+# Taxonomy
+# ---------------------------------------------------------------------------
+
+segments = client.graphql.query_segments(provider_id=PROVIDER_ID, first=5)
+show("Segments for provider", segments)
+
+page_info = (
+    (segments.get("data") or {})
+    .get("thirdPartyDataProvider", {})
+    .get("thirdPartyTargetingDataSegments", {})
+    .get("pageInfo", {})
+)
+if page_info.get("hasNextPage"):
+    show(
+        "Segments (next page, via endCursor)",
+        client.graphql.query_segments(
+            provider_id=PROVIDER_ID, first=5, after=page_info["endCursor"]
+        ),
+    )
+
+nodes = (
+    (segments.get("data") or {})
+    .get("thirdPartyDataProvider", {})
+    .get("thirdPartyTargetingDataSegments", {})
+    .get("nodes")
+) or []
+
+if nodes:
+    element_id = nodes[0]["providerElementId"]
+    show(
+        f"Taxonomy approval status for {element_id}",
+        client.graphql.query_segment_taxonomy_status(PROVIDER_ID, element_id),
+    )
+    show(
+        f"Segments filtered to {element_id}",
+        client.graphql.query_segments(
+            provider_id=PROVIDER_ID, provider_element_ids=[element_id]
+        ),
+    )
+
+# ---------------------------------------------------------------------------
+# Escape hatch: anything the typed methods do not cover
+# ---------------------------------------------------------------------------
+
+show(
+    "Arbitrary query via execute()",
+    client.graphql.execute(
+        query="""
+        query ThirdPartyDataProvider($id: ID!) {
+          thirdPartyDataProvider(id: $id) {
+            id
+            name
+          }
+        }
+        """,
+        variables={"id": PROVIDER_ID},
+    ),
+)
+
+# ---------------------------------------------------------------------------
+# Mutation — opt in explicitly; this writes to the provider's taxonomy
+# ---------------------------------------------------------------------------
+
+if os.getenv("GRAPHQL_EXAMPLE_ALLOW_UPSERT") == "1":
+    show(
+        "Upsert segment",
+        client.graphql.upsert_segment(
+            provider_id=PROVIDER_ID,
+            provider_element_id=os.environ["GRAPHQL_EXAMPLE_ELEMENT_ID"],
+            display_name="Example > SDK Test Segment",
+            parent_element_id="ROOT",
+            buyable=True,
+            description="Created by examples/graphql_example.py",
+        ),
+    )
+else:
+    print("\nSkipping upsert (set GRAPHQL_EXAMPLE_ALLOW_UPSERT=1 to run it).")

pyproject.toml

@@ -16,6 +16,7 @@ license = { text = "map[name:Apache License 2.0 shortName:Apache-2.0 url:https:/
 
 [dependency-groups]
 dev = [
+    "graphql-core >=3.2.0,<4.0.0",
     "mypy ==1.15.0",
     "pylint ==3.2.3",
     "pyright ==1.1.398",

scripts/validate_graphql.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python
+"""Validate every GraphQL document the SDK sends against the live schema.
+
+Introspects the Platform API supergraph and validates each document in
+`ttd_data.graphql.QUERY_DOCUMENTS`. Requires a platform credential:
+
+    TTD_AUTH_TOKEN=... python scripts/validate_graphql.py
+
+Introspection needs authentication (the endpoint returns 401 without it), so
+this runs on demand rather than as a pull-request check.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Dict, Tuple
+
+import httpx
+from graphql import build_client_schema, get_introspection_query, parse, print_schema, validate
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+# pylint: disable=wrong-import-position  # sys.path is set up just above
+from ttd_data.graphql import QUERY_DOCUMENTS  # noqa: E402
+from ttd_data.graphql.client import GraphQL  # noqa: E402
+
+HTTP_TIMEOUT_S = 60.0
+
+
+def resolve_auth_headers() -> Tuple[Dict[str, str], str]:
+    """Resolve credentials from the environment. Never returns the token value
+    in the description."""
+    bearer = os.environ.get("BEARER_TOKEN", "").strip()
+    if bearer:
+        value = bearer if bearer.lower().startswith("bearer ") else f"Bearer {bearer}"
+        return {"Authorization": value}, "BEARER_TOKEN"
+    ttd_auth = os.environ.get("TTD_AUTH_TOKEN", "").strip()
+    if ttd_auth:
+        return {"TTD-Auth": ttd_auth}, "TTD_AUTH_TOKEN"
+    return {}, "none"
+
+
+def fetch_schema(url: str, headers: Dict[str, str]):
+    """POST the introspection query and build a client schema from the result."""
+    query = get_introspection_query(descriptions=False, directive_is_repeatable=True)
+    try:
+        response = httpx.post(
+            url,
+            json={"query": query},
+            headers={"Content-Type": "application/json", **headers},
+            timeout=HTTP_TIMEOUT_S,
+        )
+    except httpx.HTTPError as exc:
+        sys.exit(f"Failed to reach {url}: {exc}. Check your network/VPN connection.")
+
+    if response.status_code in (401, 403):
+        sys.exit(f"Auth rejected ({response.status_code}) by {url}. Refresh your token and retry.")
+    if response.status_code != 200:
+        sys.exit(f"Introspection failed: HTTP {response.status_code} from {url}.")
+
+    body = response.json()
+    if body.get("errors"):
+        sys.exit(f"Introspection returned errors: {body['errors']}")
+    schema_root = (body.get("data") or {}).get("__schema")
+    if not schema_root:
+        sys.exit(f"Response from {url} contained no '__schema'.")
+    return build_client_schema({"__schema": schema_root})
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--url", default=GraphQL.DEFAULT_ENDPOINT, help="GraphQL endpoint")
+    parser.add_argument("--write-sdl", type=Path, help="Also write the composed SDL here")
+    args = parser.parse_args()
+
+    headers, source = resolve_auth_headers()
+    if not headers:
+        sys.exit("No credentials found. Set TTD_AUTH_TOKEN or BEARER_TOKEN.")
+
+    print(f"Introspecting {args.url} (auth from {source})...", flush=True)
+    schema = fetch_schema(args.url, headers)
+
+    if args.write_sdl:
+        args.write_sdl.parent.mkdir(parents=True, exist_ok=True)
+        args.write_sdl.write_text(print_schema(schema), encoding="utf-8")
+        print(f"Wrote SDL to {args.write_sdl}")
+
+    failed = 0
+    for name, document in sorted(QUERY_DOCUMENTS.items()):
+        errors = validate(schema, parse(document))
+        if errors:
+            failed += 1
+            print(f"\nFAIL {name}")
+            for error in errors:
+                print(f"  {error.message}")
+                if error.locations:
+                    location = error.locations[0]
+                    print(f"    at line {location.line}, column {location.column}")
+        else:
+            print(f"ok   {name}")
+
+    total = len(QUERY_DOCUMENTS)
+    if failed:
+        print(f"\n{failed}/{total} documents failed validation.")
+        return 1
+    print(f"\nAll {total} documents valid.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/ttd_data/client.py

@@ -12,6 +12,7 @@
 
 from uid2_client import IdentityMapV3Client, IdentityMapV3Input  # type: ignore[import-not-found,import-untyped]
 
+from ttd_data.graphql import GraphQL
 from ttd_data.sdk import BaseDataClient
 from ttd_data.types import BaseModel, OptionalNullable
 from ttd_data.utils import RetryConfig
@@ -59,6 +60,8 @@ class ClientConfig:
     retry_config: OptionalNullable[RetryConfig]
     timeout_ms: Optional[int]
     uid2_config: Optional[UID2Config]
+    graphql_server_url: Optional[str]
+    graphql_http_headers: Optional[Dict[str, str]]
 
 
 class DataClient:
@@ -76,10 +79,14 @@ def __init__(
         self,
         uid2_config: Optional[UID2Config] = None,
         data_client: Optional[BaseDataClient] = None,
+        graphql_server_url: Optional[str] = None,
+        graphql_http_headers: Optional[Dict[str, str]] = None,
         **data_client_kwargs: Any,
     ) -> None:
         self.uid2_config = uid2_config
         self.data_client = data_client or BaseDataClient(**data_client_kwargs)
+        self._graphql_server_url = graphql_server_url
+        self._graphql_http_headers = graphql_http_headers
 
     @property
     def config(self) -> ClientConfig:
@@ -91,6 +98,8 @@ def config(self) -> ClientConfig:
             retry_config=sdk_config.retry_config,
             timeout_ms=sdk_config.timeout_ms,
             uid2_config=self.uid2_config,
+            graphql_server_url=self._graphql_server_url,
+            graphql_http_headers=self._graphql_http_headers,
         )
 
     @classmethod
@@ -296,6 +305,20 @@ def offline_conversion(self) -> "_OfflineConversionProxy":
     def deletion_opt_out(self) -> "_DeletionOptOutProxy":
         return _DeletionOptOutProxy(self)
 
+    @cached_property
+    def graphql(self) -> GraphQL:
+        """Direct GraphQL access to third-party data taxonomy operations.
+        Not part of the UID2 pipeline: these operate on segment metadata,
+        never on user identifiers."""
+        sdk_config = self.data_client.sdk_configuration
+        return GraphQL(
+            server_url=self._graphql_server_url,
+            client=sdk_config.client,
+            http_headers=self._graphql_http_headers,
+            timeout_ms=sdk_config.timeout_ms,
+            debug_logger=sdk_config.debug_logger,
+        )
+
     # ----- Pass-through for any sub-SDK without a UID2 wrapper -----
 
     def __getattr__(self, name: str) -> Any:

src/ttd_data/graphql/__init__.py

@@ -0,0 +1,8 @@
+from ttd_data.graphql.client import GraphQL
+from ttd_data.graphql.taxonomy import QUERY_DOCUMENTS as _TAXONOMY_DOCUMENTS
+
+# Every document the typed methods send, keyed by method name. The schema
+# validator reads this so it can never drift from what callers actually send.
+QUERY_DOCUMENTS = dict(_TAXONOMY_DOCUMENTS)
+
+__all__ = ["GraphQL", "QUERY_DOCUMENTS"]

src/ttd_data/graphql/_operations.py

@@ -0,0 +1,23 @@
+"""Shared base for the typed GraphQL operation mixins."""
+
+from typing import Any, Dict, Mapping, Optional, Protocol
+
+
+class _Execute(Protocol):
+    """Signature of `GraphQL.execute`, as the operation mixins call it."""
+
+    def __call__(
+        self,
+        query: str,
+        variables: Optional[Dict[str, Any]] = None,
+        timeout_ms: Optional[int] = None,
+        http_headers: Optional[Mapping[str, str]] = None,
+    ) -> Dict[str, Any]: ...
+
+
+class GraphQLOperations:
+    """Base for the operation mixins. `execute` is supplied by `GraphQL`; it is
+    declared here so the mixins type-check on their own.
+    """
+
+    execute: _Execute