Skip to content

Commit

Permalink
more spans
Browse files Browse the repository at this point in the history
Signed-off-by: Fatih Acar <[email protected]>
  • Loading branch information
fatih-acar committed Jan 7, 2025
1 parent 62c865e commit 41f865e
Show file tree
Hide file tree
Showing 5 changed files with 34 additions and 11 deletions.
29 changes: 21 additions & 8 deletions backend/infrahub/core/initialization.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import importlib
from uuid import uuid4

from opentelemetry import trace

from infrahub import config, lock
from infrahub.core import registry
from infrahub.core.branch import Branch
Expand Down Expand Up @@ -32,7 +34,6 @@
from infrahub.menu.utils import create_menu_children
from infrahub.permissions import PermissionBackend
from infrahub.storage import InfrahubObjectStorage
from opentelemetry import trace

log = get_logger()

Expand Down Expand Up @@ -183,13 +184,13 @@ async def initialization(db: InfrahubDatabase) -> None:
default_branch = registry.get_branch_from_registry(branch=registry.default_branch)
schema_branch = registry.schema.get_schema_branch(name=default_branch.name)
with trace.get_tracer(__name__).start_as_current_span("GraphQLSchemaManager.get_manager_for_branch"):
gqlm = GraphQLSchemaManager.get_manager_for_branch(branch=default_branch, schema_branch=schema_branch)
gqlm.get_graphql_schema(
include_query=True,
include_mutation=True,
include_subscription=True,
include_types=True,
)
gqlm = GraphQLSchemaManager.get_manager_for_branch(branch=default_branch, schema_branch=schema_branch)
gqlm.get_graphql_schema(
include_query=True,
include_mutation=True,
include_subscription=True,
include_types=True,
)

# ---------------------------------------------------
# Load Default Namespace
Expand All @@ -199,6 +200,7 @@ async def initialization(db: InfrahubDatabase) -> None:
registry.default_ipnamespace = ip_namespace.id


@trace.get_tracer(__name__).start_as_current_span("create_root_node")
async def create_root_node(db: InfrahubDatabase) -> Root:
root = Root(graph_version=GRAPH_VERSION, default_branch=config.SETTINGS.initial.default_branch)
await root.save(db=db)
Expand All @@ -210,6 +212,7 @@ async def create_root_node(db: InfrahubDatabase) -> Root:
return root


