infrahub-server 1.2.0rc0__py3-none-any.whl → 1.2.2__py3-none-any.whl

infrahub/task_manager/event.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING, Any
 
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.events.schemas.events import Event as PrefectEventModel
+from prefect.exceptions import PrefectHTTPStatusError
 from pydantic import BaseModel, Field, TypeAdapter
 
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
+from infrahub.exceptions import ServiceUnavailableError
 from infrahub.log import get_logger
 from infrahub.utils import get_nested_dict
 
@@ -22,6 +26,8 @@ class PrefectEventData(PrefectEventModel):
                 continue
             if "infrahub.resource.label" not in resource:
                 continue
+            if resource.get("infrahub.resource.label") == GLOBAL_BRANCH_NAME:
+                return None
             return resource.get("infrahub.resource.label")
         return None
 
@@ -86,11 +92,10 @@ class PrefectEventData(PrefectEventModel):
 
     def _return_node_mutation(self) -> dict[str, Any]:
         attributes = []
+        relationships = []
 
         for resource in self.related:
-            if resource.get("prefect.resource.role") == "infrahub.node.field_update" and resource.get(
-                "infrahub.attribute.name"
-            ):
+            if resource.role == "infrahub.node.attribute_update" and resource.get("infrahub.attribute.name"):
                 attributes.append(
                     {
                         "name": resource.get("infrahub.attribute.name", ""),
@@ -104,15 +109,91 @@ class PrefectEventData(PrefectEventModel):
                         "action": resource.get("infrahub.attribute.action", "unchanged"),
                     }
                 )
