Add site creation subcommand

ucdavis · Dec 18, 2024 · 019750b · 019750b
1 parent 6cfc5c3
commit 019750b
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 11 deletions.
diff --git a/cheeto/cmds/database.py b/cheeto/cmds/database.py
@@ -74,13 +74,22 @@ def site_args(parser: ArgParser,
 
 @site_args.postprocessor()
 def parse_site_arg(args: Namespace):
-    args.site = query_sitename(args.site)
+    if args.site is not None:
+        args.site = query_sitename(args.site)
 
 
-@commands.register('database', 'site', 'add',
+@commands.register('database', 'site', 'new',
                    help='Add a new site')
-def site_add(args: Namespace):
+def site_new(args: Namespace):
     logger = logging.getLogger(__name__)
+    create_site(args.sitename, args.fqdn)
+
+
+@site_new.args()
+def _(parser: ArgParser):
+    parser.add_argument('--sitename', '-s', required=True)
+    parser.add_argument('--fqdn', required=True)
+
 
 
 @site_args.apply()
@@ -345,7 +354,9 @@ def cmd_site_from_puppet(args: Namespace):
             global_record = GlobalUser.from_puppet(user_name, user_record, ssh_key=ssh_key)
             global_record.save()
             global_group = GlobalGroup(groupname=user_name,
-                                       gid=user_record.gid)
+                                       gid=user_record.gid,
+                                       type='user',
+                                       user=global_record)
             global_group.save()
 
             site_record = SiteUser.from_puppet(user_name,

diff --git a/cheeto/database.py b/cheeto/database.py
@@ -18,6 +18,7 @@
 from typing import List, Mapping, Optional, no_type_check, Self, Union
 
 from bson.dbref import DBRef
+from mongoengine.context_managers import run_in_transaction
 import pyescrypt
 from mongoengine import *
 from mongoengine import signals
@@ -1090,6 +1091,14 @@ def connect_to_database(config: MongoConfig, quiet: bool = False):
                    tls=config.tls)
 
 
+def create_site(sitename: str,
+                fqdn: str):
+    if query_site_exists(sitename=sitename):
+        raise ValueError(f'Site {sitename} already exists.')
+    site = Site(sitename=sitename, fqdn=fqdn)
+    site.save()
+
+
 def query_site_exists(sitename: str, raise_exc: bool = False) -> bool:
     logger = logging.getLogger(__name__)
     try:
@@ -1578,13 +1587,21 @@ def add_site_user(sitename: str, user: global_user_t):
     logger = logging.getLogger(__name__)
     if type(user) is str:
         user = GlobalUser.objects.get(username=user)
-    group = SiteGroup.objects.get(groupname=user.username,
-                                  sitename=sitename)
-    site_user = SiteUser(username=user.username,
-                         sitename=sitename,
-                         parent=user,
-                         _groups=[group])
-    site_user.save(force_insert=True)
+    group = GlobalGroup.objects.get(groupname=user.username)
+
+    with run_in_transaction():
+        site_user = SiteUser(username=user.username,
+                             sitename=sitename,
+                             parent=user,
+                             _groups=[group])
+        site_user.save(force_insert=True)
+
+        site_group = SiteGroup(groupname=user.username,
+                               sitename=sitename,
+                               parent=group,
+                               _members=[site_user])
+        site_group.save(force_insert=True)
+
     logger.info(f'Created SiteUser {user.username} on site {sitename}')