infrahub-server 1.3.7__py3-none-any.whl → 1.4.0__py3-none-any.whl

infrahub/graphql/types/event.py

@@ -113,6 +113,65 @@ class BranchDeletedEvent(ObjectType):
     payload = Field(GenericScalar, required=True)
 
 
+# ---------------------------------------
+# Proposed change events
+# ---------------------------------------
+class ProposedChangeReviewEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    reviewer_account_id = String(required=True, description="The ID of the user who reviewed the proposed change")
+    reviewer_account_name = String(required=True, description="The name of the user who reviewed the proposed change")
+    reviewer_decision = String(required=True, description="The decision made by the reviewer")
+    payload = Field(GenericScalar, required=True)
+
+
+class ProposedChangeReviewRevokedEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    reviewer_account_id = String(required=True, description="The ID of the user who reviewed the proposed change")
+    reviewer_account_name = String(required=True, description="The name of the user who reviewed the proposed change")
+    reviewer_former_decision = String(required=True, description="The decision made by the reviewer")
+    payload = Field(GenericScalar, required=True)
+
+
+class ProposedChangeApprovalsRevokedEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    payload = Field(GenericScalar, required=True)
+
+
+class ProposedChangeReviewRequestedEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    requested_by_account_id = String(
+        required=True, description="The ID of the user who requested the proposed change to be reviewed"
+    )
+    requested_by_account_name = String(
+        required=True, description="The name of the user who requested the proposed change to be reviewed"
+    )
+    payload = Field(GenericScalar, required=True)
+
+
+class ProposedChangeMergedEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    merged_by_account_id = String(required=True, description="The ID of the user who merged the proposed change")
+    merged_by_account_name = String(required=True, description="The name of the user who merged the proposed change")
+    payload = Field(GenericScalar, required=True)
+
+
+class ProposedChangeThreadEvent(ObjectType):
+    class Meta:
+        interfaces = (EventNodeInterface,)
+
+    payload = Field(GenericScalar, required=True)
+
+
 # ---------------------------------------
 # Node/Object events
 # ---------------------------------------
@@ -163,5 +222,14 @@ EVENT_TYPES: dict[str, type[ObjectType]] = {
     events.BranchDeletedEvent.event_name: BranchDeletedEvent,
     events.GroupMemberAddedEvent.event_name: GroupEvent,
     events.GroupMemberRemovedEvent.event_name: GroupEvent,
+    events.ProposedChangeApprovedEvent.event_name: ProposedChangeReviewEvent,
+    events.ProposedChangeApprovalRevokedEvent.event_name: ProposedChangeReviewRevokedEvent,
+    events.ProposedChangeRejectedEvent.event_name: ProposedChangeReviewEvent,
+    events.ProposedChangeRejectionRevokedEvent.event_name: ProposedChangeReviewRevokedEvent,
+    events.ProposedChangeReviewRequestedEvent.event_name: ProposedChangeReviewRequestedEvent,
+    events.ProposedChangeApprovalsRevokedEvent.event_name: ProposedChangeApprovalsRevokedEvent,
+    events.ProposedChangeMergedEvent.event_name: ProposedChangeMergedEvent,
+    events.ProposedChangeThreadCreatedEvent.event_name: ProposedChangeThreadEvent,
+    events.ProposedChangeThreadUpdatedEvent.event_name: ProposedChangeThreadEvent,
     "undefined": StandardEvent,
 }

infrahub/groups/tasks.py

@@ -4,14 +4,16 @@ from prefect import flow
 
 from infrahub.core.constants import InfrahubKind
 from infrahub.groups.models import RequestGraphQLQueryGroupUpdate
-from infrahub.services import InfrahubServices
+from infrahub.workers.dependencies import get_client
 from infrahub.workflows.utils import add_tags
 
 
 @flow(name="graphql-query-group-update", flow_run_name="Update GraphQLQuery Group '{model.query_name}'")
-async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate, service: InfrahubServices) -> None:
+async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate) -> None:
     """Create or Update a GraphQLQueryGroup."""
 
+    client = get_client()
+
     # If there is only one subscriber, associate the task to it
     # If there are more than one, for now we can't associate all of them
     related_nodes = []
