infrahub-server 1.2.0b1__py3-none-any.whl → 1.2.1__py3-none-any.whl

infrahub/message_bus/operations/finalize/validator.py

@@ -1,11 +1,22 @@
+from infrahub_sdk.protocols import (
+    CoreArtifactValidator,
+    CoreDataValidator,
+    CoreGeneratorValidator,
+    CoreRepositoryValidator,
+    CoreSchemaValidator,
+    CoreUserValidator,
+    CoreValidator,
+)
 from prefect import flow
 
 from infrahub import config
+from infrahub.core.constants import InfrahubKind, ValidatorConclusion
 from infrahub.core.timestamp import Timestamp
 from infrahub.log import get_logger
 from infrahub.message_bus import messages
 from infrahub.message_bus.types import KVTTL, MessageTTL
 from infrahub.services import InfrahubServices
+from infrahub.validators.events import send_failed_validator, send_passed_validator
 
 log = get_logger()
 
@@ -20,7 +31,8 @@ async def execution(message: messages.FinalizeValidatorExecution, service: Infra
 
     The message will get rescheduled until the timeout has exceeded or until all checks are accounted for.
     """
-    validator = await service.client.get(kind=message.validator_type, id=message.validator_id)
+    validator_type = get_validator_type(validator_type=message.validator_type)
+    validator = await service.client.get(kind=validator_type, id=message.validator_id)
     checks_key = f"validator_execution_id:{message.validator_execution_id}:checks"
     current_conclusion = validator.conclusion.value
     if validator.state.value != "in_progress":
@@ -81,3 +93,41 @@ async def execution(message: messages.FinalizeValidatorExecution, service: Infra
     validator.completed_at.value = Timestamp().to_string()
     validator.conclusion.value = conclusion
     await validator.save()
+    if validator.conclusion.value == ValidatorConclusion.SUCCESS.value:
+        await send_passed_validator(
+            service=service, validator=validator, proposed_change_id=message.proposed_change, context=message.context
+        )
+    else:
+        await send_failed_validator(
+            service=service, validator=validator, proposed_change_id=message.proposed_change, context=message.context
+        )
+
+
+def get_validator_type(
+    validator_type: str,
+) -> (
+    type[CoreArtifactValidator]
+    | type[CoreDataValidator]
+    | type[CoreGeneratorValidator]
+    | type[CoreRepositoryValidator]
+    | type[CoreSchemaValidator]
+    | type[CoreUserValidator]
+    | type[CoreValidator]
+):
+    match validator_type:
+        case InfrahubKind.USERVALIDATOR:
+            validator_kind = CoreUserValidator
+        case InfrahubKind.SCHEMAVALIDATOR:
+            validator_kind = CoreSchemaValidator
+        case InfrahubKind.GENERATORVALIDATOR:
+            validator_kind = CoreGeneratorValidator
+        case InfrahubKind.REPOSITORYVALIDATOR:
+            validator_kind = CoreRepositoryValidator
+        case InfrahubKind.DATAVALIDATOR:
+            validator_kind = CoreDataValidator
+        case InfrahubKind.ARTIFACTVALIDATOR:
+            validator_kind = CoreArtifactValidator
+        case _:
+            validator_kind = CoreValidator
+
+    return validator_kind

infrahub/message_bus/operations/requests/generator_definition.py

@@ -1,12 +1,14 @@
+from infrahub_sdk.protocols import CoreGeneratorValidator
 from infrahub_sdk.uuidt import UUIDT
 from prefect import flow
 from prefect.logging import get_run_logger
 
-from infrahub.core.constants import InfrahubKind, ValidatorConclusion, ValidatorState
+from infrahub.core.constants import InfrahubKind
 from infrahub.core.timestamp import Timestamp
 from infrahub.message_bus import InfrahubMessage, Meta, messages
 from infrahub.message_bus.types import KVTTL
 from infrahub.services import InfrahubServices
+from infrahub.validators.tasks import start_validator
 from infrahub.workflows.utils import add_tags
 
 
@@ -27,31 +29,26 @@ async def check(message: messages.RequestGeneratorDefinitionCheck, service: Infr
 
     await proposed_change.validations.fetch()
 
-    validator = None
+    previous_validator: CoreGeneratorValidator | None = None
     for relationship in proposed_change.validations.peers:
         existing_validator = relationship.peer
         if (
             existing_validator.typename == InfrahubKind.GENERATORVALIDATOR
             and existing_validator.definition.id == message.generator_definition.definition_id
         ):
-            validator = existing_validator
+            previous_validator = existing_validator
 
-    if validator:
-        validator.conclusion.value = ValidatorConclusion.UNKNOWN.value
-        validator.state.value = ValidatorState.QUEUED.value
-        validator.started_at.value = ""
-        validator.completed_at.value = ""
-        await validator.save()
-    else:
-        validator = await service.client.create(
-            kind=InfrahubKind.GENERATORVALIDATOR,
-            data={
-                "label": validator_name,
-                "proposed_change": message.proposed_change,
-                "definition": message.generator_definition.definition_id,
-            },
-        )
-        await validator.save()
+    validator = await start_validator(
+        service=service,
+        validator=previous_validator,
+        validator_type=CoreGeneratorValidator,
+        proposed_change=message.proposed_change,
+        data={
+            "label": validator_name,
+            "definition": message.generator_definition.definition_id,
+        },
+        context=message.context,
+    )
 
     group = await service.client.get(
         kind=InfrahubKind.GENERICGROUP,
@@ -90,6 +87,7 @@ async def check(message: messages.RequestGeneratorDefinitionCheck, service: Infr
             log.info(f"Trigger execution of {message.generator_definition.definition_name} for {member.display_label}")
             events.append(
                 messages.CheckGeneratorRun(
+                    context=message.context,
                     generator_definition=message.generator_definition,
                     generator_instance=generator_instance,
                     commit=repository.source_commit,
@@ -119,6 +117,8 @@ async def check(message: messages.RequestGeneratorDefinitionCheck, service: Infr
             validator_id=validator.id,
             validator_execution_id=validator_execution_id,
             validator_type=InfrahubKind.GENERATORVALIDATOR,
+            context=message.context,
+            proposed_change=message.proposed_change,
         )
     )
     for event in events:

infrahub/message_bus/operations/requests/proposed_change.py

@@ -85,7 +85,9 @@ async def pipeline(message: messages.RequestProposedChangePipeline, service: Inf
                     target_branch=repo.destination_branch,
                 )
                 await service.workflow.submit_workflow(
-                    workflow=GIT_REPOSITORY_INTERNAL_CHECKS_TRIGGER, parameters={"model": model}
+                    workflow=GIT_REPOSITORY_INTERNAL_CHECKS_TRIGGER,
+                    context=message.context,
+                    parameters={"model": model},
                 )
         return
 

infrahub/pools/number.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from infrahub.core.query.resource_manager import NumberPoolGetAllocated
 from infrahub.core.registry import registry
@@ -20,9 +20,7 @@ class UsedNumber:
 
 
 class NumberUtilizationGetter:
-    def __init__(
-        self, db: InfrahubDatabase, pool: CoreNode, branch: Branch, at: Optional[Union[Timestamp, str]] = None
-    ) -> None:
+    def __init__(self, db: InfrahubDatabase, pool: CoreNode, branch: Branch, at: Timestamp | str | None = None) -> None:
         self.db = db
         self.at = at
         self.pool = pool

infrahub/proposed_change/tasks.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 
 import pytest
-from infrahub_sdk.protocols import CoreGeneratorDefinition, CoreProposedChange
+from infrahub_sdk.protocols import CoreArtifactValidator, CoreGeneratorDefinition, CoreProposedChange
 from prefect import flow, task
 from prefect.cache_policies import NONE
 from prefect.client.schemas.objects import (
@@ -22,7 +22,7 @@ from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.branch.tasks import merge_branch
-from infrahub.core.constants import InfrahubKind, RepositoryInternalStatus, ValidatorConclusion, ValidatorState
+from infrahub.core.constants import InfrahubKind, RepositoryInternalStatus, ValidatorConclusion
 from infrahub.core.diff.coordinator import DiffCoordinator
 from infrahub.core.diff.model.diff import DiffElementType, SchemaConflict
 from infrahub.core.diff.model.path import NodeDiffFieldSummary
@@ -52,8 +52,8 @@ from infrahub.proposed_change.models import (
 )
 from infrahub.pytest_plugin import InfrahubBackendPlugin
 from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
+from infrahub.validators.tasks import start_validator
 from infrahub.workflows.catalogue import (
-    COMPUTED_ATTRIBUTE_SETUP_PYTHON,
     GIT_REPOSITORIES_CHECK_ARTIFACT_CREATE,
     GIT_REPOSITORY_INTERNAL_CHECKS_TRIGGER,
     GIT_REPOSITORY_USER_CHECKS_TRIGGER,
@@ -150,7 +150,9 @@ async def merge_proposed_change(
 
         log.info("Proposed change is eligible to be merged")
         try:
-            await merge_branch(branch=source_branch.name, context=context, service=service)
+            await merge_branch(
+                branch=source_branch.name, context=context, service=service, proposed_change_id=proposed_change_id
+            )
         except MergeFailedError as exc:
             await _proposed_change_transition_state(
                 proposed_change=proposed_change, state=ProposedChangeState.OPEN, service=service
@@ -162,7 +164,6 @@ async def merge_proposed_change(
         await _proposed_change_transition_state(
             proposed_change=proposed_change, state=ProposedChangeState.MERGED, service=service
         )
-        await service.workflow.submit_workflow(workflow=COMPUTED_ATTRIBUTE_SETUP_PYTHON, context=context)
         return Completed(message="proposed change merged successfully")
 
 
@@ -275,6 +276,7 @@ async def run_generators(
 
         if select:
             msg = messages.RequestGeneratorDefinitionCheck(
+                context=context,
                 generator_definition=generator_definition,
                 branch_diff=model.branch_diff,
                 proposed_change=model.proposed_change,
@@ -415,7 +417,9 @@ async def _get_proposed_change_schema_integrity_constraints(
     name="proposed-changed-repository-checks",
     flow_run_name="Process user defined checks",
 )
-async def repository_checks(model: RequestProposedChangeRepositoryChecks, service: InfrahubServices) -> None:
+async def repository_checks(
+    model: RequestProposedChangeRepositoryChecks, service: InfrahubServices, context: InfrahubContext
+) -> None:
     await add_tags(branches=[model.source_branch], nodes=[model.proposed_change])
 
     for repository in model.branch_diff.repositories:
@@ -431,7 +435,9 @@ async def repository_checks(model: RequestProposedChangeRepositoryChecks, servic
                 target_branch=model.destination_branch,
             )
             await service.workflow.submit_workflow(
-                workflow=GIT_REPOSITORY_INTERNAL_CHECKS_TRIGGER, parameters={"model": trigger_internal_checks_model}
+                workflow=GIT_REPOSITORY_INTERNAL_CHECKS_TRIGGER,
+                context=context,
+                parameters={"model": trigger_internal_checks_model},
             )
 
         trigger_user_checks_model = TriggerRepositoryUserChecks(
@@ -444,7 +450,9 @@ async def repository_checks(model: RequestProposedChangeRepositoryChecks, servic
             branch_diff=model.branch_diff,
         )
         await service.workflow.submit_workflow(
-            workflow=GIT_REPOSITORY_USER_CHECKS_TRIGGER, parameters={"model": trigger_user_checks_model}
+            workflow=GIT_REPOSITORY_USER_CHECKS_TRIGGER,
+            context=context,
+            parameters={"model": trigger_user_checks_model},
         )
 
 
@@ -532,31 +540,26 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, s
 
     await proposed_change.validations.fetch()
 
-    validator = None
+    previous_validator: CoreArtifactValidator | None = None
     for relationship in proposed_change.validations.peers:
         existing_validator = relationship.peer
         if (
             existing_validator.typename == InfrahubKind.ARTIFACTVALIDATOR
             and existing_validator.definition.id == model.artifact_definition.definition_id
         ):
-            validator = existing_validator
-
-    if validator:
-        validator.conclusion.value = ValidatorConclusion.UNKNOWN.value
-        validator.state.value = ValidatorState.QUEUED.value
-        validator.started_at.value = ""
-        validator.completed_at.value = ""
-        await validator.save()
-    else:
-        validator = await service.client.create(
-            kind=InfrahubKind.ARTIFACTVALIDATOR,
-            data={
-                "label": validator_name,
-                "proposed_change": model.proposed_change,
-                "definition": model.artifact_definition.definition_id,
-            },
-        )
-        await validator.save()
+            previous_validator = existing_validator
+
+    validator = await start_validator(
+        service=service,
+        validator=previous_validator,
+        validator_type=CoreArtifactValidator,
+        proposed_change=model.proposed_change,
+        data={
+            "label": validator_name,
+            "definition": model.artifact_definition.definition_id,
+        },
+        context=model.context,
+    )
 
     await artifact_definition.targets.fetch()
     group = artifact_definition.targets.peer
@@ -617,7 +620,13 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, s
                 )
             )
 
-    await run_checks_and_update_validator(checks, validator)
+    await run_checks_and_update_validator(
+        checks=checks,
+        validator=validator,
+        proposed_change_id=model.proposed_change,
+        context=model.context,
+        service=service,
+    )
 
 
 def _should_render_artifact(artifact_id: str | None, managed_branch: bool, impacted_artifacts: list[str]) -> bool:  # noqa: ARG001

infrahub/pytest_plugin.py

@@ -1,9 +1,7 @@
-from typing import Optional
-
+import pytest
 from infrahub_sdk.client import Config as InfrahubClientConfig
 from infrahub_sdk.client import InfrahubClientSync
 from infrahub_sdk.node import InfrahubNodeSync
-from pytest import Config, Item, Session, TestReport
 
 from infrahub.core.constants import InfrahubKind
 from infrahub.core.timestamp import Timestamp
@@ -52,7 +50,12 @@ class InfrahubBackendPlugin:
 
         return validator, True
 
-    def pytest_collection_modifyitems(self, session: Session, config: Config, items: list[Item]) -> None:  # noqa: ARG002
+    def pytest_collection_modifyitems(
+        self,
+        session: pytest.Session,  # noqa: ARG002
+        config: pytest.Config,  # noqa: ARG002
+        items: list[pytest.Item],
+    ) -> None:
         """This function is called after item collection and gives the opportunity to work on the collection before sending the items for testing.
 
         All items without an "infrahub" marker will be discarded. Items will also be re-ordered to be run in a specific order:
@@ -64,7 +67,7 @@ class InfrahubBackendPlugin:
         """
         filtered_items = [i for i in items if i.get_closest_marker("infrahub")]
 
-        def sort_key(item: Item) -> tuple[int, int]:
+        def sort_key(item: pytest.Item) -> tuple[int, int]:
             type_cost = 99
             for marker_name, priority in ORDER_TYPE_MAP.items():
                 if item.get_closest_marker(marker_name):
@@ -82,7 +85,7 @@ class InfrahubBackendPlugin:
         filtered_items.sort(key=sort_key)
         items[:] = filtered_items
 
-    def pytest_collection_finish(self, session: Session) -> None:  # noqa: ARG002
+    def pytest_collection_finish(self, session: pytest.Session) -> None:  # noqa: ARG002
         """This function is called when tests have been collected and modified, meaning they are ready to be run."""
         self.proposed_change = self.client.get(kind=InfrahubKind.PROPOSEDCHANGE, id=self.proposed_change_id)
         self.proposed_change.validations.fetch()
@@ -95,7 +98,7 @@ class InfrahubBackendPlugin:
                 check = relationship.peer
                 self.checks[check.origin.value] = check
 
-    def pytest_runtestloop(self, session: Session) -> Optional[object]:  # noqa: ARG002
+    def pytest_runtestloop(self, session: pytest.Session) -> object | None:  # noqa: ARG002
         """This function is called when the test loop is being run."""
         self.validator.conclusion.value = "unknown"
         self.validator.state.value = "in_progress"
@@ -104,7 +107,7 @@ class InfrahubBackendPlugin:
 
         return None
 
-    def pytest_runtest_setup(self, item: Item) -> None:
+    def pytest_runtest_setup(self, item: pytest.Item) -> None:
         """Create a StandardCheck for each test item to later record its details.
 
         If a check already exists, reset it to its default values.
@@ -132,7 +135,7 @@ class InfrahubBackendPlugin:
 
         check.save()
 
-    def pytest_runtest_logreport(self, report: TestReport) -> None:
+    def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
         """This function is called 3 times per test: setup, call, teardown."""
         if report.when != "call":
             return
@@ -144,7 +147,7 @@ class InfrahubBackendPlugin:
         # Workaround for https://github.com/opsmill/infrahub/issues/2184
         check.update(do_full_update=True)
 
-    def pytest_sessionfinish(self, session: Session) -> None:  # noqa: ARG002
+    def pytest_sessionfinish(self, session: pytest.Session) -> None:  # noqa: ARG002
         """Set the final RepositoryValidator details after completing the test session."""
         conclusion = "success"
 

infrahub/server.py

@@ -17,7 +17,6 @@ from fastapi.templating import Jinja2Templates
 from infrahub_sdk.exceptions import TimestampFormatError
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.trace import Span
-from pydantic import ValidationError
 from starlette_exporter import PrometheusMiddleware, handle_metrics
 
 from infrahub import __version__, config
@@ -28,7 +27,7 @@ from infrahub.core.graph.index import node_indexes, rel_indexes
 from infrahub.core.initialization import initialization
 from infrahub.database import InfrahubDatabase, InfrahubDatabaseMode, get_db
 from infrahub.dependencies.registry import build_component_registry
-from infrahub.exceptions import Error
+from infrahub.exceptions import Error, ValidationError
 from infrahub.graphql.api.endpoints import router as graphql_router
 from infrahub.lock import initialize_lock
 from infrahub.log import clear_log_context, get_logger, set_log_data

infrahub/services/adapters/event/__init__.py

@@ -31,7 +31,7 @@ class InfrahubEventService:
     async def _send_prefect(self, event: InfrahubEvent) -> None:
         emit_event(
             id=event.meta.id,
-            event=event.get_name(),
+            event=event.event_name,
             resource=event.get_resource(),
             related=event.get_related(),
             payload=event.get_event_payload(),

infrahub/task_manager/event.py

@@ -5,6 +5,7 @@ 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
@@ -91,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", ""),
@@ -109,8 +109,19 @@ 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}
+        return {"attributes": attributes, "relationships": relationships}
 
     def _get_branch_name_from_resource(self) -> str:
         return self.resource.get("infrahub.branch.name") or ""
@@ -220,14 +231,17 @@ class PrefectEvent:
 
         # Retry due to https://github.com/PrefectHQ/prefect/issues/16299
         for _ in range(1, 5):
-            response = await client._client.post("/infrahub/events/filter", json=body)
-            if response.status_code == 200:
+            prefect_error: PrefectHTTPStatusError | None = None
+            try:
+                response = await client._client.post("/infrahub/events/filter", json=body)
                 break
-            await asyncio.sleep(0.1)
+            except PrefectHTTPStatusError as exc:
+                prefect_error = exc
+                await asyncio.sleep(0.1)
 
-        if response.status_code != 200:
+        if prefect_error:
             raise ServiceUnavailableError(
-                message=f"Unable to query prefect due to invalid response from the server (status_code={response.status_code})"
+                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()
 

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,7 +11,7 @@ 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
+    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
@@ -41,6 +39,6 @@ async def define_artifact(
                         "content_type": model.content_type,
                     },
                 )
-                await artifact.save()
+                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