diff --git a/netengine/cli/main.py b/netengine/cli/main.py index 06820bd..c032184 100644 --- a/netengine/cli/main.py +++ b/netengine/cli/main.py @@ -78,6 +78,21 @@ async def _up(spec_file: str, mock: bool, skip_migrations: bool) -> None: orchestrator = Orchestrator(spec, mock_mode=mock) await orchestrator.execute_phases() + # Start background consumers (DNS updates, etc.) + if orchestrator.consumer_supervisor.consumers: + logger.info("All phases complete. Starting background consumers.") + await orchestrator.start_consumers() + # Keep the event loop alive for background consumers + logger.info("Background consumers started. Keeping event loop alive (Ctrl+C to exit).") + try: + await asyncio.sleep(float("inf")) + except KeyboardInterrupt: + logger.info("Interrupted by user") + await orchestrator.consumer_supervisor.stop_all() + logger.info("Consumers stopped") + else: + logger.info("All phases complete. No background consumers to start.") + @cli.command() def status() -> None: diff --git a/netengine/core/consumer_supervisor.py b/netengine/core/consumer_supervisor.py new file mode 100644 index 0000000..bcf4bbb --- /dev/null +++ b/netengine/core/consumer_supervisor.py @@ -0,0 +1,63 @@ +import asyncio +import logging +from typing import Any, Callable, Coroutine, Dict + +logger = logging.getLogger(__name__) + + +class ConsumerSupervisor: + """Manages long-running consumer tasks with automatic restart on failure.""" + + def __init__(self) -> None: + self.tasks: Dict[str, asyncio.Task[Any]] = {} + self.consumers: Dict[str, Callable[[], Coroutine[Any, Any, None]]] = {} + + def register(self, name: str, consumer_coro: Callable[[], Coroutine[Any, Any, None]]) -> None: + """Register a consumer coroutine function.""" + self.consumers[name] = consumer_coro + + async def start_all(self) -> None: + """Start all registered consumers.""" + for name, consumer_func in self.consumers.items(): + await self.start_consumer(name, consumer_func) + + async def start_consumer( + self, name: str, consumer_func: Callable[[], Coroutine[Any, Any, None]] + ) -> None: + """Start a single consumer with automatic restart on failure.""" + + async def supervised_consumer() -> None: + while True: + try: + logger.info(f"Starting consumer: {name}") + await consumer_func() + except asyncio.CancelledError: + logger.info(f"Consumer {name} cancelled") + break + except Exception as e: + logger.error(f"Consumer {name} crashed: {e}. Restarting in 5s...") + await asyncio.sleep(5) + + task: asyncio.Task[None] = asyncio.create_task(supervised_consumer()) + self.tasks[name] = task + logger.info(f"Consumer {name} started") + + async def stop_all(self) -> None: + """Stop all consumers gracefully.""" + for name, task in self.tasks.items(): + logger.info(f"Stopping consumer: {name}") + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + def get_status(self) -> Dict[str, str]: + """Get status of all consumers.""" + status: Dict[str, str] = {} + for name, task in self.tasks.items(): + if task.done(): + status[name] = "crashed" if task.exception() else "completed" + else: + status[name] = "running" + return status diff --git a/netengine/core/orchestrator.py b/netengine/core/orchestrator.py index a54663b..078bfee 100644 --- a/netengine/core/orchestrator.py +++ b/netengine/core/orchestrator.py @@ -4,6 +4,7 @@ from pydantic import ValidationError +from netengine.core.consumer_supervisor import ConsumerSupervisor from netengine.core.state import RuntimeState from netengine.handlers._base import BasePhaseHandler from netengine.handlers.context import PhaseContext @@ -70,12 +71,16 @@ def __init__(self, spec: NetEngineSpec | dict[str, Any], mock_mode: Optional[boo logger.warning(f"Docker unavailable, falling back to mock mode: {exc}") effective_mock = True + # Initialize consumer supervisor for background tasks + self.consumer_supervisor = ConsumerSupervisor() + self.context = PhaseContext( spec=self.spec, runtime_state=self.runtime_state, logger=logger, docker_client=docker_client, mock_mode=effective_mock, + consumer_supervisor=self.consumer_supervisor, ) @staticmethod @@ -136,6 +141,10 @@ async def execute_phases(self, up_to_phase: int = 8) -> None: self.runtime_state.save() raise + async def start_consumers(self) -> None: + """Start all registered consumer tasks.""" + await self.consumer_supervisor.start_all() + def _mark_phase_complete(self, phase_num: int, handler: BasePhaseHandler) -> None: """Record user-facing phase completion for a completed handler. diff --git a/netengine/errors.py b/netengine/errors.py index 543e423..9d69348 100644 --- a/netengine/errors.py +++ b/netengine/errors.py @@ -1,10 +1,17 @@ # TODO: NetEngine exceptions, error handling/reporting/etc module -from .logging import get_logger +from typing import Any + from omegaconf import DictConfig as Config + class BaseNetEngineException(Exception): # TODO: Comprehensive Engine Exception Base - def __init__(self, message: str = "An unknown NetEngine exception occurred.", *args, **kwargs): + def __init__( + self, + message: str = "An unknown NetEngine exception occurred.", + *args: Any, + **kwargs: Any, + ) -> None: """ :param message: @@ -13,15 +20,17 @@ def __init__(self, message: str = "An unknown NetEngine exception occurred.", *a """ self._msg = message self._code: int | str | None = None - self._log_rules: Config | dict = Config({ - "log_on_init": True, - "at_lvl": "TRACE", - "with_msg": self.message, - }) - self._log_xt: Config | dict = Config({}) + self._log_rules: Config | dict[str, Any] = Config( + { + "log_on_init": True, + "at_lvl": "TRACE", + "with_msg": self.message, + } + ) + self._log_xt: Config | dict[str, Any] = Config({}) if kwargs: - for k, v in kwargs: + for k, v in kwargs.items(): self._log_xt.update({k: v}) super().__init__(self.message) @@ -29,4 +38,3 @@ def __init__(self, message: str = "An unknown NetEngine exception occurred.", *a @property def message(self) -> str: return self._msg or "An unknown NetEngine exception occurred." - diff --git a/netengine/handlers/context.py b/netengine/handlers/context.py index 136d894..58cfeec 100644 --- a/netengine/handlers/context.py +++ b/netengine/handlers/context.py @@ -4,11 +4,14 @@ import os from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from netengine.core.state import RuntimeState from netengine.spec.models import NetEngineSpec +if TYPE_CHECKING: + from netengine.core.consumer_supervisor import ConsumerSupervisor + # Default directory for CoreDNS Corefile and zone files. # Overridden by NETENGINE_ZONE_DIR env var. DEFAULT_ZONE_DIR = str(Path.cwd() / "data" / "coredns") @@ -31,6 +34,7 @@ class PhaseContext: kubernetes_client: Any = None supabase_client: Any = None pgmq_client: Any = None + consumer_supervisor: Optional["ConsumerSupervisor"] = None # Phase-specific config phase_name: Optional[str] = None diff --git a/netengine/handlers/oidc_handler.py b/netengine/handlers/oidc_handler.py index a8abf43..914df28 100644 --- a/netengine/handlers/oidc_handler.py +++ b/netengine/handlers/oidc_handler.py @@ -143,6 +143,26 @@ async def create_user( users = await self._admin_request("GET", f"realms/{realm}/users?username={username}") return users[0]["id"] + async def add_token_mapper( + self, + realm: str, + client_id: str, + mapper_name: str, + protocol_mapper_type: str, + config: dict, + ) -> str: + """Add a protocol mapper (token claim) to a client.""" + payload = { + "name": mapper_name, + "protocolMapper": protocol_mapper_type, + "protocol": "openid-connect", + "config": config, + } + await self._admin_request( + "POST", f"realms/{realm}/clients/{client_id}/protocol-mappers/models", json=payload + ) + return mapper_name + def _generate_secret(self) -> str: import secrets diff --git a/netengine/phases/phase_platform_identity.py b/netengine/phases/phase_platform_identity.py index 34f0116..5343fb5 100644 --- a/netengine/phases/phase_platform_identity.py +++ b/netengine/phases/phase_platform_identity.py @@ -83,12 +83,40 @@ async def execute(self, context: PhaseContext) -> None: email="admin@platform.internal", password=admin_password, ) + + # Create platform client for API authentication + client_id = await oidc.create_client( + realm="platform", + client_id="platform-api", + name="Platform API", + redirect_uris=["https://api.platform.internal/callback"], + public=False, + ) + + # Add token mapper to include org claim in JWT + await oidc.add_token_mapper( + realm="platform", + client_id=client_id, + mapper_name="org-claim-mapper", + protocol_mapper_type="oidc-usermodel-property-mapper", + config={ + "user.attribute": "org", + "claim.name": "org", + "jsonType.label": "String", + "id.token.claim": "true", + "access.token.claim": "true", + "userinfo.token.claim": "true", + }, + ) + context.runtime_state.platform_realm_id = realm_id context.runtime_state.admin_user_id = user_id + context.runtime_state.platform_client_id = client_id context.runtime_state.identity_platform_output = { "keycloak_container_id": container_id, "platform_realm_id": realm_id, "admin_user_id": user_id, + "platform_client_id": client_id, "deployed_at": datetime.utcnow().isoformat(), } context.runtime_state.phase_completed["4"] = True diff --git a/netengine/phases/phase_registries.py b/netengine/phases/phase_registries.py index 89991e6..c555f73 100644 --- a/netengine/phases/phase_registries.py +++ b/netengine/phases/phase_registries.py @@ -54,9 +54,14 @@ async def execute(self, context: PhaseContext) -> None: value=tld["listen_ip"], ) - # 5. Wire pgmq consumers (stub – in production, run a loop) - # We'll set up a consumer for DNS updates in a background task. - asyncio.create_task(self._consume_dns_updates(context)) + # 5. Wire pgmq consumers + # Register consumer with supervisor to run after all phases complete + if context.consumer_supervisor: + # Create a wrapper that captures the context + async def dns_consumer(): + await self._consume_dns_updates(context) + + context.consumer_supervisor.register("dns_updates", dns_consumer) # 6. Update state context.runtime_state.world_registry_output = { @@ -92,26 +97,38 @@ async def should_skip(self, context: PhaseContext) -> bool: async def _consume_dns_updates(self, context: PhaseContext): """pgmq consumer: domain.registered -> DNSHandler.add_zone_record.""" + logger = context.logger pgmq = PGMQClient() dns = DNSHandler() while True: - msg = await pgmq.receive("dns_updates") - if not msg: - await asyncio.sleep(1) - continue try: - envelope = EventEnvelope(**json.loads(msg["message"])) - payload = envelope.payload - # Add zone record for the domain (e.g., acme.internal -> IP) - # For MVP, we add an A record pointing to a placeholder or to the AND gateway. - # In real use, the IP would come from the AND allocation. - await dns.add_zone_record( - context=context, - zone=payload["domain"], - record_type="A", - name="@", - value="10.0.0.1", # placeholder – would be replaced with actual IP from AND handler - ) - await pgmq.delete("dns_updates", msg["msg_id"]) + msg = await pgmq.receive("dns_updates") + if not msg: + await asyncio.sleep(1) + continue + + try: + envelope = EventEnvelope(**json.loads(msg["message"])) + payload = envelope.payload + logger.info(f"Processing DNS update for domain: {payload.get('domain')}") + + # Add zone record for the domain (e.g., acme.internal -> IP) + # For MVP, we add an A record pointing to a placeholder or to the AND gateway. + # In real use, the IP would come from the AND allocation. + await dns.add_zone_record( + context=context, + zone=payload["domain"], + record_type="A", + name="@", + value="10.0.0.1", # placeholder – would be replaced with actual IP from AND handler + ) + await pgmq.delete("dns_updates", msg["msg_id"]) + logger.info( + f"Successfully processed DNS update for domain: {payload.get('domain')}" + ) + except Exception as e: + logger.error(f"Error processing DNS update: {e}") + await pgmq.archive_to_dlq("dns_updates", msg["msg_id"], str(e)) except Exception as e: - await pgmq.archive_to_dlq("dns_updates", msg["msg_id"], str(e)) + logger.error(f"Error in DNS update consumer loop: {e}") + await asyncio.sleep(5) # backoff before retrying diff --git a/tests/test_cli.py b/tests/test_cli.py index fe6b778..f4f4de4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -20,6 +20,8 @@ def test_up_invokes_execute_phases_with_example_spec(): with patch("netengine.cli.main.Orchestrator") as mock_orchestrator_class: mock_orchestrator = mock_orchestrator_class.return_value mock_orchestrator.execute_phases = AsyncMock() + # Mock consumer_supervisor as empty (no consumers registered) + mock_orchestrator.consumer_supervisor.consumers = {} result = CliRunner().invoke(cli_main.cli, ["up", str(spec_file)])