@@ -23,7 +25,7 @@ async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate, serv
     params_hash = dict_hash(model.params)
     group_name = f"{model.query_name}__{params_hash}"
     group_label = f"Query {model.query_name} Hash({params_hash[:8]})"
-    group = await service.client.create(
+    group = await client.create(
         kind=InfrahubKind.GRAPHQLQUERYGROUP,
         branch=model.branch,
         name=group_name,
@@ -36,6 +38,4 @@ async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate, serv
     await group.save(allow_upsert=True)
 
     if model.subscribers:
-        await group_add_subscriber(
-            client=service.client, group=group, subscribers=model.subscribers, branch=model.branch
-        )
+        await group_add_subscriber(client=client, group=group, subscribers=model.subscribers, branch=model.branch)

infrahub/lock.py

@@ -12,6 +12,7 @@ from prometheus_client import Histogram
 from redis.asyncio.lock import Lock as GlobalLock
 
 from infrahub import config
+from infrahub.core.timestamp import current_timestamp
 
 if TYPE_CHECKING:
     from types import TracebackType
@@ -91,7 +92,7 @@ class NATSLock:
         await self.release()
 
     async def acquire(self) -> None:
-        token = uuid.uuid1().hex
+        token = current_timestamp()
         while True:
             if await self.do_acquire(token):
                 self.token = token
@@ -155,7 +156,7 @@ class InfrahubLock:
     async def acquire(self) -> None:
         with LOCK_ACQUIRE_TIME_METRICS.labels(self.name, self.lock_type).time():
             if not self.use_local:
-                await self.remote.acquire()
+                await self.remote.acquire(token=current_timestamp())
             else:
                 await self.local.acquire()
         self.acquire_time = time.time_ns()

infrahub/menu/generator.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from infrahub.core import registry
+from infrahub.core.constants import InfrahubKind
 from infrahub.core.protocols import CoreMenuItem
 from infrahub.log import get_logger
 
@@ -133,4 +134,11 @@ async def generate_menu(db: InfrahubDatabase, branch: Branch, menu_items: list[C
         default_menu.children[str(menu_item.identifier)] = menu_item
         items_to_add[item_name] = True
 
+    builtin_ipaddress = registry.schema.get_generic_schema(name=InfrahubKind.IPADDRESS, branch=branch, duplicate=False)
+    builtin_ipprefix = registry.schema.get_generic_schema(name=InfrahubKind.IPPREFIX, branch=branch, duplicate=False)
+    ipam_missing = len(builtin_ipaddress.used_by + builtin_ipprefix.used_by) == 0
+
+    if ipam_missing:
+        structure.data.pop("BuiltinIPAM")
+
     return structure

infrahub/message_bus/operations/__init__.py

@@ -1,14 +1,11 @@
 import ujson
 from prefect import Flow
 
+from infrahub.log import get_logger
 from infrahub.message_bus import RPCErrorResponse, messages
-from infrahub.message_bus.operations import (
-    git,
-    refresh,
-    send,
-)
+from infrahub.message_bus.operations import git, refresh, send
 from infrahub.message_bus.types import MessageTTL
-from infrahub.services import InfrahubServices
+from infrahub.services.adapters.message_bus import InfrahubMessageBus
 from infrahub.tasks.check import set_check_status
 
 COMMAND_MAP = {
@@ -22,7 +19,7 @@ COMMAND_MAP = {
 
 
 async def execute_message(
-    routing_key: str, message_body: bytes, service: InfrahubServices, skip_flow: bool = False
+    routing_key: str, message_body: bytes, message_bus: InfrahubMessageBus, skip_flow: bool = False
 ) -> MessageTTL | None:
     message_data = ujson.loads(message_body)
     message = messages.MESSAGE_MAP[routing_key](**message_data)
@@ -31,16 +28,16 @@ async def execute_message(
         func = COMMAND_MAP[routing_key]
         if skip_flow and isinstance(func, Flow):
             func = func.fn
-        await func(message=message, service=service)
+        await func(message=message)
     except Exception as exc:
         if message.reply_requested:
             response = RPCErrorResponse(errors=[str(exc)], initial_message=message.model_dump())
-            await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
+            await message_bus.reply_if_initiator_meta(message=response, initiator=message)
             return None
         if message.reached_max_retries:
-            service.log.exception("Message failed after maximum number of retries", error=exc)
-            await set_check_status(message, conclusion="failure", service=service)
+            get_logger().exception("Message failed after maximum number of retries", error=exc)
+            await set_check_status(message, conclusion="failure")
             return None
         message.increase_retry_count()
-        await service.message_bus.send(message, delay=MessageTTL.FIVE, is_retry=True)
+        await message_bus.send(message, delay=MessageTTL.FIVE, is_retry=True)
         return MessageTTL.FIVE

infrahub/message_bus/operations/git/file.py

@@ -6,29 +6,30 @@ from infrahub.message_bus.messages.git_file_get import (
     GitFileGetResponse,
     GitFileGetResponseData,
 )
-from infrahub.services import InfrahubServices
+from infrahub.workers.dependencies import get_client, get_message_bus
 
 log = get_logger()
 
 
-async def get(message: messages.GitFileGet, service: InfrahubServices) -> None:
+async def get(message: messages.GitFileGet) -> None:
     log.info("Collecting file from repository", repository=message.repository_name, file=message.file)
 
     repo = await get_initialized_repo(
+        client=get_client(),
         repository_id=message.repository_id,
         name=message.repository_name,
-        service=service,
         repository_kind=message.repository_kind,
         commit=message.commit,
     )
 
+    message_bus = await get_message_bus()
     try:
         content = await repo.get_file(commit=message.commit, location=message.file)
     except (FileOutOfRepositoryError, RepositoryFileNotFoundError) as e:
         if message.reply_requested:
             response = GitFileGetResponse(data=GitFileGetResponseData(error_message=e.message, http_code=e.HTTP_CODE))
-            await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
+            await message_bus.reply_if_initiator_meta(message=response, initiator=message)
     else:
         if message.reply_requested:
             response = GitFileGetResponse(data=GitFileGetResponseData(content=content))
-            await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
+            await message_bus.reply_if_initiator_meta(message=response, initiator=message)

infrahub/message_bus/operations/git/repository.py

@@ -1,21 +1,21 @@
 from prefect import flow
 
 from infrahub.exceptions import RepositoryError
-from infrahub.git.repository import InfrahubRepository, get_initialized_repo, initialize_repo
+from infrahub.git.repository import InfrahubRepository, get_initialized_repo
 from infrahub.log import get_logger
 from infrahub.message_bus import messages
 from infrahub.message_bus.messages.git_repository_connectivity import (
     GitRepositoryConnectivityResponse,
     GitRepositoryConnectivityResponseData,
 )
-from infrahub.services import InfrahubServices
 from infrahub.worker import WORKER_IDENTITY
+from infrahub.workers.dependencies import get_client, get_message_bus
 
 log = get_logger()
 
 
 @flow(name="git-repository-check-connectivity", flow_run_name="Check connectivity for {message.repository_name}")
-async def connectivity(message: messages.GitRepositoryConnectivity, service: InfrahubServices) -> None:
+async def connectivity(message: messages.GitRepositoryConnectivity) -> None:
     response_data = GitRepositoryConnectivityResponseData(message="Successfully accessed repository", success=True)
 
     try:
@@ -28,30 +28,22 @@ async def connectivity(message: messages.GitRepositoryConnectivity, service: Inf
         response = GitRepositoryConnectivityResponse(
             data=response_data,
         )
-        await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
+        message_bus = await get_message_bus()
+        await message_bus.reply_if_initiator_meta(message=response, initiator=message)
 
 
 @flow(name="refresh-git-fetch", flow_run_name="Fetch git repository {message.repository_name} on " + WORKER_IDENTITY)
-async def fetch(message: messages.RefreshGitFetch, service: InfrahubServices) -> None:
+async def fetch(message: messages.RefreshGitFetch) -> None:
     if message.meta and message.meta.initiator_id == WORKER_IDENTITY:
         log.info("Ignoring git fetch request originating from self", worker=WORKER_IDENTITY)
         return
 
-    try:
-        repo = await get_initialized_repo(
-            repository_id=message.repository_id,
-            name=message.repository_name,
-            service=service,
-            repository_kind=message.repository_kind,
-        )
-    except RepositoryError:
-        repo = await initialize_repo(
-            location=message.location,
-            repository_id=message.repository_id,
-            name=message.repository_name,
-            service=service,
-            repository_kind=message.repository_kind,
-        )
+    repo = await get_initialized_repo(
+        client=get_client(),
+        repository_id=message.repository_id,
+        name=message.repository_name,
+        repository_kind=message.repository_kind,
+    )
 
     await repo.fetch()
     await repo.pull(

infrahub/message_bus/operations/refresh/registry.py

@@ -1,28 +1,34 @@
+from infrahub.log import get_logger
 from infrahub.message_bus import messages
-from infrahub.services import InfrahubServices
 from infrahub.tasks.registry import refresh_branches
 from infrahub.worker import WORKER_IDENTITY
+from infrahub.workers.dependencies import get_component, get_database
 
 
-async def branches(message: messages.RefreshRegistryBranches, service: InfrahubServices) -> None:
+async def branches(message: messages.RefreshRegistryBranches) -> None:
     if message.meta and message.meta.initiator_id == WORKER_IDENTITY:
-        service.log.info("Ignoring refresh registry refresh request originating from self", worker=WORKER_IDENTITY)
+        get_logger().info("Ignoring refresh registry refresh request originating from self", worker=WORKER_IDENTITY)
         return
 
-    async with service.database.start_session(read_only=True) as db:
+    database = await get_database()
+    async with database.start_session(read_only=False) as db:
         await refresh_branches(db=db)
 
-    await service.component.refresh_schema_hash()
+    component = await get_component()
+    await component.refresh_schema_hash()
 
 
-async def rebased_branch(message: messages.RefreshRegistryRebasedBranch, service: InfrahubServices) -> None:
+async def rebased_branch(message: messages.RefreshRegistryRebasedBranch) -> None:
     if message.meta and message.meta.initiator_id == WORKER_IDENTITY:
-        service.log.info(
+        get_logger().info(
             "Ignoring refresh registry refreshed branch for request originating from self", worker=WORKER_IDENTITY
         )
         return
 
-    async with service.database.start_session(read_only=True) as db:
+    database = await get_database()
+
+    async with database.start_session(read_only=True) as db:
         await refresh_branches(db=db)
 
-    await service.component.refresh_schema_hash()
+    component = await get_component()
+    await component.refresh_schema_hash()

infrahub/message_bus/operations/send/echo.py

@@ -1,13 +1,16 @@
 from prefect import flow
 
+from infrahub.log import get_logger
 from infrahub.message_bus import messages
 from infrahub.message_bus.messages.send_echo_request import SendEchoRequestResponse, SendEchoRequestResponseData
-from infrahub.services import InfrahubServices
+from infrahub.workers.dependencies import get_message_bus
 
 
 @flow(name="echo-request")
-async def request(message: messages.SendEchoRequest, service: InfrahubServices) -> None:
-    service.log.info(f"Received message: {message.message}")
+async def request(message: messages.SendEchoRequest) -> None:
+    get_logger().info(f"Received message: {message.message}")
+
     if message.reply_requested:
         response = SendEchoRequestResponse(data=SendEchoRequestResponseData(response=f"Reply to: {message.message}"))
-        await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
+        message_bus = await get_message_bus()
+        await message_bus.reply_if_initiator_meta(message=response, initiator=message)

infrahub/message_bus/types.py

@@ -31,6 +31,7 @@ class KVTTL(int, Enum):
     ONE = 1
     TEN = 10
     FIFTEEN = 15
+    ONE_MINUTE = 60
     TWO_HOURS = 7200
 
     @classmethod

infrahub/permissions/__init__.py

@@ -1,5 +1,5 @@
 from infrahub.permissions.backend import PermissionBackend
-from infrahub.permissions.globals import define_global_permission_from_branch
+from infrahub.permissions.globals import define_global_permission_from_branch, get_or_create_global_permission
 from infrahub.permissions.local_backend import LocalPermissionBackend
 from infrahub.permissions.manager import PermissionManager
 from infrahub.permissions.report import report_schema_permissions
@@ -12,5 +12,6 @@ __all__ = [
     "PermissionManager",
     "define_global_permission_from_branch",
     "get_global_permission_for_kind",
+    "get_or_create_global_permission",
     "report_schema_permissions",
 ]

infrahub/permissions/constants.py

@@ -25,8 +25,21 @@ GLOBAL_PERMISSION_DENIAL_MESSAGE = {
     GlobalPermissions.EDIT_DEFAULT_BRANCH.value: "You are not allowed to change data in the default branch",
     GlobalPermissions.MERGE_BRANCH.value: "You are not allowed to merge a branch",
     GlobalPermissions.MERGE_PROPOSED_CHANGE.value: "You are not allowed to merge proposed changes",
+    GlobalPermissions.REVIEW_PROPOSED_CHANGE.value: "You are not allowed to review proposed changes",
     GlobalPermissions.MANAGE_SCHEMA.value: "You are not allowed to manage the schema",
     GlobalPermissions.MANAGE_ACCOUNTS.value: "You are not allowed to manage user accounts, groups or roles",
     GlobalPermissions.MANAGE_PERMISSIONS.value: "You are not allowed to manage permissions",
     GlobalPermissions.MANAGE_REPOSITORIES.value: "You are not allowed to manage repositories",
 }
+
+GLOBAL_PERMISSION_DESCRIPTION = {
+    GlobalPermissions.EDIT_DEFAULT_BRANCH: "Allow a user to change data in the default branch",
+    GlobalPermissions.MERGE_BRANCH: "Allow a user to merge branches",
+    GlobalPermissions.MERGE_PROPOSED_CHANGE: "Allow a user to merge proposed changes",
+    GlobalPermissions.REVIEW_PROPOSED_CHANGE: "Allow a user to approve or reject proposed changes",
+    GlobalPermissions.MANAGE_SCHEMA: "Allow a user to manage the schema",
+    GlobalPermissions.MANAGE_ACCOUNTS: "Allow a user to manage accounts, account roles and account groups",
+    GlobalPermissions.MANAGE_PERMISSIONS: "Allow a user to manage permissions",
+    GlobalPermissions.MANAGE_REPOSITORIES: "Allow a user to manage repositories",
+    GlobalPermissions.SUPER_ADMIN: "Allow a user to do anything",
+}

infrahub/permissions/globals.py

@@ -1,7 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from infrahub.core.account import GlobalPermission
 from infrahub.core.constants import GLOBAL_BRANCH_NAME, GlobalPermissions, PermissionDecision
+from infrahub.core.manager import NodeManager
+from infrahub.core.node import Node
+from infrahub.core.protocols import CoreGlobalPermission
 from infrahub.core.registry import registry
 
+from .constants import GLOBAL_PERMISSION_DESCRIPTION
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
 
 def define_global_permission_from_branch(permission: GlobalPermissions, branch_name: str) -> GlobalPermission:
     if branch_name in (GLOBAL_BRANCH_NAME, registry.default_branch):
@@ -9,7 +21,24 @@ def define_global_permission_from_branch(permission: GlobalPermissions, branch_n
     else:
         decision = PermissionDecision.ALLOW_OTHER
 
-    return GlobalPermission(
+    return GlobalPermission(action=permission.value, decision=decision.value)
+
+
+async def get_or_create_global_permission(db: InfrahubDatabase, permission: GlobalPermissions) -> CoreGlobalPermission:
+    permissions = await NodeManager.query(
+        db=db, schema=CoreGlobalPermission, filters={"action__value": permission.value}, limit=1
+    )
+
+    if permissions:
+        return permissions[0]
+
+    p = await Node.init(db=db, schema=CoreGlobalPermission)
+    await p.new(
+        db=db,
         action=permission.value,
-        decision=decision.value,
+        decision=PermissionDecision.ALLOW_ALL.value,
+        description=GLOBAL_PERMISSION_DESCRIPTION[permission],
     )
+    await p.save(db=db)
+
+    return p

infrahub/permissions/manager.py

@@ -48,6 +48,13 @@ class PermissionManager:
             specificity += 1
         return specificity
 
+    def is_super_admin(self) -> bool:
+        return self.resolve_global_permission(
+            permission_to_check=GlobalPermission(
+                action=GlobalPermissions.SUPER_ADMIN, decision=PermissionDecision.ALLOW_ALL
+            ),
+        )
+
     def report_object_permission(self, namespace: str, name: str, action: str) -> PermissionDecisionFlag:
         """Given a set of permissions, return the permission decision for a given kind and action."""
         highest_specificity: int = -1
@@ -94,11 +101,7 @@ class PermissionManager:
 
     def has_permission(self, permission: GlobalPermission | ObjectPermission) -> bool:
         """Tell if a permission is granted given the permissions loaded in memory."""
-        is_super_admin = self.resolve_global_permission(
-            permission_to_check=GlobalPermission(
-                action=GlobalPermissions.SUPER_ADMIN, decision=PermissionDecision.ALLOW_ALL
-            ),
-        )
+        is_super_admin = self.is_super_admin()
 
         if isinstance(permission, GlobalPermission):
             return self.resolve_global_permission(permission_to_check=permission) or is_super_admin

infrahub/pools/prefix.py

@@ -9,7 +9,9 @@ if TYPE_CHECKING:
     from infrahub.core.ipam.constants import IPNetworkType
 
 
-def get_next_available_prefix(pool: IPSet, prefix_length: int, prefix_ver: Literal[4, 6] = 4) -> IPNetworkType:
+def get_next_available_prefix(
+    pool: IPSet, prefix_length: int | None = None, prefix_ver: Literal[4, 6] = 4
+) -> IPNetworkType:
     """Get the next available prefix of a given prefix length from an IPSet.
 
     Args:
@@ -20,10 +22,7 @@ def get_next_available_prefix(pool: IPSet, prefix_length: int, prefix_ver: Liter
     Raises:
         ValueError: If there are no available subnets in the pool
     """
-    prefix_ver_map = {
-        4: ipaddress.IPv4Network,
-        6: ipaddress.IPv6Network,
-    }
+    prefix_ver_map = {4: ipaddress.IPv4Network, 6: ipaddress.IPv6Network}
 
     filtered_pool = IPSet([])
     for subnet in pool.iter_cidrs():
@@ -31,6 +30,9 @@ def get_next_available_prefix(pool: IPSet, prefix_length: int, prefix_ver: Liter
             filtered_pool.add(subnet)
 
     for cidr in filtered_pool.iter_cidrs():
+        if prefix_length is None:
+            return cidr
+
         if cidr.prefixlen <= prefix_length:
             next_available = ipaddress.ip_network(f"{cidr.network}/{prefix_length}")
             return next_available

infrahub/prefect_server/app.py

@@ -1,16 +1,47 @@
 from __future__ import annotations
 
+import asyncio
+import os
+
 from fastapi import APIRouter, FastAPI
 from prefect.server.api.server import create_app
 
 from . import events
+from .bootstrap import init_prefect
+
+GLOBAL_TASKMGR_INIT_LOCK = "global.taskmgr.init"
 
 router = APIRouter(prefix="/infrahub")
 
 router.include_router(events.router)
 
 
+async def _init_prefect() -> None:
+    # Import there in case we are running Prefect within a testsuite using the original Prefect container
+    from infrahub import lock
+    from infrahub.lock import initialize_lock
+    from infrahub.services import InfrahubServices
+    from infrahub.workers.dependencies import get_cache
+
+    cache = await get_cache()
+    service = await InfrahubServices.new(cache=cache)
+    initialize_lock(service=service)
+
+    async with lock.registry.get(name=GLOBAL_TASKMGR_INIT_LOCK):
+        await init_prefect()
+
+
 def create_infrahub_prefect() -> FastAPI:
+    if (
+        os.getenv("PREFECT_API_BLOCKS_REGISTER_ON_START") == "false"
+        and os.getenv("PREFECT_API_DATABASE_MIGRATE_ON_START") == "false"
+    ):
+        # We are probably running distributed mode
+        from infrahub import config
+
+        config.SETTINGS.initialize_and_exit()
+        asyncio.run(_init_prefect())
+
     app = create_app()
     api_app: FastAPI = app.__dict__["api_app"]
     api_app.include_router(router=router)

infrahub/prefect_server/bootstrap.py

@@ -0,0 +1,18 @@
+import asyncio
+
+from prefect.server.database import provide_database_interface
+from prefect.server.models.block_registration import run_block_auto_registration
+
+
+async def init_prefect() -> None:
+    db = provide_database_interface()
+
+    await db.create_db()
+    session = await db.session()
+
+    async with session:
+        await run_block_auto_registration(session=session)
+
+
+if __name__ == "__main__":
+    asyncio.run(init_prefect())