fixing the company updates

src/scheduler/database/employee_repository.py

@@ -4,14 +4,12 @@ class EmployeeRepository:
 
     def list(self):
         with self.connection.cursor() as cursor:
-            cursor.execute(
-                """
+            cursor.execute("""
                 SELECT users.id, company_users.company_id, users.name, users.last_name, users.preference
                 FROM users
                 INNER JOIN company_users ON company_users.user_id = users.id
                 ORDER BY users.id
-                """
-            )
+                """)
             return cursor.fetchall()
 
     def create(self, company_id: int, name: str, last_name: str, preference: str):
@@ -72,7 +70,13 @@ class EmployeeRepository:
         self.connection.commit()
         return self.get(employee_id)
 
-    def update(self, employee_id: int, name: str | None, last_name: str | None, preference: str | None):
+    def update(
+        self,
+        employee_id: int,
+        name: str | None,
+        last_name: str | None,
+        preference: str | None,
+    ):
         updates: list[str] = []
         params: list[object] = []
         if name is not None:
@@ -101,7 +105,9 @@ class EmployeeRepository:
 
     def update_company(self, employee_id: int, company_id: int):
         with self.connection.cursor() as cursor:
-            cursor.execute("DELETE FROM company_users WHERE user_id = %s", (employee_id,))
+            cursor.execute(
+                "DELETE FROM company_users WHERE user_id = %s", (employee_id,)
+            )
             cursor.execute(
                 """
                 INSERT INTO company_users (user_id, company_id)
@@ -114,17 +120,22 @@ class EmployeeRepository:
         return self.get(employee_id)
 
     def assign_to_company(self, company_id: int, employee_id: int):
+        employee = self.get(employee_id)
         with self.connection.cursor() as cursor:
             cursor.execute(
                 """
                 INSERT INTO company_users (user_id, company_id)
                 VALUES (%s, %s)
-                ON CONFLICT (user_id, company_id) DO UPDATE SET date_updated = NOW()
+                ON CONFLICT (user_id)
+                DO UPDATE SET
+                    company_id = EXCLUDED.company_id,
+                    date_updated = NOW()
                 RETURNING id, user_id, company_id, date_created, date_updated
                 """,
                 (employee_id, company_id),
             )
             company_user = cursor.fetchone()
+
         self.connection.commit()
         return company_user