cartography 0.105.0__py3-none-any.whl → 0.106.0__py3-none-any.whl

cartography/intel/airbyte/__init__.py

@@ -0,0 +1,105 @@
+import logging
+
+import neo4j
+
+import cartography.intel.airbyte.connections
+import cartography.intel.airbyte.destinations
+import cartography.intel.airbyte.organizations
+import cartography.intel.airbyte.sources
+import cartography.intel.airbyte.tags
+import cartography.intel.airbyte.users
+import cartography.intel.airbyte.workspaces
+from cartography.config import Config
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def start_airbyte_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
+    """
+    If this module is configured, perform ingestion of Airbyte data. Otherwise warn and exit
+    :param neo4j_session: Neo4J session for database interface
+    :param config: A cartography.config object
+    :return: None
+    """
+    if (
+        not config.airbyte_api_url
+        or not config.airbyte_client_id
+        or not config.airbyte_client_secret
+    ):
+        logger.info(
+            "Airbyte import is not configured - skipping this module. "
+            "See docs to configure.",
+        )
+        return
+
+    # Create api session
+    api_client = AirbyteClient(
+        base_url=config.airbyte_api_url,
+        client_id=config.airbyte_client_id,
+        client_secret=config.airbyte_client_secret,
+    )
+
+    common_job_parameters = {
+        "UPDATE_TAG": config.update_tag,
+    }
+
+    organizations = cartography.intel.airbyte.organizations.sync(
+        neo4j_session,
+        api_client,
+        common_job_parameters,
+    )
+
+    for organization in organizations:
+        org_common_job_parameters = {
+            "UPDATE_TAG": config.update_tag,
+            "ORG_ID": organization["organizationId"],
+        }
+        workspaces = cartography.intel.airbyte.workspaces.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            org_common_job_parameters,
+        )
+        workspace_ids = [workspace["workspaceId"] for workspace in workspaces]
+
+        cartography.intel.airbyte.users.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            org_common_job_parameters,
+        )
+
+        cartography.intel.airbyte.sources.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            workspace_ids,
+            org_common_job_parameters,
+        )
+
+        cartography.intel.airbyte.destinations.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            workspace_ids,
+            org_common_job_parameters,
+        )
+
+        cartography.intel.airbyte.tags.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            workspace_ids,
+            org_common_job_parameters,
+        )
+
+        cartography.intel.airbyte.connections.sync(
+            neo4j_session,
+            api_client,
+            organization["organizationId"],
+            workspace_ids,
+            org_common_job_parameters,
+        )

cartography/intel/airbyte/connections.py

