Owners ownership

"""
Contains helper classes for authorization to the SYNQ server.
"""

import requests
import time
import grpc


class TokenAuth(grpc.AuthMetadataPlugin):
    """AuthMetadataPlugin which adds the access token to the outgoing context metadata."""

    def __init__(self, token_source):
        self._token_source = token_source

    def __call__(self, context, callback):
        try:
            token = self._token_source.get_token()
            callback([("authorization", f"Bearer {token}")], None)
        except Exception as e:
            callback(None, e)


class TokenSource:
    """Token source which maintains the access token and refreshes it when it is expired."""

    def __init__(self, client_id, client_secret, api_endpoint):
        self.api_endpoint = api_endpoint
        self.token_url = f"https://{self.api_endpoint}/oauth2/token"
        self.client_id = client_id
        self.client_secret = client_secret
        self.token = self.obtain_token()

    def obtain_token(self):
        resp = requests.post(
            self.token_url,
            data={
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "grant_type": "client_credentials",
            },
        )
        resp.raise_for_status()
        self.token = resp.json()
        self.expires_at = time.time() + self.token["expires_in"]
        return self.token

    def get_token(self) -> str:
        if time.time() > self.expires_at:
            self.obtain_token()
        return self.token["access_token"]

"""
Data-products maintenance lifecycle using the Coalesce Quality
DataproductsService (v2):

    upsert -> set-definition -> add/remove parts -> list-members -> batch-get ->
    partial update (with etag) -> delete

A data product is a named, owned grouping of assets with a membership
definition. The id is a caller-supplied UUID, so every write is idempotent —
re-running this example converges instead of duplicating.

The membership definition is authored here in ResolverQL, the compact text
query language. The server compiles and stores it canonically and echoes it
back on reads as `rendered_resolver_ql`.

Prerequisites:
- SYNQ_CLIENT_ID and SYNQ_CLIENT_SECRET (scopes: Read/Edit Data Products).
- Optionally API_ENDPOINT (defaults to developer.synq.io; US: api.us.synq.io).
"""

import os
import sys
import uuid

import grpc
from dotenv import load_dotenv

sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))

from auth import TokenSource, TokenAuth
from synq.dataproducts.v2 import (
    dataproducts_service_pb2 as svc,
    dataproducts_service_pb2_grpc as svc_grpc,
    dataproduct_pb2 as dp,
    dataproduct_definition_pb2 as dp_def,
)
from synq.v1 import pagination_pb2

load_dotenv()

CLIENT_ID = os.getenv("SYNQ_CLIENT_ID")
CLIENT_SECRET = os.getenv("SYNQ_CLIENT_SECRET")
API_ENDPOINT = os.getenv("API_ENDPOINT", "developer.synq.io")

if not CLIENT_ID or not CLIENT_SECRET:
    raise SystemExit("SYNQ_CLIENT_ID and SYNQ_CLIENT_SECRET must be set (scopes: Read/Edit Data Products)")


def print_definition(product):
    parts = product.definition.parts
    print(f"definition has {len(parts)} part(s), etag={product.etag}")
    for p in parts:
        if p.HasField("query"):
            print(f"  - part {p.id}: {p.query.rendered_resolver_ql}")
        elif p.entity_id:
            print(f"  - part {p.id}: entity_id={p.entity_id}")
    print()


