Newer
Older
cortex-hub / scripts / register_test_nodes.py
import sys
import os

# Ensure we can import from app
sys.path.append("/app")

from app.db.database import SessionLocal
from app.db.models import AgentNode, NodeGroupAccess

db = SessionLocal()

nodes = [
    {
        "node_id": "test-node-1",
        "display_name": "Test Node 1",
        "description": "Scaled Pod 1",
        "registered_by": "9a333ccd-9c3f-432f-a030-7b1e1284a436",
        "invite_token": "cortex-secret-shared-key",
        "is_active": True,
        "skill_config": {
            "shell": {"enabled": True},
            "browser": {"enabled": True},
            "sync": {"enabled": True}
        },
        "capabilities": {"shell": "v1", "browser": "playwright-sync-bridge"}
    },
    {
        "node_id": "test-node-2",
        "display_name": "Test Node 2",
        "description": "Scaled Pod 2",
        "registered_by": "9a333ccd-9c3f-432f-a030-7b1e1284a436",
        "invite_token": "cortex-secret-shared-key",
        "is_active": True,
        "skill_config": {
            "shell": {"enabled": True},
            "browser": {"enabled": True},
            "sync": {"enabled": True}
        },
        "capabilities": {"shell": "v1", "browser": "playwright-sync-bridge"}
    }
]

group_id = None
# Fetch the user's group to assign access
from app.db.models import User
user = db.query(User).filter(User.id == "9a333ccd-9c3f-432f-a030-7b1e1284a436").first()
if user:
    group_id = user.group_id

try:
    for node_data in nodes:
        # Check if exists
        existing = db.query(AgentNode).filter(AgentNode.node_id == node_data['node_id']).first()
        if existing:
            # Update
            for k, v in node_data.items():
                setattr(existing, k, v)
        else:
            new_node = AgentNode(**node_data)
            db.add(new_node)
            
        db.commit()
        
        # Grant Group Access
        if group_id:
            access = db.query(NodeGroupAccess).filter(
                NodeGroupAccess.node_id == node_data['node_id'],
                NodeGroupAccess.group_id == group_id
            ).first()
            if not access:
                new_access = NodeGroupAccess(
                    node_id=node_data['node_id'],
                    group_id=group_id,
                    access_level="admin",
                    granted_by="9a333ccd-9c3f-432f-a030-7b1e1284a436"
                )
                db.add(new_access)
                db.commit()

    print("Successfully registered test nodes directly via ORM!")
except Exception as e:
    db.rollback()
    print(f"Error registering nodes: {e}")
finally:
    db.close()