feat: load user groups

querybook/server/const/metastore.py

@@ -55,7 +55,9 @@ class DataTable(NamedTuple):
     location: str = None
 
     # Json arrays of partitions
-    partitions: list = []
+    partitions: list[str] = []
+    earliest_partitions: list[str] = None
+    latest_partitions: list[str] = None
 
     # Store the raw info here
     raw_description: str = None

querybook/server/const/user.py

@@ -0,0 +1,10 @@
+from typing import NamedTuple
+
+
+class UserGroup(NamedTuple):
+    name: str
+    display_name: str
+    description: str
+    email: str
+    # list of user names
+    members: list[str]

querybook/server/lib/metastore/base_metastore_loader.py

@@ -353,8 +353,12 @@ def _create_table_table(
             create_table_information(
                 data_table_id=table_id,
                 description=table.description,
-                latest_partitions=json.dumps((table.partitions or [])[-10:]),
-                earliest_partitions=json.dumps((table.partitions or [])[:10]),
+                latest_partitions=json.dumps(
+                    table.latest_partitions or (table.partitions or [])[-10:]
+                ),
+                earliest_partitions=json.dumps(
+                    table.earliest_partitions or (table.partitions or [])[:10]
+                ),
                 hive_metastore_description=table.raw_description,
                 partition_keys=table.partition_keys,
                 custom_properties=table.custom_properties,

querybook/server/logic/user.py

@@ -1,18 +1,14 @@
 from sqlalchemy import func
 
 from app.db import with_session
+from const.elasticsearch import ElasticsearchItem
+from const.user import UserGroup
+from const.user_roles import UserRoleType
 from lib.config import get_config_value
 from lib.logger import get_logger
-from const.user_roles import UserRoleType
-from const.elasticsearch import ElasticsearchItem
-from models.user import (
-    User,
-    UserSetting,
-    UserRole,
-)
+from models.user import User, UserRole, UserSetting, UserGroupMember
 from tasks.sync_elasticsearch import sync_elasticsearch
 
-
 LOG = get_logger(__file__)
 user_settings_config = get_config_value("user_setting")
 
@@ -144,6 +140,66 @@ def delete_user(uid, session=None):
     pass
 
 
+@with_session
+def create_or_update_user_group(user_group: UserGroup, commit=True, session=None):
+    group = get_user_by_name(user_group.name, session=session)
+    fields = {
+        "username": user_group.name,
+        "fullname": user_group.display_name,
+        "email": user_group.email,
+        "is_group": True,
+        "properties": {"description": user_group.description},
+    }
+
+    if not group:
+        # create a new group
+        group = User.create(
+            fields=fields,
+            commit=commit,
+            session=session,
+        )
+    else:
+        # update the group
+        group = User.update(
+            id=group.id,
+            fields=fields,
+            commit=commit,
+            session=session,
+        )
+
+    # get current existing member user ids
+    existing_group_members = (
+        session.query(UserGroupMember).filter(UserGroupMember.gid == group.id).all()
+    )
+    existing_group_member_ids = set([m.uid for m in existing_group_members])
+
+    # get the latest member user ids by name
+    group_members = (
+        session.query(User).filter(User.username.in_(user_group.members)).all()
+    )
+    group_member_ids = set([m.id for m in group_members])
+
+    members_to_delete = list(existing_group_member_ids - group_member_ids)
+    members_to_add = list(group_member_ids - existing_group_member_ids)
+
+    # delete group members not in the group anymore
+    if members_to_delete:
+        session.query(UserGroupMember).filter(UserGroupMember.gid == group.id).filter(
+            UserGroupMember.uid.in_(members_to_delete)
+        ).delete()
+
+    # add new group members
+    if members_to_add:
+        session.add_all(
+            [UserGroupMember(gid=group.id, uid=user_id) for user_id in members_to_add]
+        )
+
+    if commit:
+        session.commit()
+    else:
+        session.flush()
+
+
 """
     ----------------------------------------------------------------------------------------------------------
     USER SETTINGS