def main():
    token_source = TokenSource(CLIENT_ID, CLIENT_SECRET, API_ENDPOINT)
    creds = grpc.composite_channel_credentials(
        grpc.ssl_channel_credentials(),
        grpc.metadata_call_credentials(TokenAuth(token_source)),
    )
    with grpc.secure_channel(
        f"{API_ENDPOINT}:443", creds, options=(("grpc.default_authority", API_ENDPOINT),)
    ) as channel:
        grpc.channel_ready_future(channel).result(timeout=10)
        print(f"Connected to {API_ENDPOINT}\n")
        stub = svc_grpc.DataproductsServiceStub(channel)

        # The caller owns the id. A fresh UUID creates a new product; re-using a
        # known id updates that product in place.
        product_id = str(uuid.uuid4())

        # --- Step 1: Create ---
        # Only `title` is required on create. `folder` and `priority` are optional.
        print("=== Step 1: Create data product ===")
        created = stub.Upsert(svc.UpsertRequest(
            id=product_id,
            title="API Example — Orders",
            description="Created by the owners_ownership Python example.",
            folder="API Examples",
            priority=dp.Dataproduct.PRIORITY_P1,
        )).dataproduct
        print(f"Created id={created.id}")
        # entity_id is server-derived (dataproduct-<uuid>) — the value other APIs
        # (lineage, entities, alerts) accept as a reference to this product.
        print(f"entity_id={created.entity_id} (use this to reference the product elsewhere)")
        print(f"etag={created.etag} priority={dp.Dataproduct.Priority.Name(created.priority)}\n")

        # --- Step 2: Replace the whole definition (SetDefinition) ---
        # A definition is a list of parts OR'd together. Here a single ResolverQL
        # query part. Replace the placeholder query with one matching your workspace.
        print("=== Step 2: Set definition (ResolverQL) ===")
        set_resp = stub.SetDefinition(svc.SetDefinitionRequest(
            id=product_id,
            etag=created.etag,
            definition=dp_def.DataproductDefinition(parts=[
                dp_def.DataproductDefinition.Part(
                    id=str(uuid.uuid4()),
                    query=dp_def.DataproductQuery(resolver_ql='with_type("table", filter=with_name("orders"))'),
                )
            ]),
        )).dataproduct
        print_definition(set_resp)
        product = set_resp

        # --- Step 3: Add a second part (UpsertDefinitionPart) ---
        # Parts have caller-supplied ids so an upsert is idempotent. Definition
        # writes bump the etag, so pass the latest one.
        print("=== Step 3: Add a second part ===")
        part_id = str(uuid.uuid4())
        product = stub.UpsertDefinitionPart(svc.UpsertDefinitionPartRequest(
            id=product_id,
            etag=product.etag,
            part=dp_def.DataproductDefinition.Part(
                id=part_id,
                query=dp_def.DataproductQuery(resolver_ql='with_type("view", filter=with_name("orders"))'),
            ),
        )).dataproduct
        print_definition(product)

        # --- Step 4: Remove that part (RemoveDefinitionPart) ---
        print("=== Step 4: Remove the part ===")
        product = stub.RemoveDefinitionPart(svc.RemoveDefinitionPartRequest(
            id=product_id, part_id=part_id, etag=product.etag,
        )).dataproduct
        print_definition(product)

        # --- Step 5: List resolved members ---
        # The definition resolves to concrete assets (opaque entity ids). Empty is
        # normal if the placeholder query matches nothing in your workspace.
        print("=== Step 5: List members ===")
        members = stub.ListMembers(svc.ListMembersRequest(id=product_id, pagination=pagination_pb2.Pagination()))
        print(f"Resolved {len(members.entity_ids)} member(s)")
        for e in members.entity_ids[:5]:
            print(f"  - {e}")
        if len(members.entity_ids) > 5:
            print("  …")
        print()

        # --- Step 6: BatchGet (exclude the definition for a lighter read) ---
        print("=== Step 6: BatchGet ===")
        got = stub.BatchGet(svc.BatchGetRequest(ids=[product_id], exclude_definition=True)).dataproducts[product_id]
        print(f'title="{got.title}" folder="{got.folder}" '
              f"priority={dp.Dataproduct.Priority.Name(got.priority)} "
              f"(definition excluded: {len(got.definition.parts)} parts)\n")

        # --- Step 7: Partial update with optimistic concurrency ---
        # Upsert only the fields you set; omitted fields are left unchanged. Pass
        # the etag you last read so a concurrent edit is not silently overwritten.
        print("=== Step 7: Partial update (title only) with etag ===")
        stale_etag = product.etag
        product = stub.Upsert(svc.UpsertRequest(
            id=product_id,
            title="API Example — Orders (P2)",
            priority=dp.Dataproduct.PRIORITY_P2,
            etag=product.etag,
        )).dataproduct
        print(f'Updated title="{product.title}" '
              f"priority={dp.Dataproduct.Priority.Name(product.priority)} new etag={product.etag}")

        # Re-using the now-stale etag is rejected — the concurrency guard.
        try:
            stub.Upsert(svc.UpsertRequest(id=product_id, title="should be rejected", etag=stale_etag))
            print("Unexpected: stale etag accepted\n")
        except grpc.RpcError as e:
            ok = e.code() == grpc.StatusCode.ABORTED
            print(f"Stale etag rejected with {e.code().name} ({'expected' if ok else 'unexpected'})\n")

        # --- Step 8: Delete (purge to release the id) ---
        print("=== Step 8: Delete ===")
        stub.Delete(svc.DeleteRequest(id=product_id, purge=True, etag=product.etag))
        after = stub.BatchGet(svc.BatchGetRequest(ids=[product_id]))
        print(f"Deleted {product_id}; BatchGet after delete returned {len(after.dataproducts)} product(s)")

        print("\nDone: data-product maintenance lifecycle exercised end to end.")


if __name__ == "__main__":
    main()