@trace.get_tracer(__name__).start_as_current_span("create_default_branch")
async def create_default_branch(db: InfrahubDatabase) -> Branch:
branch = Branch(
name=registry.default_branch,
Expand All @@ -227,6 +230,7 @@ async def create_default_branch(db: InfrahubDatabase) -> Branch:
return branch


@trace.get_tracer(__name__).start_as_current_span("create_global_branch")
async def create_global_branch(db: InfrahubDatabase) -> Branch:
branch = Branch(
name=GLOBAL_BRANCH_NAME,
Expand All @@ -244,6 +248,7 @@ async def create_global_branch(db: InfrahubDatabase) -> Branch:
return branch


@trace.get_tracer(__name__).start_as_current_span("initialization.create_branch")
async def create_branch(
branch_name: str, db: InfrahubDatabase, description: str = "", isolated: bool = True, at: str | None = None
) -> Branch:
Expand Down Expand Up @@ -276,6 +281,7 @@ async def create_branch(
return branch


@trace.get_tracer(__name__).start_as_current_span("create_account")
async def create_account(
db: InfrahubDatabase,
name: str = "admin",
Expand All @@ -296,6 +302,7 @@ async def create_account(
return obj


@trace.get_tracer(__name__).start_as_current_span("create_ipam_namespace")
async def create_ipam_namespace(
db: InfrahubDatabase,
name: str = DEFAULT_IP_NAMESPACE,
Expand All @@ -309,6 +316,7 @@ async def create_ipam_namespace(
return obj


@trace.get_tracer(__name__).start_as_current_span("create_default_menu")
async def create_default_menu(db: InfrahubDatabase) -> None:
for item in default_menu:
obj = await item.to_node(db=db)
Expand All @@ -317,6 +325,7 @@ async def create_default_menu(db: InfrahubDatabase) -> None:
await create_menu_children(db=db, parent=obj, children=item.children)


@trace.get_tracer(__name__).start_as_current_span("create_super_administrator_role")
async def create_super_administrator_role(db: InfrahubDatabase) -> Node:
permission = await Node.init(db=db, schema=InfrahubKind.GLOBALPERMISSION)
await permission.new(
Expand All @@ -337,6 +346,7 @@ async def create_super_administrator_role(db: InfrahubDatabase) -> Node:
return obj


@trace.get_tracer(__name__).start_as_current_span("create_default_roles")
async def create_default_roles(db: InfrahubDatabase) -> Node:
repo_permission = await Node.init(db=db, schema=InfrahubKind.GLOBALPERMISSION)
await repo_permission.new(
Expand Down Expand Up @@ -428,6 +438,7 @@ async def create_default_roles(db: InfrahubDatabase) -> Node:
return role


@trace.get_tracer(__name__).start_as_current_span("create_anonymous_role")
async def create_anonymous_role(db: InfrahubDatabase) -> Node:
deny_permission = await Node.init(db=db, schema=InfrahubKind.OBJECTPERMISSION)
await deny_permission.new(
Expand All @@ -451,6 +462,7 @@ async def create_anonymous_role(db: InfrahubDatabase) -> Node:
return role


@trace.get_tracer(__name__).start_as_current_span("create_super_administrators_group")
async def create_super_administrators_group(
db: InfrahubDatabase, role: Node, admin_accounts: list[CoreAccount]
) -> Node:
Expand All @@ -468,6 +480,7 @@ async def create_super_administrators_group(
return group


@trace.get_tracer(__name__).start_as_current_span("first_time_initialization")
async def first_time_initialization(db: InfrahubDatabase) -> None:
# --------------------------------------------------
# Create the default Branch
Expand Down
1 change: 1 addition & 0 deletions backend/infrahub/core/schema/basenode_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from typing import TYPE_CHECKING, Any, Callable, Iterable, Literal, Optional, Union, overload

from infrahub_sdk.utils import compare_lists, intersection
from opentelemetry import trace
from pydantic import field_validator

from infrahub.core.constants import RelationshipCardinality, RelationshipKind
Expand Down
5 changes: 5 additions & 0 deletions backend/infrahub/core/schema/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

from typing import TYPE_CHECKING, Any, Optional, Union

from opentelemetry import trace

from infrahub import lock
from infrahub.core.manager import NodeManager
from infrahub.core.models import (
Expand Down Expand Up @@ -237,6 +239,7 @@ async def update_schema_to_db(
removed_generics=removed_generics,
)

@trace.get_tracer(__name__).start_as_current_span("load_schema_to_db")
async def load_schema_to_db(
self,
schema: SchemaBranch,
Expand All @@ -259,6 +262,7 @@ async def load_schema_to_db(
node = await self.update_node_in_db(node=item, branch=branch, db=db)
schema.set(name=item_kind, schema=node)

@trace.get_tracer(__name__).start_as_current_span("SchemaManager.load_node_to_db")
async def load_node_to_db(
self,
node: Union[NodeSchema, GenericSchema],
Expand Down Expand Up @@ -307,6 +311,7 @@ async def load_node_to_db(
self.set(name=new_node.kind, schema=new_node, branch=branch.name)
return new_node

@trace.get_tracer(__name__).start_as_current_span("SchemaManager.load_node_to_db")
async def update_node_in_db(
self,
db: InfrahubDatabase,
Expand Down
4 changes: 4 additions & 0 deletions backend/infrahub/core/schema/schema_branch.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from infrahub_sdk.topological_sort import DependencyCycleExistsError, topological_sort
from infrahub_sdk.utils import compare_lists, deep_merge_dict, duplicates, intersection
from opentelemetry import trace
from typing_extensions import Self

from infrahub.computed_attribute.constants import VALID_KINDS as VALID_COMPUTED_ATTRIBUTE_KINDS
Expand Down Expand Up @@ -138,6 +139,7 @@ def get_hash(self) -> str:

return md5hash.hexdigest()

@trace.get_tracer(__name__).start_as_current_span("SchemaBranch.get_hash_full")
def get_hash_full(self) -> SchemaBranchHash:
return SchemaBranchHash(main=self.get_hash(), nodes=self.nodes, generics=self.generics)

Expand Down Expand Up @@ -445,6 +447,7 @@ def generate_fields_for_display_label(self, name: str) -> Optional[dict]:

return fields or None

@trace.get_tracer(__name__).start_as_current_span("SchemaBranch.load_schema")
def load_schema(self, schema: SchemaRoot) -> None:
"""Load a SchemaRoot object and store all NodeSchema or GenericSchema.
Expand Down Expand Up @@ -476,6 +479,7 @@ def load_schema(self, schema: SchemaRoot) -> None:
new_item.update(node_extension)
self.set(name=node_extension.kind, schema=new_item)

@trace.get_tracer(__name__).start_as_current_span("SchemaBranch.process")
def process(self, validate_schema: bool = True) -> None:
self.process_pre_validation()
if validate_schema:
Expand Down
6 changes: 3 additions & 3 deletions backend/infrahub/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,15 +63,15 @@ async def app_initialization(application: FastAPI, enable_scheduler: bool = True
# Initialize database Driver and load local registry
database = application.state.db = InfrahubDatabase(mode=InfrahubDatabaseMode.DRIVER, driver=await get_db())
database.manager.index.init(nodes=node_indexes, rels=rel_indexes)

build_component_registry()

workflow = config.OVERRIDE.workflow or (
WorkflowWorkerExecution()
if config.SETTINGS.workflow.driver == config.WorkflowDriver.WORKER
else WorkflowLocalExecution()
)

message_bus = config.OVERRIDE.message_bus or (
NATSMessageBus() if config.SETTINGS.broker.driver == config.BrokerDriver.NATS else RabbitMQMessageBus()
)
Expand Down

0 comments on commit 41f865e

Please sign in to comment.