@@ -0,0 +1,120 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.intel.airbyte.util import list_to_string
+from cartography.models.airbyte.connection import AirbyteConnectionSchema
+from cartography.models.airbyte.stream import AirbyteStreamSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    workspace_ids: List[str],
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    connections = get(api_session, workspace_ids)
+    transformed_connections, transformed_streams = transform(connections)
+    load_connections(
+        neo4j_session,
+        transformed_connections,
+        org_id,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    load_streams(
+        neo4j_session, transformed_streams, org_id, common_job_parameters["UPDATE_TAG"]
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    workspace_ids: List[str],
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/connections",
+        {"workspaceIds": ",".join(workspace_ids)} if workspace_ids else None,
+    )
+
+
+def transform(
+    connections: List[Dict[str, Any]],
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+    transformed_connections = []
+    transformed_streams = []
+    for connection in connections:
+        connection["tags_ids"] = [tag["tagId"] for tag in connection.get("tags", [])]
+        transformed_connections.append(connection)
+        for stream in connection.get("configurations", {}).get("streams", []):
+            formated_stream = {
+                "connectionId": connection["connectionId"],
+                "streamId": f"{connection['connectionId']}_{stream['name']}",
+                "name": stream["name"],
+                "syncMode": stream["syncMode"],
+                "cursorField": list_to_string(stream.get("cursorField", [])),
+                "primaryKey": list_to_string(stream.get("primaryKey", [])),
+                "includeFiles": stream.get("includeFiles", False),
+                "selectedFields": list_to_string(stream.get("selectedFields", [])),
+                "mappers": list_to_string(stream.get("mappers", [])),
+            }
+            transformed_streams.append(formated_stream)
+    return transformed_connections, transformed_streams
+
+
+@timeit
+def load_connections(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Connections into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteConnectionSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def load_streams(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Streams into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteStreamSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteStreamSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(AirbyteConnectionSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/destinations.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.intel.airbyte.util import normalize_airbyte_config
+from cartography.models.airbyte.destination import AirbyteDestinationSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    workspace_ids: List[str],
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    destinations = get(api_session, workspace_ids)
+    transformed_destinations = transform(destinations)
+    load_destinations(
+        neo4j_session,
+        transformed_destinations,
+        org_id,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    workspace_ids: List[str],
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/destinations",
+        params={"workspaceIds": ",".join(workspace_ids)} if workspace_ids else None,
+    )
+
+
+def transform(destinations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    transformed_destinations = []
+    for destination in destinations:
+        destination["configuration"] = normalize_airbyte_config(
+            destination.get("configuration", {})
+        )
+        transformed_destinations.append(destination)
+    return transformed_destinations
+
+
+@timeit
+def load_destinations(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Destinations into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteDestinationSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteDestinationSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/organizations.py

@@ -0,0 +1,59 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.models.airbyte.organization import AirbyteOrganizationSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    common_job_parameters: Dict[str, Any],
+) -> List[Dict]:
+    organizations = get(api_session)
+    load_organizations(
+        neo4j_session, organizations, common_job_parameters["UPDATE_TAG"]
+    )
+    cleanup(neo4j_session, common_job_parameters)
+    return organizations
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+) -> List[Dict[str, Any]]:
+    return api_session.get("/organizations")
+
+
+@timeit
+def load_organizations(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Organizations into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteOrganizationSchema(),
+        data,
+        lastupdated=update_tag,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteOrganizationSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/sources.py

@@ -0,0 +1,78 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.intel.airbyte.util import normalize_airbyte_config
+from cartography.models.airbyte.source import AirbyteSourceSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    workspace_ids: List[str],
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    sources = get(api_session, workspace_ids)
+    transformed_sources = transform(sources)
+    load_sources(
+        neo4j_session, transformed_sources, org_id, common_job_parameters["UPDATE_TAG"]
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    workspace_ids: List[str],
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/sources",
+        params={"workspaceIds": ",".join(workspace_ids)} if workspace_ids else None,
+    )
+
+
+def transform(sources: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    transformed_sources = []
+    for source in sources:
+        source["configuration"] = normalize_airbyte_config(
+            source.get("configuration", {})
+        )
+        transformed_sources.append(source)
+    return transformed_sources
+
+
+@timeit
+def load_sources(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Sources into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteSourceSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteSourceSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/tags.py

@@ -0,0 +1,64 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.models.airbyte.tag import AirbyteTagSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    workspace_ids: List[str],
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    tags = get(api_session, workspace_ids)
+    load_tags(neo4j_session, tags, org_id, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    workspace_ids: List[str],
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/tags",
+        params={"workspaceIds": ",".join(workspace_ids)} if workspace_ids else None,
+    )
+
+
+@timeit
+def load_tags(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Tags into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteTagSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteTagSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/users.py

@@ -0,0 +1,106 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.models.airbyte.user import AirbyteUserSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    users = get(api_session, org_id=org_id)
+    for user in users:
+        permissions = get_permissions(api_session, user["id"], org_id=org_id)
+        org_admin, workspace_admin, workspace_member = transform_permissions(
+            permissions
+        )
+        user["adminOfOrganization"] = org_admin
+        user["adminOfWorkspace"] = workspace_admin
+        user["memberOfWorkspace"] = workspace_member
+    load_users(neo4j_session, users, org_id, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    org_id: str,
+) -> List[Dict[str, Any]]:
+    return api_session.get("/users", params={"organizationId": org_id})
+
+
+@timeit
+def get_permissions(
+    api_session: AirbyteClient,
+    user_id: str,
+    org_id: str,
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/permissions",
+        params={"organizationId": org_id, "userId": user_id},
+    )
+
+
+@timeit
+def transform_permissions(
+    permissions: List[Dict[str, Any]],
+) -> Tuple[
+    List[str],  # org_admin
+    List[str],  # workspace_admin
+    List[str],  # workspace_member
+]:
+    org_admin: list[str] = []
+    workspace_admin: list[str] = []
+    workspace_member: list[str] = []
+
+    for permission in permissions:
+        # workspace
+        if permission["scope"] == "workspace":
+            if permission["permissionType"] in ("workspace_owner", "workspace_admin"):
+                workspace_admin.append(permission["scopeId"])
+            workspace_member.append(permission["scopeId"])
+        # organization
+        elif permission["scope"] == "organization":
+            if permission["permissionType"] in ("organization_admin",):
+                org_admin.append(permission["scopeId"])
+    return org_admin, workspace_admin, workspace_member
+
+
+@timeit
+def load_users(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Users into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteUserSchema(),
+        data,
+        ORG_ID=org_id,
+        lastupdated=update_tag,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteUserSchema(), common_job_parameters).run(
+        neo4j_session
+    )