"""
Owners & ownership maintenance lifecycle — "alert routing as code" — using the
Coalesce Quality OwnersService (v1):

    create owner (+contacts) -> assign assets (ownership + alert config) ->
    list -> partial update (with etag) -> delete

An owner is a named responsible party with notification channels (contacts). An
ownership assigns a set of assets to an owner and configures the alerts routed
to it. Owner is the resource; ownership is its sub-resource — deleting an owner
deletes its ownerships.

Both ids are caller-supplied UUIDs, so every write is idempotent. Together they
replace clicking owners and alert rules together in the UI: define the
responsible party, its channels, what it owns, and how it should be alerted —
all as code.

To make the "assign a data product" path concrete, the example first creates a
throwaway data product to own, and deletes it at the end.

Prerequisites:
- SYNQ_CLIENT_ID and SYNQ_CLIENT_SECRET (scopes: Read/Edit Owners, Read/Edit
  Ownership, Read/Edit Data Products).
- Optionally API_ENDPOINT (defaults to developer.synq.io; US: api.us.synq.io).
"""

import os
import sys
import uuid

import grpc
from dotenv import load_dotenv

sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))

from auth import TokenSource, TokenAuth
from synq.dataproducts.v2 import (
    dataproducts_service_pb2 as dp_svc,
    dataproducts_service_pb2_grpc as dp_grpc,
    dataproduct_pb2 as dp,
    dataproduct_definition_pb2 as dp_def,
)
from synq.owners.v1 import (
    owners_service_pb2 as svc,
    owners_service_pb2_grpc as svc_grpc,
    ownership_pb2 as own,
    contact_pb2 as contact,
)
from synq.alerts.v1 import alerts_pb2
from synq.v1 import severity_pb2, pagination_pb2

load_dotenv()

CLIENT_ID = os.getenv("SYNQ_CLIENT_ID")
CLIENT_SECRET = os.getenv("SYNQ_CLIENT_SECRET")
API_ENDPOINT = os.getenv("API_ENDPOINT", "developer.synq.io")

if not CLIENT_ID or not CLIENT_SECRET:
    raise SystemExit("SYNQ_CLIENT_ID and SYNQ_CLIENT_SECRET must be set (scopes: Read/Edit Owners, Ownership, Data Products)")


def print_ownership(label, o):
    print(f"{label} id={o.id} etag={o.etag}")
    if o.selection.HasField("dataproduct_id"):
        print(f"  selection: data product {o.selection.dataproduct_id}")
    elif o.selection.HasField("query"):
        print(f'  selection: query "{o.selection.query.name}" -> {o.selection.query.rendered_resolver_ql}')
    a = o.alert
    sevs = [severity_pb2.Severity.Name(s) for s in a.severities]
    print(f"  alert: severities={sevs} notify_upstream={a.notify_upstream} disabled={a.is_disabled}\n")


