collibra-connector 1.0.18__py3-none-any.whl → 1.1.0__py3-none-any.whl

collibra_connector/lineage.py

@@ -0,0 +1,716 @@
+"""
+Lineage Builder - Declarative data lineage creation.
+
+This module provides a fluent API for building and committing
+technical data lineage in Collibra, including assets and relations.
+
+Example:
+    >>> from collibra_connector import CollibraConnector
+    >>> from collibra_connector.lineage import LineageBuilder, LineageNode
+    >>>
+    >>> conn = CollibraConnector(...)
+    >>> builder = LineageBuilder(conn)
+    >>>
+    >>> # Define nodes
+    >>> s3_table = LineageNode("s3://bucket/raw/customers", "Table")
+    >>> glue_job = LineageNode("etl_transform_customers", "Data Pipeline")
+    >>> redshift_table = LineageNode("warehouse.customers", "Table")
+    >>>
+    >>> # Build lineage
+    >>> builder.add_edge(s3_table, glue_job, "is source for")
+    >>> builder.add_edge(glue_job, redshift_table, "is target for")
+    >>>
+    >>> # Commit to Collibra
+    >>> result = builder.commit(domain_id="lineage-domain-uuid")
+    >>> print(f"Created {result.assets_created} assets, {result.relations_created} relations")
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Tuple, Union, TYPE_CHECKING
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    from .connector import CollibraConnector
+
+
+class LineageDirection(str, Enum):
+    """Direction of data flow in lineage."""
+    UPSTREAM = "upstream"      # Data flows TO this node
+    DOWNSTREAM = "downstream"  # Data flows FROM this node
+    BIDIRECTIONAL = "bidirectional"
+
+
+class LineageRelationType(str, Enum):
+    """Common lineage relation types."""
+    SOURCE_FOR = "is source for"
+    TARGET_FOR = "is target for"
+    TRANSFORMS = "transforms"
+    DERIVED_FROM = "is derived from"
+    CONTAINS = "contains"
+    PART_OF = "is part of"
+    USES = "uses"
+    PRODUCES = "produces"
+
+
+@dataclass
+class LineageNode:
+    """
+    Represents a node in the lineage graph.
+
+    A node can represent any data asset: a table, a file, an ETL job,
+    a dashboard, etc. Nodes can either reference existing Collibra assets
+    by ID or define new assets to be created.
+
+    Attributes:
+        name: The name of the node/asset.
+        asset_type: The asset type name (e.g., "Table", "Data Pipeline").
+        asset_id: Optional existing asset ID (if referencing existing asset).
+        display_name: Optional display name.
+        description: Optional description.
+        attributes: Optional dict of attribute name -> value.
+        metadata: Optional dict of additional metadata.
+
+    Example:
+        >>> # Reference existing asset
+        >>> existing = LineageNode.from_id("existing-asset-uuid")
+        >>>
+        >>> # Define new asset to create
+        >>> new_table = LineageNode(
+        ...     name="raw.customers",
+        ...     asset_type="Table",
+        ...     description="Raw customer data from CRM",
+        ...     attributes={"Data Source": "Salesforce"}
+        ... )
+    """
+    name: str
+    asset_type: str = "Data Asset"
+    asset_id: Optional[str] = None
+    display_name: Optional[str] = None
+    description: Optional[str] = None
+    attributes: Dict[str, Any] = field(default_factory=dict)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    _internal_id: str = field(default_factory=lambda: str(uuid4()))
+
+    @classmethod
+    def from_id(cls, asset_id: str, name: str = "") -> "LineageNode":
+        """Create a node referencing an existing Collibra asset."""
+        return cls(name=name, asset_id=asset_id)
+
+    @classmethod
+    def table(
+        cls,
+        name: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+        **kwargs: Any
+    ) -> "LineageNode":
+        """Create a Table node."""
+        full_name = name
+        if schema:
+            full_name = f"{schema}.{name}"
+        if database:
+            full_name = f"{database}.{full_name}"
+        return cls(name=full_name, asset_type="Table", **kwargs)
+
+    @classmethod
+    def column(cls, name: str, table: Optional[str] = None, **kwargs: Any) -> "LineageNode":
+        """Create a Column node."""
+        full_name = f"{table}.{name}" if table else name
+        return cls(name=full_name, asset_type="Column", **kwargs)
+
+    @classmethod
+    def pipeline(cls, name: str, **kwargs: Any) -> "LineageNode":
+        """Create a Data Pipeline node."""
+        return cls(name=name, asset_type="Data Pipeline", **kwargs)
+
+    @classmethod
+    def report(cls, name: str, **kwargs: Any) -> "LineageNode":
+        """Create a Report node."""
+        return cls(name=name, asset_type="Report", **kwargs)
+
+    @classmethod
+    def dashboard(cls, name: str, **kwargs: Any) -> "LineageNode":
+        """Create a Dashboard node."""
+        return cls(name=name, asset_type="Dashboard", **kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self._internal_id)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, LineageNode):
+            return self._internal_id == other._internal_id
+        return False
+
+
+@dataclass
+class LineageEdge:
+    """
+    Represents an edge (relation) in the lineage graph.
+
+    An edge connects two nodes with a specific relation type,
+    representing data flow or dependency.
+
+    Attributes:
+        source: The source node.
+        target: The target node.
+        relation_type: The type of relation.
+        relation_type_id: Optional specific relation type UUID.
+        metadata: Optional additional metadata for the edge.
+    """
+    source: LineageNode
+    target: LineageNode
+    relation_type: str = "is source for"
+    relation_type_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class LineageCommitResult:
+    """
+    Result of committing lineage to Collibra.
+
+    Attributes:
+        success: Whether the commit was successful.
+        assets_created: Number of new assets created.
+        assets_updated: Number of existing assets updated.
+        relations_created: Number of relations created.
+        assets: Dict mapping node internal IDs to created asset IDs.
+        relations: List of created relation IDs.
+        errors: List of error messages.
+    """
+    success: bool = True
+    assets_created: int = 0
+    assets_updated: int = 0
+    relations_created: int = 0
+    assets: Dict[str, str] = field(default_factory=dict)
+    relations: List[str] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+
+
+class LineageBuilder:
+    """
+    Fluent builder for creating data lineage in Collibra.
+
+    This class provides a declarative API for defining lineage graphs
+    and committing them to Collibra. It handles asset creation,
+    relation creation, and error handling automatically.
+
+    Example:
+        >>> builder = LineageBuilder(connector)
+        >>>
+        >>> # Define the lineage
+        >>> source = LineageNode.table("raw.orders", database="s3")
+        >>> transform = LineageNode.pipeline("transform_orders")
+        >>> target = LineageNode.table("orders", schema="warehouse")
+        >>>
+        >>> builder.add_edge(source, transform, "is source for")
+        >>> builder.add_edge(transform, target, "is source for")
+        >>>
+        >>> # Or use method chaining
+        >>> builder.source(source).through(transform).to(target)
+        >>>
+        >>> # Commit
+        >>> result = builder.commit(domain_id="lineage-domain-uuid")
+
+    Advanced Example with multiple paths:
+        >>> # Multiple sources to one target
+        >>> builder.add_edges([
+        ...     (table_a, etl_job, "is source for"),
+        ...     (table_b, etl_job, "is source for"),
+        ...     (table_c, etl_job, "is source for"),
+        ...     (etl_job, output_table, "is source for"),
+        ... ])
+    """
+
+    def __init__(
+        self,
+        connector: "CollibraConnector",
+        default_relation_type: str = "is source for"
+    ) -> None:
+        """
+        Initialize the LineageBuilder.
+
+        Args:
+            connector: The CollibraConnector instance.
+            default_relation_type: Default relation type for edges.
+        """
+        self.connector = connector
+        self.default_relation_type = default_relation_type
+        self._nodes: Dict[str, LineageNode] = {}
+        self._edges: List[LineageEdge] = []
+        self._current_source: Optional[LineageNode] = None
+        self._type_id_cache: Dict[str, str] = {}
+        self._relation_type_cache: Dict[str, str] = {}
+
+    def add_node(self, node: LineageNode) -> "LineageBuilder":
+        """
+        Add a node to the lineage graph.
+
+        Args:
+            node: The LineageNode to add.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._nodes[node._internal_id] = node
+        return self
+
+    def add_edge(
+        self,
+        source: LineageNode,
+        target: LineageNode,
+        relation_type: Optional[str] = None,
+        relation_type_id: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        Add an edge (relation) between two nodes.
+
+        Args:
+            source: The source node.
+            target: The target node.
+            relation_type: Relation type name (uses default if not specified).
+            relation_type_id: Optional specific relation type UUID.
+
+        Returns:
+            Self for method chaining.
+        """
+        # Auto-add nodes
+        self.add_node(source)
+        self.add_node(target)
+
+        edge = LineageEdge(
+            source=source,
+            target=target,
+            relation_type=relation_type or self.default_relation_type,
+            relation_type_id=relation_type_id
+        )
+        self._edges.append(edge)
+        return self
+
+    def add_edges(
+        self,
+        edges: List[Tuple[LineageNode, LineageNode, str]]
+    ) -> "LineageBuilder":
+        """
+        Add multiple edges at once.
+
+        Args:
+            edges: List of (source, target, relation_type) tuples.
+
+        Returns:
+            Self for method chaining.
+        """
+        for source, target, relation_type in edges:
+            self.add_edge(source, target, relation_type)
+        return self
+
+    def source(self, node: LineageNode) -> "LineageBuilder":
+        """
+        Set the current source node for fluent API.
+
+        Args:
+            node: The source node.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> builder.source(table_a).to(table_b)
+        """
+        self.add_node(node)
+        self._current_source = node
+        return self
+
+    def through(
+        self,
+        node: LineageNode,
+        relation_type: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        Add an intermediate node (like an ETL job).
+
+        Args:
+            node: The intermediate node.
+            relation_type: Relation type from source to this node.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> builder.source(raw_table).through(etl_job).to(warehouse_table)
+        """
+        if self._current_source is None:
+            raise ValueError("Must call source() before through()")
+
+        self.add_edge(self._current_source, node, relation_type)
+        self._current_source = node
+        return self
+
+    def to(
+        self,
+        node: LineageNode,
+        relation_type: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        Add the target node and create edge from current source.
+
+        Args:
+            node: The target node.
+            relation_type: Relation type to target.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> builder.source(table_a).to(table_b)
+        """
+        if self._current_source is None:
+            raise ValueError("Must call source() before to()")
+
+        self.add_edge(self._current_source, node, relation_type)
+        return self
+
+    def chain(
+        self,
+        *nodes: LineageNode,
+        relation_type: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        Create a chain of nodes with edges between consecutive pairs.
+
+        Args:
+            *nodes: Nodes to chain together.
+            relation_type: Relation type for all edges.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> builder.chain(source, transform1, transform2, target)
+        """
+        for i in range(len(nodes) - 1):
+            self.add_edge(nodes[i], nodes[i + 1], relation_type)
+        return self
+
+    def fan_in(
+        self,
+        sources: List[LineageNode],
+        target: LineageNode,
+        relation_type: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        Multiple sources feeding into one target.
+
+        Args:
+            sources: List of source nodes.
+            target: The target node.
+            relation_type: Relation type for all edges.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> # Multiple tables feeding into one ETL job
+            >>> builder.fan_in([table_a, table_b, table_c], etl_job)
+        """
+        for source in sources:
+            self.add_edge(source, target, relation_type)
+        return self
+
+    def fan_out(
+        self,
+        source: LineageNode,
+        targets: List[LineageNode],
+        relation_type: Optional[str] = None
+    ) -> "LineageBuilder":
+        """
+        One source feeding into multiple targets.
+
+        Args:
+            source: The source node.
+            targets: List of target nodes.
+            relation_type: Relation type for all edges.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            >>> # One ETL job producing multiple tables
+            >>> builder.fan_out(etl_job, [table_a, table_b, table_c])
+        """
+        for target in targets:
+            self.add_edge(source, target, relation_type)
+        return self
+
+    def _resolve_asset_type_id(self, type_name: str) -> Optional[str]:
+        """Get asset type ID by name, with caching."""
+        if type_name in self._type_id_cache:
+            return self._type_id_cache[type_name]
+
+        try:
+            # Try to find the asset type
+            result = self.connector.metadata.get_asset_types(name=type_name, limit=1)
+            types = result.get("results", [])
+            if types:
+                type_id = types[0].get("id")
+                self._type_id_cache[type_name] = type_id
+                return type_id
+        except Exception:
+            pass
+
+        return None
+
+    def _resolve_relation_type_id(self, role: str) -> Optional[str]:
+        """Get relation type ID by role name, with caching."""
+        if role in self._relation_type_cache:
+            return self._relation_type_cache[role]
+
+        try:
+            # Search for relation type by role
+            result = self.connector.metadata.get_relation_types(role=role, limit=1)
+            types = result.get("results", [])
+            if types:
+                type_id = types[0].get("id")
+                self._relation_type_cache[role] = type_id
+                return type_id
+        except Exception:
+            pass
+
+        return None
+
+    def commit(
+        self,
+        domain_id: str,
+        status_id: Optional[str] = None,
+        dry_run: bool = False,
+        create_missing_types: bool = False
+    ) -> LineageCommitResult:
+        """
+        Commit the lineage graph to Collibra.
+
+        This method:
+        1. Creates any new assets defined in nodes
+        2. Creates relations between assets
+        3. Sets attributes on new assets
+
+        Args:
+            domain_id: The domain to create assets in.
+            status_id: Optional status for new assets.
+            dry_run: If True, validate but don't create anything.
+            create_missing_types: If True, log warnings for missing types.
+
+        Returns:
+            LineageCommitResult with success status and created IDs.
+
+        Example:
+            >>> result = builder.commit(domain_id="lineage-domain")
+            >>> if result.success:
+            ...     print(f"Created {result.assets_created} assets")
+            ...     print(f"Created {result.relations_created} relations")
+            >>> else:
+            ...     for error in result.errors:
+            ...         print(f"Error: {error}")
+        """
+        result = LineageCommitResult()
+
+        if dry_run:
+            result.assets_created = len([n for n in self._nodes.values() if not n.asset_id])
+            result.relations_created = len(self._edges)
+            return result
+
+        # Phase 1: Create or resolve all assets
+        for internal_id, node in self._nodes.items():
+            try:
+                if node.asset_id:
+                    # Existing asset - just map the ID
+                    result.assets[internal_id] = node.asset_id
+                else:
+                    # New asset - create it
+                    type_id = self._resolve_asset_type_id(node.asset_type)
+
+                    if not type_id:
+                        result.errors.append(
+                            f"Asset type not found: {node.asset_type} for node {node.name}"
+                        )
+                        continue
+
+                    asset_data = {
+                        "name": node.name,
+                        "domain_id": domain_id,
+                        "type_id": type_id,
+                    }
+                    if node.display_name:
+                        asset_data["display_name"] = node.display_name
+                    if status_id:
+                        asset_data["status_id"] = status_id
+
+                    created = self.connector.asset.add_asset(**asset_data)
+                    asset_id = created.get("id")
+                    result.assets[internal_id] = asset_id
+                    result.assets_created += 1
+
+                    # Set description attribute if provided
+                    if node.description:
+                        try:
+                            self.connector.asset.set_asset_attributes(
+                                asset_id=asset_id,
+                                type_public_id="Description",
+                                values=[node.description]
+                            )
+                        except Exception:
+                            pass  # Description is optional
+
+                    # Set custom attributes
+                    for attr_name, attr_value in node.attributes.items():
+                        try:
+                            self.connector.asset.set_asset_attributes(
+                                asset_id=asset_id,
+                                type_public_id=attr_name,
+                                values=[attr_value]
+                            )
+                        except Exception:
+                            pass  # Custom attributes are optional
+
+            except Exception as e:
+                result.errors.append(f"Failed to create asset {node.name}: {str(e)}")
+                result.success = False
+
+        # Phase 2: Create relations
+        for edge in self._edges:
+            try:
+                source_id = result.assets.get(edge.source._internal_id)
+                target_id = result.assets.get(edge.target._internal_id)
+
+                if not source_id:
+                    result.errors.append(
+                        f"Source asset not found for edge: {edge.source.name}"
+                    )
+                    continue
+
+                if not target_id:
+                    result.errors.append(
+                        f"Target asset not found for edge: {edge.target.name}"
+                    )
+                    continue
+
+                # Resolve relation type
+                relation_type_id = edge.relation_type_id
+                if not relation_type_id:
+                    relation_type_id = self._resolve_relation_type_id(edge.relation_type)
+
+                if not relation_type_id:
+                    result.errors.append(
+                        f"Relation type not found: {edge.relation_type}"
+                    )
+                    continue
+
+                created = self.connector.relation.add_relation(
+                    source_id=source_id,
+                    target_id=target_id,
+                    type_id=relation_type_id
+                )
+                relation_id = created.get("id")
+                result.relations.append(relation_id)
+                result.relations_created += 1
+
+            except Exception as e:
+                result.errors.append(
+                    f"Failed to create relation {edge.source.name} -> {edge.target.name}: {str(e)}"
+                )
+                result.success = False
+
+        if result.errors:
+            result.success = False
+
+        return result
+
+    def clear(self) -> "LineageBuilder":
+        """Clear all nodes and edges."""
+        self._nodes.clear()
+        self._edges.clear()
+        self._current_source = None
+        return self
+
+    def get_nodes(self) -> List[LineageNode]:
+        """Get all nodes in the graph."""
+        return list(self._nodes.values())
+
+    def get_edges(self) -> List[LineageEdge]:
+        """Get all edges in the graph."""
+        return list(self._edges)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Export the lineage graph as a dictionary."""
+        return {
+            "nodes": [
+                {
+                    "id": n._internal_id,
+                    "name": n.name,
+                    "asset_type": n.asset_type,
+                    "asset_id": n.asset_id,
+                    "description": n.description,
+                    "attributes": n.attributes
+                }
+                for n in self._nodes.values()
+            ],
+            "edges": [
+                {
+                    "source": e.source._internal_id,
+                    "target": e.target._internal_id,
+                    "relation_type": e.relation_type
+                }
+                for e in self._edges
+            ]
+        }
+
+    def from_dict(self, data: Dict[str, Any]) -> "LineageBuilder":
+        """Import lineage graph from a dictionary."""
+        self.clear()
+
+        node_map: Dict[str, LineageNode] = {}
+
+        for node_data in data.get("nodes", []):
+            node = LineageNode(
+                name=node_data["name"],
+                asset_type=node_data.get("asset_type", "Data Asset"),
+                asset_id=node_data.get("asset_id"),
+                description=node_data.get("description"),
+                attributes=node_data.get("attributes", {})
+            )
+            node._internal_id = node_data.get("id", node._internal_id)
+            node_map[node._internal_id] = node
+            self.add_node(node)
+
+        for edge_data in data.get("edges", []):
+            source = node_map.get(edge_data["source"])
+            target = node_map.get(edge_data["target"])
+            if source and target:
+                self.add_edge(source, target, edge_data.get("relation_type"))
+
+        return self
+
+    def visualize(self) -> str:
+        """
+        Generate a simple ASCII visualization of the lineage.
+
+        Returns:
+            ASCII representation of the lineage graph.
+        """
+        lines = ["Lineage Graph:", "=" * 40]
+
+        # Group edges by source
+        by_source: Dict[str, List[LineageEdge]] = {}
+        for edge in self._edges:
+            key = edge.source.name
+            if key not in by_source:
+                by_source[key] = []
+            by_source[key].append(edge)
+
+        for source_name, edges in by_source.items():
+            lines.append(f"\n[{source_name}]")
+            for edge in edges:
+                lines.append(f"  --({edge.relation_type})--> [{edge.target.name}]")
+
+        return "\n".join(lines)