+            elif resource.role == "infrahub.node.relationship_update":
+                relationships.append(
+                    {
+                        "name": resource.get("infrahub.relationship.name"),
+                        "action": resource.get("infrahub.relationship.peer_status"),
+                        "peer": {
+                            "id": resource.get("infrahub.relationship.peer_id"),
+                            "kind": resource.get("infrahub.relationship.peer_kind"),
+                        },
+                    }
+                )
+
+        return {"attributes": attributes, "relationships": relationships}
+
+    def _get_branch_name_from_resource(self) -> str:
+        return self.resource.get("infrahub.branch.name") or ""
+
+    def _return_artifact_event(self) -> dict[str, Any]:
+        checksum = ""
+        checksum_previous: str | None = None
+        storage_id = ""
+        storage_id_previous: str | None = None
+        artifact_definition_id = ""
+        for resource in self.related:
+            if resource.role == "infrahub.artifact":
+                checksum = resource.get("infrahub.artifact.checksum") or ""
+                checksum_previous = resource.get("infrahub.artifact.checksum_previous")
+                storage_id = resource.get("infrahub.artifact.storage_id") or ""
+                storage_id_previous = resource.get("infrahub.artifact.storage_id_previous")
+                artifact_definition_id = resource.get("infrahub.artifact.artifact_definition_id") or ""
+
+        return {
+            "checksum": checksum,
+            "checksum_previous": checksum_previous,
+            "storage_id": storage_id,
+            "storage_id_previous": storage_id_previous,
+            "artifact_definition_id": artifact_definition_id,
+        }
+
+    def _return_branch_created(self) -> dict[str, Any]:
+        return {"created_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_deleted(self) -> dict[str, Any]:
+        return {"deleted_branch": self._get_branch_name_from_resource()}
 
-        return {"attributes": attributes}
+    def _return_branch_merged(self) -> dict[str, Any]:
+        return {"source_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_rebased(self) -> dict[str, Any]:
+        return {"rebased_branch": self._get_branch_name_from_resource()}
+
+    def _return_group_event(self) -> dict[str, Any]:
+        members = []
+        ancestors = []
+
+        for resource in self.related:
+            if resource.role == "infrahub.group.member" and resource.get("infrahub.node.kind"):
+                members.append({"id": resource.id, "kind": resource.get("infrahub.node.kind")})
+            elif resource.role == "infrahub.group.ancestor" and resource.get("infrahub.node.kind"):
+                ancestors.append({"id": resource.id, "kind": resource.get("infrahub.node.kind")})
+
+        return {"members": members, "ancestors": ancestors}
 
     def _return_event_specifics(self) -> dict[str, Any]:
+        """Return event specific data based on the type of event being processed"""
+
+        event_specifics = {}
+
         match self.event:
+            case "infrahub.artifact.created" | "infrahub.artifact.updated":
+                event_specifics = self._return_artifact_event()
             case "infrahub.node.created" | "infrahub.node.updated" | "infrahub.node.deleted":
-                return self._return_node_mutation()
-
-        return {}
+                event_specifics = self._return_node_mutation()
+            case "infrahub.branch.created":
+                event_specifics = self._return_branch_created()
+            case "infrahub.branch.deleted":
+                event_specifics = self._return_branch_deleted()
+            case "infrahub.branch.merged":
+                event_specifics = self._return_branch_merged()
+            case "infrahub.branch.rebased":
+                event_specifics = self._return_branch_rebased()
+            case "infrahub.group.member_added" | "infrahub.group.member_removed":
+                event_specifics = self._return_group_event()
+
+        return event_specifics
 
     def to_graphql(self) -> dict[str, Any]:
         response = {
@@ -148,8 +229,20 @@ class PrefectEvent:
     ) -> PrefectEventResponse:
         body = {"limit": limit, "filter": filters.model_dump(mode="json", exclude_none=True), "offset": offset}
 
-        response = await client._client.post("/infrahub/events/filter", json=body)
-        response.raise_for_status()
+        # Retry due to https://github.com/PrefectHQ/prefect/issues/16299
+        for _ in range(1, 5):
+            prefect_error: PrefectHTTPStatusError | None = None
+            try:
+                response = await client._client.post("/infrahub/events/filter", json=body)
+                break
+            except PrefectHTTPStatusError as exc:
+                prefect_error = exc
+                await asyncio.sleep(0.1)
+
+        if prefect_error:
+            raise ServiceUnavailableError(
+                message=f"Unable to query prefect due to invalid response from the server (status_code={prefect_error.response.status_code})"
+            )
         data: dict[str, Any] = response.json()
 
         return PrefectEventResponse(

infrahub/task_manager/models.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import uuid
 from collections import defaultdict
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 from prefect.client.schemas.objects import Log as PrefectLog  # noqa: TC002
@@ -31,7 +31,7 @@ class RelatedNodeInfo(BaseModel):
 
 
 class RelatedNodesInfo(BaseModel):
-    flows: dict[UUID, dict[str, RelatedNodeInfo]] = Field(default_factory=lambda: defaultdict(dict))
+    flows: dict[UUID, dict[str, RelatedNodeInfo]] = Field(default_factory=lambda: defaultdict(dict))  # type: ignore[arg-type]
     nodes: dict[str, RelatedNodeInfo] = Field(default_factory=dict)
 
     def add_nodes(self, flow_id: UUID, node_ids: list[str]) -> None:
@@ -64,7 +64,7 @@ class RelatedNodesInfo(BaseModel):
 
 
 class FlowLogs(BaseModel):
-    logs: defaultdict[UUID, list[PrefectLog]] = Field(default_factory=lambda: defaultdict(list))
+    logs: defaultdict[UUID, list[PrefectLog]] = Field(default_factory=lambda: defaultdict(list))  # type: ignore[arg-type]
 
     def to_graphql(self, flow_id: UUID) -> list[dict]:
         return [
@@ -122,7 +122,26 @@ class InfrahubEventFilter(EventFilter):
         if ids:
             self.id = EventIDFilter(id=[uuid.UUID(id) for id in ids])
 
-    def add_event_type_filter(self, event_type: list[str] | None = None) -> None:
+    def add_event_type_filter(
+        self, event_type: list[str] | None = None, event_type_filter: dict[str, Any] | None = None
+    ) -> None:
+        event_type = event_type or []
+        event_type_filter = event_type_filter or {}
+
+        if branch_merged := event_type_filter.get("branch_merged"):
+            branches: list[str] = branch_merged.get("branches") or []
+            if "infrahub.branch.created" not in event_type:
+                event_type.append("infrahub.branch.merged")
+            if branches:
+                self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.branch.name": branches}))
+
+        if branch_rebased := event_type_filter.get("branch_rebased"):
+            branches = branch_rebased.get("branches") or []
+            if "infrahub.branch.created" not in event_type:
+                event_type.append("infrahub.branch.rebased")
+            if branches:
+                self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.branch.name": branches}))
+
         if event_type:
             self.event = EventNameFilter(name=event_type)
 
@@ -159,6 +178,7 @@ class InfrahubEventFilter(EventFilter):
         parent__ids: list[str] | None = None,
         primary_node__ids: list[str] | None = None,
         event_type: list[str] | None = None,
+        event_type_filter: dict[str, Any] | None = None,
         branches: list[str] | None = None,
         level: int | None = None,
         has_children: bool | None = None,
@@ -167,10 +187,10 @@ class InfrahubEventFilter(EventFilter):
     ) -> InfrahubEventFilter:
         occurred_filter = {}
         if since:
-            occurred_filter["since"] = Timestamp(since.isoformat()).obj
+            occurred_filter["since"] = Timestamp(since.isoformat()).to_datetime()
 
         if until:
-            occurred_filter["until"] = Timestamp(until.isoformat()).obj
+            occurred_filter["until"] = Timestamp(until.isoformat()).to_datetime()
 
         if occurred_filter:
             filters = cls(occurred=EventOccurredFilter(**occurred_filter))
@@ -179,7 +199,7 @@ class InfrahubEventFilter(EventFilter):
 
         filters.add_event_filter(level=level, has_children=has_children)
         filters.add_event_id_filter(ids=ids)
-        filters.add_event_type_filter(event_type=event_type)
+        filters.add_event_type_filter(event_type=event_type, event_type_filter=event_type_filter)
         filters.add_branch_filter(branches=branches)
         filters.add_account_filter(account__ids=account__ids)
         filters.add_parent_filter(parent__ids=parent__ids)

infrahub/tasks/artifact.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 from infrahub_sdk.node import InfrahubNode
 from prefect import task
 from prefect.cache_policies import NONE
@@ -13,8 +11,10 @@ from infrahub.services import InfrahubServices
 
 @task(name="define-artifact", task_run_name="Define Artifact", cache_policy=NONE)  # type: ignore[arg-type]
 async def define_artifact(
-    model: Union[CheckArtifactCreate, RequestArtifactGenerate], service: InfrahubServices
-) -> InfrahubNode:
+    model: CheckArtifactCreate | RequestArtifactGenerate, service: InfrahubServices
+) -> tuple[InfrahubNode, bool]:
+    """Return an artifact together with a flag to indicate if the artifact is created now or already existed."""
+    created = False
     if model.artifact_id:
         artifact = await service.client.get(kind=InfrahubKind.ARTIFACT, id=model.artifact_id, branch=model.branch_name)
     else:
@@ -39,5 +39,6 @@ async def define_artifact(
                         "content_type": model.content_type,
                     },
                 )
-                await artifact.save()
-    return artifact
+                await artifact.save(request_context=model.context.to_request_context())
+                created = True
+    return artifact, created

infrahub/telemetry/constants.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+TELEMETRY_KIND: str = "community"
+TELEMETRY_VERSION: str = "20250318"
+
+
+class InfrahubType(str, Enum):
+    COMMUNITY = "community"
+    ENTERPRISE = "enterprise"

infrahub/telemetry/database.py

@@ -0,0 +1,86 @@
+from neo4j.exceptions import Neo4jError
+from prefect import task
+from prefect.cache_policies import NONE
+
+from infrahub.core import utils
+from infrahub.core.graph.schema import GRAPH_SCHEMA
+from infrahub.core.query import QueryType
+from infrahub.database import DatabaseType, InfrahubDatabase
+
+from .models import TelemetryDatabaseData, TelemetryDatabaseServerData, TelemetryDatabaseSystemInfoData
+
+
+async def get_server_info(db: InfrahubDatabase) -> list[TelemetryDatabaseServerData]:
+    data: list[TelemetryDatabaseServerData] = []
+
+    try:
+        results = await db.execute_query(query="SHOW SERVERS YIELD *", name="get_server_info", type=QueryType.READ)
+    except Neo4jError:
+        return []
+
+    for result in results:
+        data.append(
+            TelemetryDatabaseServerData(
+                name=result["name"],
+                version=result["version"],
+            )
+        )
+
+    return data
+
+
+async def get_system_info(db: InfrahubDatabase) -> TelemetryDatabaseSystemInfoData:
+    query = """
+    CALL dbms.queryJmx("java.lang:type=OperatingSystem")
+    YIELD attributes
+    RETURN
+        attributes.AvailableProcessors as processor_available,
+        attributes.TotalMemorySize as memory_total,
+        attributes.FreeMemorySize as memory_available
+    """
+    results = await db.execute_query(query=query, name="get_system_info", type=QueryType.READ)
+
+    return TelemetryDatabaseSystemInfoData(
+        memory_total=results[0]["memory_total"]["value"],
+        memory_available=results[0]["memory_available"]["value"],
+        processor_available=results[0]["processor_available"]["value"],
+    )
+
+
+@task(name="telemetry-gather-db", task_run_name="Gather Database Information", cache_policy=NONE)
+async def gather_database_information(db: InfrahubDatabase) -> TelemetryDatabaseData:
+    async with db.start_session() as dbs:
+        server_info = []
+        system_info = None
+        database_type = db.db_type.value
+
+        if db.db_type == DatabaseType.NEO4J:
+            server_info = await get_server_info(db=dbs)
+            system_info = await get_system_info(db=dbs)
+
+            # server_info is only available on Neo4j Enterprise
+            #  so if it's not empty, we can assume the database is of type Enterprise
+            if len(server_info) == 0:
+                database_type = f"{database_type}-community"
+            else:
+                database_type = f"{database_type}-enterprise"
+
+        data = TelemetryDatabaseData(
+            database_type=database_type,
+            relationship_count={
+                "total": await utils.count_relationships(db=dbs),
+            },
+            node_count={
+                "total": await utils.count_nodes(db=dbs),
+            },
+            servers=server_info,
+            system_info=system_info,
+        )
+
+        for name in GRAPH_SCHEMA["relationships"]:
+            data.relationship_count[name] = await utils.count_relationships(db=dbs, label=name)
+
+        for name in GRAPH_SCHEMA["nodes"]:
+            data.node_count[name] = await utils.count_nodes(db=dbs, label=name)
+
+        return data

infrahub/telemetry/models.py

@@ -0,0 +1,65 @@
+from pydantic import BaseModel
+
+from .constants import InfrahubType
+
+
+class TelemetryWorkerData(BaseModel):
+    total: int
+    active: int
+
+
+class TelemetryBranchData(BaseModel):
+    total: int
+
+
+class TelemetrySchemaData(BaseModel):
+    node_count: int
+    generic_count: int
+    last_update: str
+
+
+class TelemetryDatabaseServerData(BaseModel):
+    name: str
+    version: str
+
+
+class TelemetryDatabaseSystemInfoData(BaseModel):
+    memory_total: int
+    memory_available: int
+    processor_available: int
+
+
+class TelemetryDatabaseData(BaseModel):
+    database_type: str
+    relationship_count: dict[str, int]
+    node_count: dict[str, int]
+    servers: list[TelemetryDatabaseServerData]
+    system_info: TelemetryDatabaseSystemInfoData | None
+
+
+class TelemetryWorkPoolData(BaseModel):
+    name: str
+    type: str
+    total_workers: int
+    active_workers: int
+
+
+class TelemetryPrefectData(BaseModel):
+    events: dict[str, int]
+    automations: dict[str, int]
+    work_pools: list[TelemetryWorkPoolData]
+
+
+class TelemetryData(BaseModel):
+    deployment_id: str | None
+    execution_time: float | None
+    infrahub_version: str
+    infrahub_type: InfrahubType
+    python_version: str
+    platform: str
+    workers: TelemetryWorkerData
+    branches: TelemetryBranchData
+    features: dict[str, int]
+    schema_info: TelemetrySchemaData
+    database: TelemetryDatabaseData
+    prefect: TelemetryPrefectData

infrahub/telemetry/task_manager.py

@@ -0,0 +1,77 @@
+from typing import Any
+
+from prefect import task
+from prefect.cache_policies import NONE
+from prefect.client.orchestration import PrefectClient, get_client
+from prefect.client.schemas.objects import WorkerStatus
+
+from infrahub.events.utils import get_all_events
+from infrahub.trigger.constants import NAME_SEPARATOR
+from infrahub.trigger.models import TriggerType
+
+from .models import TelemetryPrefectData, TelemetryWorkPoolData
+
+
+@task(name="telemetry-gather-work-pools", task_run_name="Gather Work Pools", cache_policy=NONE)
+async def gather_prefect_work_pools(client: PrefectClient) -> list[TelemetryWorkPoolData]:
+    work_pools = await client.read_work_pools()
+    data: list[TelemetryWorkPoolData] = []
+
+    for pool in work_pools:
+        workers = await client.read_workers_for_work_pool(work_pool_name=pool.name)
+        data.append(
+            TelemetryWorkPoolData(
+                name=pool.name,
+                type=pool.type,
+                total_workers=len(workers),
+                active_workers=len([item for item in workers if item.status == WorkerStatus.ONLINE]),
+            )
+        )
+
+    return data
+
+
+@task(name="telemetry-gather-events", task_run_name="Gather Events", cache_policy=NONE)
+async def gather_prefect_events(client: PrefectClient) -> dict[str, Any]:
+    infrahub_events = get_all_events()
+    events: dict[str, int] = {}
+
+    async def count_events(event_name: str) -> int:
+        payload = {"filter": {"event": {"name": [event_name]}}}
+        response = await client._client.post("/events/count-by/event", json=payload)
+        response.raise_for_status()
+        data = response.json()
+        if not isinstance(data, list) or len(data) == 0:
+            return 0
+        return data[0]["count"]
+
+    for event in infrahub_events:
+        events[event.event_name] = await count_events(event_name=event.event_name)
+
+    return events
+
+
+@task(name="telemetry-gather-automations", task_run_name="Gather Automations", cache_policy=NONE)
+async def gather_prefect_automations(client: PrefectClient) -> dict[str, Any]:
+    automations = await client.read_automations()
+
+    data: dict[str, Any] = {}
+
+    for trigger_type in TriggerType:
+        data[trigger_type.value] = len(
+            [item for item in automations if item.name.startswith(f"{trigger_type.value}{NAME_SEPARATOR}")]
+        )
+
+    return data
+
+
+@task(name="telemetry-gather-prefect-information", task_run_name="Gather Prefect Information", cache_policy=NONE)
+async def gather_prefect_information() -> TelemetryPrefectData:
+    async with get_client(sync_client=False) as client:
+        data = TelemetryPrefectData(
+            work_pools=await gather_prefect_work_pools(client=client),
+            events=await gather_prefect_events(client=client),
+            automations=await gather_prefect_automations(client=client),
+        )
+
+        return data

infrahub/{tasks/telemetry.py → telemetry/tasks.py}

@@ -12,44 +12,32 @@ from infrahub import __version__, config
 from infrahub.core import registry, utils
 from infrahub.core.branch import Branch
 from infrahub.core.constants import InfrahubKind
-from infrahub.core.graph.schema import GRAPH_SCHEMA
 from infrahub.services import InfrahubServices
 
-TELEMETRY_KIND: str = "community"
-TELEMETRY_VERSION: str = "20240524"
+from .constants import TELEMETRY_KIND, TELEMETRY_VERSION
+from .database import gather_database_information
+from .models import (
+    TelemetryBranchData,
+    TelemetryData,
+    TelemetrySchemaData,
+    TelemetryWorkerData,
+)
+from .task_manager import gather_prefect_information
+from .utils import determine_infrahub_type
 
 
-@task(name="telemetry-gather-db", task_run_name="Gather Database Information", cache_policy=NONE)  # type: ignore[arg-type]
-async def gather_database_information(service: InfrahubServices) -> dict:
-    async with service.database.start_session() as db:
-        data: dict[str, Any] = {
-            "database_type": db.db_type.value,
-            "relationship_count": {"total": await utils.count_relationships(db=db)},
-            "node_count": {"total": await utils.count_nodes(db=db)},
-        }
-
-        for name in GRAPH_SCHEMA["relationships"]:
-            data["relationship_count"][name] = await utils.count_relationships(db=db, label=name)
-
-        for name in GRAPH_SCHEMA["nodes"]:
-            data["node_count"][name] = await utils.count_nodes(db=db, label=name)
-
-        return data
-
-
-@task(name="telemetry-schema-information", task_run_name="Gather Schema Information", cache_policy=NONE)  # type: ignore[arg-type]
-async def gather_schema_information(branch: Branch) -> dict:
-    data: dict[str, Any] = {}
+@task(name="telemetry-schema-information", task_run_name="Gather Schema Information", cache_policy=NONE)
+async def gather_schema_information(branch: Branch) -> TelemetrySchemaData:
     main_schema = registry.schema.get_schema_branch(name=branch.name)
-    data["node_count"] = len(main_schema.node_names)
-    data["generic_count"] = len(main_schema.generic_names)
-    data["last_update"] = branch.schema_changed_at
-
-    return data
+    return TelemetrySchemaData(
+        node_count=len(main_schema.node_names),
+        generic_count=len(main_schema.generic_names),
+        last_update=branch.schema_changed_at or "",
+    )
 
 
-@task(name="telemetry-feature-information", task_run_name="Gather Feature Information", cache_policy=NONE)  # type: ignore[arg-type]
-async def gather_feature_information(service: InfrahubServices) -> dict:
+@task(name="telemetry-feature-information", task_run_name="Gather Feature Information", cache_policy=NONE)
+async def gather_feature_information(service: InfrahubServices) -> dict[str, int]:
     async with service.database.start_session() as db:
         data = {}
         features_to_count = [
@@ -59,7 +47,9 @@ async def gather_feature_information(service: InfrahubServices) -> dict:
             InfrahubKind.GENERICGROUP,
             InfrahubKind.PROFILE,
             InfrahubKind.PROPOSEDCHANGE,
+            InfrahubKind.OBJECTTEMPLATE,
             InfrahubKind.TRANSFORM,
+            InfrahubKind.WEBHOOK,
         ]
         for kind in features_to_count:
             data[kind] = await utils.count_nodes(db=db, label=kind)
@@ -67,37 +57,39 @@ async def gather_feature_information(service: InfrahubServices) -> dict:
         return data
 
 
-@task(name="telemetry-gather-data", task_run_name="Gather Anonynous Data", cache_policy=NONE)  # type: ignore[arg-type]
-async def gather_anonymous_telemetry_data(service: InfrahubServices) -> dict:
+@task(name="telemetry-gather-data", task_run_name="Gather Anonynous Data", cache_policy=NONE)
+async def gather_anonymous_telemetry_data(service: InfrahubServices) -> TelemetryData:
     start_time = time.time()
 
     default_branch = registry.get_branch_from_registry()
     workers = await service.component.list_workers(branch=default_branch.name, schema_hash=False)
 
-    data: dict[str, Any] = {
-        "deployment_id": registry.id,
-        "execution_time": None,
-        "infrahub_version": __version__,
-        "python_version": platform.python_version(),
-        "platform": platform.machine(),
-        "workers": {
-            "total": len(workers),
-            "active": len([w for w in workers if w.active]),
-        },
-        "branches": {
-            "total": len(registry.branch),
-        },
-        "features": await gather_feature_information(service=service),
-        "schema": await gather_schema_information(branch=default_branch),
-        "database": await gather_database_information(service=service),
-    }
-
-    data["execution_time"] = time.time() - start_time
+    data = TelemetryData(
+        deployment_id=registry.id,
+        execution_time=None,
+        infrahub_version=__version__,
+        infrahub_type=determine_infrahub_type(),
+        python_version=platform.python_version(),
+        platform=platform.machine(),
+        workers=TelemetryWorkerData(
+            total=len(workers),
+            active=len([w for w in workers if w.active]),
+        ),
+        branches=TelemetryBranchData(
+            total=len(registry.branch),
+        ),
+        features=await gather_feature_information(service=service),
+        schema_info=await gather_schema_information(branch=default_branch),
+        database=await gather_database_information(db=service.database),
+        prefect=await gather_prefect_information(),
+    )
+
+    data.execution_time = time.time() - start_time
 
     return data
 
 
-@task(name="telemetry-post-data", task_run_name="Upload data", retries=5, cache_policy=NONE)  # type: ignore[arg-type]
+@task(name="telemetry-post-data", task_run_name="Upload data", retries=5, cache_policy=NONE)
 async def post_telemetry_data(service: InfrahubServices, url: str, payload: dict[str, Any]) -> None:
     """Send the telemetry data to the specified URL, using HTTP POST."""
     response = await service.http.post(url=url, json=payload)
@@ -114,13 +106,14 @@ async def send_telemetry_push(service: InfrahubServices) -> None:
     log.info(f"Pushing anonymous telemetry data to {config.SETTINGS.main.telemetry_endpoint}...")
 
     data = await gather_anonymous_telemetry_data(service=service)
-    log.info(f"Anonymous usage telemetry gathered in {data['execution_time']} seconds. | {data}")
+    data_dict = data.model_dump(mode="json")
+    log.info(f"Anonymous usage telemetry gathered in {data.execution_time} seconds. | {data_dict}")
 
     payload = {
         "kind": TELEMETRY_KIND,
         "payload_format": TELEMETRY_VERSION,
-        "data": data,
-        "checksum": hashlib.sha256(json.dumps(data).encode()).hexdigest(),
+        "data": data_dict,
+        "checksum": hashlib.sha256(json.dumps(data_dict).encode()).hexdigest(),
     }
 
     await post_telemetry_data(service=service, url=config.SETTINGS.main.telemetry_endpoint, payload=payload)

infrahub/telemetry/utils.py

@@ -0,0 +1,11 @@
+import importlib.metadata
+
+from .constants import InfrahubType
+
+
+def determine_infrahub_type() -> InfrahubType:
+    try:
+        importlib.metadata.version("infrahub-enterprise")
+        return InfrahubType.ENTERPRISE
+    except importlib.metadata.PackageNotFoundError:
+        return InfrahubType.COMMUNITY