def main():
    token_source = TokenSource(CLIENT_ID, CLIENT_SECRET, API_ENDPOINT)
    creds = grpc.composite_channel_credentials(
        grpc.ssl_channel_credentials(),
        grpc.metadata_call_credentials(TokenAuth(token_source)),
    )
    with grpc.secure_channel(
        f"{API_ENDPOINT}:443", creds, options=(("grpc.default_authority", API_ENDPOINT),)
    ) as channel:
        grpc.channel_ready_future(channel).result(timeout=10)
        print(f"Connected to {API_ENDPOINT}\n")
        owners = svc_grpc.OwnersServiceStub(channel)
        products = dp_grpc.DataproductsServiceStub(channel)

        # --- Setup: a data product to own ---
        print("=== Setup: create a data product to own ===")
        product_id = str(uuid.uuid4())
        products.Upsert(dp_svc.UpsertRequest(
            id=product_id,
            title="API Example — Owned Product",
            folder="API Examples",
            priority=dp.Dataproduct.PRIORITY_P1,
            definition=dp_def.DataproductDefinition(parts=[
                dp_def.DataproductDefinition.Part(
                    id=str(uuid.uuid4()),
                    query=dp_def.DataproductQuery(resolver_ql='with_type("table", filter=with_name("orders"))'),
                )
            ]),
        ))
        print(f"Created product {product_id}\n")

        # --- Step 1: Create an owner with notification channels ---
        # Contacts are the channels a fired alert is delivered to. Replace the
        # placeholders with real channels/addresses for your workspace.
        print("=== Step 1: Create owner (with contacts) ===")
        owner_id = str(uuid.uuid4())
        owner = owners.UpsertOwner(svc.UpsertOwnerRequest(
            id=owner_id,
            title="API Example — Data Platform",
            contacts=svc.ContactList(contacts=[
                contact.Contact(slack=contact.SlackChannelContact(channel="#data-alerts")),
                contact.Contact(email=contact.EmailContact(recipient_emails=["data-team@example.com"])),
            ]),
        )).owner
        print(f"Created owner id={owner.id}")
        # entity_id is server-derived (owner-<uuid>) — the value the alerts API
        # and other surfaces accept as an owner reference.
        print(f"entity_id={owner.entity_id} contacts={len(owner.contacts)} etag={owner.etag}\n")

        # --- Step 2: Assign the data product to the owner (ownership #1) ---
        # The ownership's AlertConfig is the "routing as code": which severities
        # fire, whether upstream issues count, and the repeat strategy.
        print("=== Step 2: Ownership #1 — own the data product, with alert routing ===")
        ownership_product_id = str(uuid.uuid4())
        up1 = owners.UpsertOwnership(svc.UpsertOwnershipRequest(
            owner_id=owner_id,
            id=ownership_product_id,
            selection=own.OwnershipSelection(dataproduct_id=product_id),
            alert=own.AlertConfig(
                severities=[severity_pb2.SEVERITY_ERROR, severity_pb2.SEVERITY_FATAL],
                notify_upstream=True,
                ongoing=alerts_pb2.OngoingAlertsStrategy(disabled=alerts_pb2.OngoingAlertsStrategy.Disabled()),
            ),
        )).ownership
        print_ownership("ownership #1", up1)

        # --- Step 3: A second ownership selecting assets by query (ownership #2) ---
        # Unlike a data-product definition (a leaf), an ownership query MAY
        # reference data products and domains.
        print("=== Step 3: Ownership #2 — select by ResolverQL, weekly digest ===")
        ownership_query_id = str(uuid.uuid4())
        up2 = owners.UpsertOwnership(svc.UpsertOwnershipRequest(
            owner_id=owner_id,
            id=ownership_query_id,
            selection=own.OwnershipSelection(
                query=own.OwnershipQuery(name="Critical models", resolver_ql='with_type("model", filter=with_name("revenue"))'),
            ),
            alert=own.AlertConfig(
                severities=[severity_pb2.SEVERITY_FATAL],
                ongoing=alerts_pb2.OngoingAlertsStrategy(
                    schedule=alerts_pb2.OngoingAlertsStrategy.Schedule(cron="0 9 * * MON"),
                ),
            ),
        )).ownership
        print_ownership("ownership #2", up2)

        # --- Step 4: List the owner's ownerships ---
        print("=== Step 4: List ownerships ===")
        listed = owners.ListOwnerships(svc.ListOwnershipsRequest(owner_id=owner_id, pagination=pagination_pb2.Pagination()))
        print(f"Owner has {len(listed.ownerships)} ownership(s)\n")

        # --- Step 5: Partial owner update with optimistic concurrency ---
        # Contacts are replace-semantics behind a presence wrapper: omit `contacts`
        # to leave them unchanged, or send the full desired set to replace them.
        print("=== Step 5: Replace contacts (add MS Teams) with etag ===")
        stale_etag = owner.etag
        owner = owners.UpsertOwner(svc.UpsertOwnerRequest(
            id=owner_id,
            etag=owner.etag,
            contacts=svc.ContactList(contacts=[
                contact.Contact(slack=contact.SlackChannelContact(channel="#data-alerts")),
                contact.Contact(email=contact.EmailContact(recipient_emails=["data-team@example.com"])),
                contact.Contact(ms_teams=contact.MsTeamsContact(channel_id="19:example-teams-channel-id@thread.tacv2")),
            ]),
        )).owner
        print(f"Updated owner now has {len(owner.contacts)} contacts, new etag={owner.etag}")

        try:
            owners.UpsertOwner(svc.UpsertOwnerRequest(id=owner_id, title="should be rejected", etag=stale_etag))
            print("Unexpected: stale etag accepted\n")
        except grpc.RpcError as e:
            ok = e.code() == grpc.StatusCode.ABORTED
            print(f"Stale etag rejected with {e.code().name} ({'expected' if ok else 'unexpected'})\n")

        # --- How this ties to the alerts API ---
        print("=== Reference: how alerts point back here ===")
        print(f"owner entity_id : {owner.entity_id}")
        print(f"ownership ids   : {ownership_product_id}, {ownership_query_id}\n")

        # --- Cleanup ---
        print("=== Cleanup ===")
        owners.DeleteOwnership(svc.DeleteOwnershipRequest(id=ownership_query_id))
        owners.DeleteOwner(svc.DeleteOwnerRequest(id=owner_id, purge=True))
        products.Delete(dp_svc.DeleteRequest(id=product_id, purge=True))
        print(f"Deleted ownerships, owner {owner_id}, and product {product_id}")

        print("\nDone: owners & ownership lifecycle exercised end to end.")


if __name__ == "__main__":
    main()