wip

Author: tazend

pyslurm/db/account.pxd

@@ -30,6 +30,7 @@ from pyslurm.slurm cimport (
     slurmdb_assoc_cond_t,
     slurmdb_account_cond_t,
     slurmdb_accounts_get,
+    slurmdb_accounts_add,
     slurmdb_destroy_account_rec,
     slurmdb_destroy_account_cond,
     try_xmalloc,
@@ -54,7 +55,7 @@ from pyslurm.xcollections cimport MultiClusterMap
 from pyslurm.utils.uint cimport u16_set_bool_flag
 
 
-cdef class Accounts(MultiClusterMap):
+cdef class Accounts(dict):
     pass
 
 
@@ -80,6 +81,7 @@ cdef class Account:
     cdef public:
         associations
         coordinators
+        association
 
     @staticmethod
     cdef Account from_ptr(slurmdb_account_rec_t *in_ptr)

pyslurm/db/account.pyx

@@ -22,7 +22,7 @@
 # cython: c_string_type=unicode, c_string_encoding=default
 # cython: language_level=3
 
-from pyslurm.core.error import RPCError
+from pyslurm.core.error import RPCError, verify_rpc
 from pyslurm.utils.helpers import (
     instance_to_dict,
     user_to_uid,
@@ -33,14 +33,10 @@ from pyslurm import settings
 from pyslurm import xcollections
 
 
-cdef class Accounts(MultiClusterMap):
+cdef class Accounts(dict):
 
     def __init__(self, accounts=None):
-        super().__init__(data=accounts,
-                         typ="Accounts",
-                         val_type=Account,
-                         id_attr=Account.name,
-                         key_type=str)
+        super().__init__()
 
     @staticmethod
     def load(AccountFilter db_filter=None, Connection db_connection=None):
@@ -64,9 +60,7 @@ cdef class Accounts(MultiClusterMap):
             cond.with_assocs = True
 
         cond._create()
-
         conn = _open_conn_or_error(db_connection)
-
         account_data = SlurmList.wrap(slurmdb_accounts_get(conn.ptr, cond.ptr))
 
         if account_data.is_null:
@@ -79,22 +73,46 @@ cdef class Accounts(MultiClusterMap):
 
         for account_ptr in SlurmList.iter_and_pop(account_data):
             account = Account.from_ptr(<slurmdb_account_rec_t*>account_ptr.data)
-
-            cluster = account.cluster
-            if cluster not in out.data:
-                out.data[cluster] = {}
-            out.data[cluster][account.name] = account
+            out[account.name] = account
 
             assoc_data = SlurmList.wrap(account.ptr.assoc_list, owned=False)
             for assoc_ptr in SlurmList.iter_and_pop(assoc_data):
                 assoc = Association.from_ptr(<slurmdb_assoc_rec_t*>assoc_ptr.data)
                 assoc.qos_data = qos_data
                 assoc.tres_data = tres_data
                 _parse_assoc_ptr(assoc)
-                account.associations.append(assoc)
+
+                if not assoc.user:
+                    # This is the Association of the account itself.
+                    account.association = assoc
+                else:
+                    # These must be User Associations.
+                    account.associations.append(assoc)
 
         return out
 
+    @staticmethod
+    def create(accounts, Connection db_connection=None):
+        cdef:
+            Connection conn
+            Account account
+            SlurmList account_list
+            list assocs_to_add = []
+
+        account_list = SlurmList.create(slurmdb_destroy_account_rec, owned=False)
+
+        for account in accounts:
+            assocs_to_add.extend(account.associations)
+            slurm.slurm_list_append(account_list.info, account.ptr)
+
+        conn = _open_conn_or_error(db_connection)
+        verify_rpc(slurmdb_accounts_add(conn.ptr, account_list.info))
+        Associations.create(assocs_to_add, conn)
+
+        if not db_connection:
+            # Autocommit if no connection was explicitly specified.
+            conn.commit()
+
 
 cdef class AccountFilter:
 
@@ -198,14 +216,26 @@ cdef class Account:
     def name(self):
         return cstr.to_unicode(self.ptr.name)
 
+    @name.setter
+    def name(self, val):
+        cstr.fmalloc(&self.ptr.name, val)
+
     @property
     def description(self):
         return cstr.to_unicode(self.ptr.description)
 
+    @description.setter
+    def description(self, val):
+        cstr.fmalloc(&self.ptr.description, val)
+
     @property
     def organization(self):
         return cstr.to_unicode(self.ptr.organization)
 
+    @organization.setter
+    def organization(self, val):
+        cstr.fmalloc(&self.ptr.organization, val)
+
     @property
     def is_deleted(self):
         if self.ptr.flags & slurm.SLURMDB_ACCT_FLAG_DELETED:

pyslurm/db/assoc.pxd

@@ -31,6 +31,7 @@ from pyslurm.slurm cimport (
     slurmdb_destroy_assoc_cond,
     slurmdb_init_assoc_rec,
     slurmdb_associations_modify,
+    slurmdb_associations_add,
     try_xmalloc,
 )
 from pyslurm.db.util cimport (
@@ -78,6 +79,7 @@ cdef class Association:
         slurmdb_assoc_rec_t *umsg
         QualitiesOfService qos_data
         TrackableResources tres_data
+        owned
 
     cdef public:
         group_tres
@@ -88,7 +90,11 @@ cdef class Association:
         max_tres_per_job
         max_tres_per_node
         qos
+        default_qos
 
     @staticmethod
     cdef Association from_ptr(slurmdb_assoc_rec_t *in_ptr)
 
+
+cdef class AssociationList(SlurmList):
+    pass

pyslurm/db/assoc.pyx

@@ -22,7 +22,7 @@
 # cython: c_string_type=unicode, c_string_encoding=default
 # cython: language_level=3
 
-from pyslurm.core.error import RPCError
+from pyslurm.core.error import RPCError, verify_rpc
 from pyslurm.utils.helpers import (
     instance_to_dict,
     user_to_uid,
@@ -33,6 +33,47 @@ from pyslurm import settings
 from pyslurm import xcollections
 
 
+cdef class AssociationList(SlurmList):
+
+    def __init__(self, owned=True):
+        self.info = slurm.slurm_list_create(slurm.slurmdb_destroy_assoc_rec)
+        self.owned = owned
+
+    def append(self, Association assoc):
+        slurm.slurm_list_append(self.info, assoc.ptr)
+        assoc.owned = False
+        self.cnt = slurm.slurm_list_count(self.info)
+
+    def __iter__(self):
+        return super().__iter__()
+
+    def __next__(self):
+        if self.is_null or self.is_itr_null:
+            raise StopIteration
+
+        if self.itr_cnt < self.cnt:
+            self.itr_cnt += 1
+            assoc = Association.from_ptr(<slurmdb_assoc_rec_t*>slurm.slurm_list_next(self.itr))
+            assoc.owned = False
+            return assoc
+
+        self._dealloc_itr()
+        raise StopIteration
+
+    def extend(self, list_in):
+        for item in list_in:
+            self.append(<Association>item)
+
+
+class AssociationModifyResponse:
+
+    def __init__(self, user=None, account=None, cluster=None, partition=None):
+        self.user = user
+        self.account = account
+        self.cluster = cluster
+        self.partition = partition
+
+
 cdef class Associations(MultiClusterMap):
 
     def __init__(self, assocs=None):
@@ -63,7 +104,7 @@ cdef class Associations(MultiClusterMap):
             conn.ptr, cond.ptr))
 
         if assoc_data.is_null:
-            raise RPCError(msg="Failed to get Association data from slurmdbd")
+            raise RPCError(msg="Failed to get Association data from slurmdbd.")
 
         # Fetch other necessary dependencies needed for translating some
         # attributes (i.e QoS IDs to its name)
@@ -104,16 +145,13 @@ cdef class Associations(MultiClusterMap):
             afilter = <AssociationFilter>db_filter
         afilter._create()
 
-        # Setup DB conn
         conn = _open_conn_or_error(db_connection)
 
         # Any data that isn't parsed yet or needs validation is done in this
         # function.
         _create_assoc_ptr(changes, conn)
 
-        # Modify associations, get the result
-        # This returns a List of char* with the associations that were
-        # modified
+        # Returns a List of char* with the associations that were modified
         response = SlurmList.wrap(slurmdb_associations_modify(
             conn.ptr, afilter.ptr, changes.ptr))
 
@@ -128,7 +166,7 @@ cdef class Associations(MultiClusterMap):
 
         elif not response.is_null:
             # There was no real error, but simply nothing has been modified
-            raise RPCError(msg="Nothing was modified")
+            return None
         else:
             # Autodetects the last slurm error
             raise RPCError()
@@ -139,6 +177,30 @@ cdef class Associations(MultiClusterMap):
 
         return out
 
+    @staticmethod
+    def create(associations, Connection db_connection=None):
+        cdef:
+            Connection conn
+            Association assoc
+            AssociationList assoc_list = AssociationList(owned=False)
+
+        if not associations:
+            return
+
+        for i, assoc in enumerate(associations):
+            # Make sure to remove any duplicate associations, i.e. associations
+            # having the same account name set. For some reason, the slurmdbd
+            # doesn't like that.
+            if assoc not in assoc_list:
+                assoc_list.append(assoc)
+
+        conn = _open_conn_or_error(db_connection)
+        verify_rpc(slurmdb_associations_add(conn.ptr, assoc_list.info))
+
+        if not db_connection:
+            # Autocommit if no connection was explicitly specified.
+            conn.commit()
+
 
 cdef class AssociationFilter:
 
@@ -172,19 +234,21 @@ cdef class AssociationFilter:
         cdef slurmdb_assoc_cond_t *ptr = self.ptr
 
         make_char_list(&ptr.user_list, self.users)
-        make_char_list(&ptr.user_list, self.ids)
+        make_char_list(&ptr.id_list, self.ids)
         make_char_list(&ptr.acct_list, self.accounts)
         make_char_list(&ptr.parent_acct_list, self.parent_accounts)
         make_char_list(&ptr.cluster_list, self.clusters)
         make_char_list(&ptr.partition_list, self.partitions)
-        # TODO: These are QOS ids, not names
+        # TODO: These should be QOS ids, not names
         make_char_list(&ptr.qos_list, self.qos)
+        # TODO: ASSOC_COND_FLAGS
 
 
 cdef class Association:
 
     def __cinit__(self):
         self.ptr = NULL
+        self.owned = True
 
     def __init__(self, **kwargs):
         self._alloc_impl()
@@ -194,7 +258,8 @@ cdef class Association:
             setattr(self, k, v)
 
     def __dealloc__(self):
-        self._dealloc_impl()
+        if self.owned:
+            self._dealloc_impl()
 
     def _dealloc_impl(self):
         slurmdb_destroy_assoc_rec(self.ptr)
@@ -228,7 +293,8 @@ cdef class Association:
 
     def __eq__(self, other):
         if isinstance(other, Association):
-            return self.id == other.id and self.cluster == other.cluster
+#            return self.id == other.id and self.cluster == other.cluster
+            return self.cluster == other.cluster and self.partition == other.partition and self.account == other.account and self.user == other.user
         return NotImplemented
 
     @property
@@ -351,6 +417,10 @@ cdef class Association:
     def parent_account_id(self):
         return u32_parse(self.ptr.parent_id, zero_is_noval=False)
 
+    @property
+    def lineage(self):
+        return cstr.to_unicode(self.ptr.lineage)
+
     @property
     def partition(self):
         return cstr.to_unicode(self.ptr.partition)
@@ -383,6 +453,10 @@ cdef class Association:
     def user(self, val):
         cstr.fmalloc(&self.ptr.user, val)
 
+    @property
+    def user_id(self):
+        return u32_parse(self.ptr.uid, zero_is_noval=False)
+
 
 cdef _parse_assoc_ptr(Association ass):
     cdef:
@@ -397,13 +471,15 @@ cdef _parse_assoc_ptr(Association ass):
             ass.ptr.grp_tres_mins, tres)
     ass.max_tres_mins_per_job = TrackableResourceLimits.from_ids(
             ass.ptr.max_tres_mins_pj, tres)
+    # TODO rename, remove _per_user
     ass.max_tres_run_mins_per_user = TrackableResourceLimits.from_ids(
             ass.ptr.max_tres_run_mins, tres)
     ass.max_tres_per_job = TrackableResourceLimits.from_ids(
             ass.ptr.max_tres_pj, tres)
     ass.max_tres_per_node = TrackableResourceLimits.from_ids(
             ass.ptr.max_tres_pn, tres)
     ass.qos = qos_list_to_pylist(ass.ptr.qos_list, qos)
+    # TODO: default_qos
 
 
 cdef _create_assoc_ptr(Association ass, conn=None):