from metadata.generated.schema.entity.data.table import Table
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    OpenMetadataConnection,
)
from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
    OpenMetadataJWTClientConfig,
)

server_config = OpenMetadataConnection(
    hostPort="https://your-company.getcollate.io/api",
    authProvider="openmetadata",
    securityConfig=OpenMetadataJWTClientConfig(
        jwtToken="<YOUR-JWT-TOKEN>"
    ),
)
metadata = OpenMetadata(server_config)

# Get lineage by fully qualified name
lineage = metadata.get_lineage_by_name(
    entity=Table,
    fqn="sample_data.ecommerce_db.shopify.dim_customer",
    up_depth=1,
    down_depth=1,
)

print(f"Entity: {lineage['entity']['fullyQualifiedName']}")
print(f"Upstream edges: {len(lineage.get('upstreamEdges', []))}")
print(f"Downstream edges: {len(lineage.get('downstreamEdges', []))}")

# Get lineage by ID
lineage = metadata.get_lineage_by_id(
    entity=Table,
    entity_id="455e3d9d-dbbf-455e-b3be-7191daa825f3",
    up_depth=3,
    down_depth=2,
)

{
  "entity": {
    "id": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
    "type": "table",
    "name": "dim_customer",
    "fullyQualifiedName": "sample_data.ecommerce_db.shopify.dim_customer",
    "deleted": false,
    "href": "http://localhost:8585/api/v1/tables/455e3d9d-dbbf-455e-b3be-7191daa825f3"
  },
  "nodes": [
    {
      "id": "800caa0f-a149-48d2-a0ce-6ca84501767e",
      "type": "table",
      "name": "raw_customer",
      "fullyQualifiedName": "sample_data.ecommerce_db.shopify.raw_customer",
      "deleted": false,
      "href": "http://localhost:8585/api/v1/tables/800caa0f-a149-48d2-a0ce-6ca84501767e"
    },
    {
      "id": "c3d4e5f6-a1b2-7890-abcd-ef1234567890",
      "type": "dashboard",
      "name": "customer_dashboard",
      "fullyQualifiedName": "sample_superset.customer_dashboard",
      "deleted": false,
      "href": "http://localhost:8585/api/v1/dashboards/c3d4e5f6-a1b2-7890-abcd-ef1234567890"
    }
  ],
  "upstreamEdges": [
    {
      "fromEntity": "800caa0f-a149-48d2-a0ce-6ca84501767e",
      "toEntity": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
      "lineageDetails": {
        "sqlQuery": "INSERT INTO dim_customer SELECT id, name FROM raw_customer",
        "columnsLineage": [
          {
            "fromColumns": [
              "sample_data.ecommerce_db.shopify.raw_customer.id"
            ],
            "toColumn": "sample_data.ecommerce_db.shopify.dim_customer.customer_id"
          }
        ]
      }
    }
  ],
  "downstreamEdges": [
    {
      "fromEntity": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
      "toEntity": "c3d4e5f6-a1b2-7890-abcd-ef1234567890"
    }
  ]
}