diff --git a/classes/assess_type.php b/classes/assess_type.php
index b6da166..0ac8888 100644
--- a/classes/assess_type.php
+++ b/classes/assess_type.php
@@ -117,14 +117,15 @@ public static function is_locked(int $cmid): bool {
     /**
      * Update the assess type.
      *
-     * @param int $cmid - The mod id.
      * @param int $courseid - The course id.
      * @param int $type - formative/summative/dummy.
+     * @param int $cmid - The mod id.
+     * @param int $gradeitemid - The grade item id.
      * @param int $locked - Lock field.
      *
      * @throws \dml_exception
      */
-    public static function update_type(int $cmid, int $courseid, int $type, int $locked = 0): void {
+    public static function update_type(int $courseid, int $type, int $cmid = 0, int $gradeitemid = 0,  int $locked = 0): void {
         global $DB;
         $table = 'local_assess_type';
 
@@ -132,12 +133,13 @@ public static function update_type(int $cmid, int $courseid, int $type, int $loc
         $record = (object) [
           'type' => $type,
           'cmid' => $cmid,
+          'gradeitemid' => $gradeitemid,
           'courseid' => $courseid,
           'locked' => $locked,
         ];
 
         // Check if the record already exists.
-        $existingrecord = $DB->get_record($table, ['cmid' => $cmid], 'id, type, locked');
+        $existingrecord = $DB->get_record($table, ['cmid' => $cmid, 'gradeitemid' => $gradeitemid], 'id, type, locked');
 
         // If the record exists and has changed, update it.
         if ($existingrecord) {
@@ -151,4 +153,22 @@ public static function update_type(int $cmid, int $courseid, int $type, int $loc
             $DB->insert_record($table, $record, false);
         }
     }
+
+    /**
+     * Get all assess type records by course id.
+     *
+     * @param int $courseid
+     * @param int|null $type
+     *
+     * @return array
+     * @throws \dml_exception
+     */
+    public static function get_assess_type_records_by_courseid(int $courseid, ?int $type = null): array {
+        global $DB;
+        $params = ['courseid' => $courseid];
+        if ($type) {
+            $params['type'] = $type;
+        }
+        return $DB->get_records('local_assess_type', $params);
+    }
 }
diff --git a/db/install.xml b/db/install.xml
index a6a57fb..b7b3f2c 100644
--- a/db/install.xml
+++ b/db/install.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8" ?>
-<XMLDB PATH="local/assess_type/db" VERSION="20240701" COMMENT="XMLDB file for UCL"
+<XMLDB PATH="local/assess_type/db" VERSION="20240902" COMMENT="XMLDB file for UCL"
     xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
     xsi:noNamespaceSchemaLocation="../../../lib/xmldb/xmldb.xsd"
 >
@@ -8,7 +8,8 @@
       <FIELDS>
         <FIELD NAME="id" TYPE="int" LENGTH="10" NOTNULL="true" SEQUENCE="true"/>
         <FIELD NAME="courseid" TYPE="int" LENGTH="10" NOTNULL="true" SEQUENCE="false" COMMENT="Foreign key of course id"/>
-        <FIELD NAME="cmid" TYPE="int" LENGTH="10" NOTNULL="true" SEQUENCE="false" COMMENT="Foreign key cm id"/>
+        <FIELD NAME="cmid" TYPE="int" LENGTH="10" NOTNULL="false" SEQUENCE="false" COMMENT="Foreign key cm id"/>
+        <FIELD NAME="gradeitemid" TYPE="int" LENGTH="10" NOTNULL="false" SEQUENCE="false"/>
         <FIELD NAME="type" TYPE="int" LENGTH="1" NOTNULL="false" SEQUENCE="false" COMMENT="The value of this datum"/>
         <FIELD NAME="locked" TYPE="int" LENGTH="1" NOTNULL="false" SEQUENCE="false" COMMENT="Locked field"/>
       </FIELDS>
@@ -17,4 +18,4 @@
       </KEYS>
     </TABLE>
   </TABLES>
-</XMLDB>
\ No newline at end of file
+</XMLDB>
diff --git a/lib.php b/lib.php
index 2e025b3..92fe375 100644
--- a/lib.php
+++ b/lib.php
@@ -121,7 +121,7 @@ function local_assess_type_coursemodule_edit_post_actions($data, $course): stdCl
 
     // We have data, update the assessment type.
     // N.B. Casting (int)assessment_type to stop core Behat error.
-    assess_type::update_type($data->coursemodule, $course->id, (int)$data->assessment_type);
+    assess_type::update_type($course->id, (int)$data->assessment_type, $data->coursemodule);
 
     // Carry on with other form things.
     return $data;