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()