Refractor code to reflect folder to organization relationship

Signed-off-by: ZhuoweiWen <[email protected]>

back-end/controllers/celery_controller/celery_tasks.py

@@ -182,7 +182,7 @@ def toggle_tiles(self, markers, userid, folderid, polygonfeatures):
     status_code = 0
     session = Session()
     try:
-        user_folder = folder_ops.get_folder_with_id(userid=userid, folderid=folderid, session=session)
+        user_folder = folder_ops.get_folder_with_id(folderid=folderid, session=session)
         if user_folder:
             # Process each polygon feature
             for index, feature in enumerate(polygonfeatures):
@@ -277,9 +277,9 @@ def async_folder_copy_for_export(self, userid, folderid, serialized_csv):
         newfolder_name = f"Exported Filing at {current_datetime}"
 
 
-        csv_name = EXPORT_CSV_NAME_TEMPLATE.format(brand_name=userVal.brand_name, current_datetime=current_datetime)
+        csv_name = EXPORT_CSV_NAME_TEMPLATE.format(brand_name=userVal.organization.brand_name, current_datetime=current_datetime)
 
-        original_folder = folder_ops.get_folder_with_id(userid=userid, folderid=folderid, session=session)
+        original_folder = folder_ops.get_folder_with_id(folderid=folderid, session=session)
         new_folder = original_folder.copy(name=newfolder_name,type='export', deadline=current_datetime, export=True, session=session)
         csv_file = file_ops.create_file(filename=csv_name, content=serialized_csv.encode('utf-8'), folderid=new_folder.id, filetype='export', session=session)
         session.add(csv_file)
@@ -296,7 +296,7 @@ def async_folder_copy_for_import(self, userid, folderid, deadline):
         session = Session()
         newfolder_name = f"Filing for Deadline {deadline}"
 
-        original_folder = folder_ops.get_folder_with_id(userid=userid, folderid=folderid, session=session)
+        original_folder = folder_ops.get_folder_with_id(folderid=folderid, session=session)
         new_folder = original_folder.copy(name=newfolder_name, type='upload', deadline=deadline, export=False, session=session)
         session.commit()
         return new_folder.id
@@ -382,7 +382,8 @@ def preview_fabric_locaiton_coverage(self, data, userid, outfile_name):
         logger.debug("Executing:", gdal_polygonize_cmd)
         subprocess.run(gdal_polygonize_cmd, shell=True)
 
-        folderVal = folder_ops.get_upload_folder(userid, session=session)
+        userVal = user_ops.get_user_with_id(userid=userid)
+        folderVal = folder_ops.get_upload_folder(userVal.organization_id, session=session)
         if folderVal is None:
              return {'error': 'Folder not found'}
 
@@ -434,12 +435,12 @@ def raster2vector(self, data, userid, outfile_name):
         subprocess.run(gdal_polygonize_cmd, shell=True)
 
         userVal = user_ops.get_user_with_id(userid, session=session)
-        folderVal = folder_ops.get_upload_folder(userVal.id, session=session)
+        folderVal = folder_ops.get_upload_folder(userVal.organization_id, session=session)
         if folderVal is None:
             num_folders = folder_ops.get_number_of_folders_for_user(userVal.id, session=session)
             folder_name = f"{userVal.username}-{num_folders + 1}"
             deadline = "September 2024"
-            folderVal = folder_ops.create_folder(folder_name, userVal.id, deadline, 'upload', session=session)
+            folderVal = folder_ops.create_folder(folder_name, userVal.organization_id, deadline, 'upload', session=session)
             session.commit()
 
         vector_file_name = outfile_name + '.kml'

back-end/controllers/database_controller/fabric_ops.py

@@ -52,7 +52,6 @@ def write_to_db(fileid):
             cur.copy_expert("COPY temp_fabric FROM STDIN CSV HEADER DELIMITER ','", output)
             output.seek(0)
 
-             # Insert data from temporary table to final table with user_id
             try:
                 cur.execute(f'INSERT INTO fabric_data SELECT *, {fileid} as file_id FROM temp_fabric;')
                 connection.commit()

back-end/controllers/database_controller/file_ops.py

@@ -1,6 +1,6 @@
 import psycopg2
 from database.sessions import ScopedSession, Session
-from database.models import file, kml_data, file_editfile_link
+from database.models import user, file, kml_data, file_editfile_link
 from threading import Lock
 from datetime import datetime
 from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
@@ -236,4 +236,28 @@ def delete_file(fileid, session=None):
         raise
     finally:
         if owns_session:
-            session.close()
+            session.close()
+
+
+def file_belongs_to_organization(file_id, user_id, session):
+    # Retrieve the user based on user_id
+    userVal = session.query(user).filter(user.id == user_id).first()
+    if not userVal:
+        return False
+
+    # Get the user's organization
+    organization = userVal.organization
+    if not organization:
+        return False
+
+    # Check if the file belongs to this organization
+    fileVal = session.query(file).filter(file.id == file_id).first()
+    if not fileVal:
+        return False
+
+    # Check if the file's folder belongs to the same organization
+    folder = fileVal.folder
+    if not folder:
+        return False
+
+    return folder.user.organization_id == organization.id

back-end/controllers/database_controller/folder_ops.py

@@ -6,85 +6,85 @@
 from sqlalchemy.exc import SQLAlchemyError
 from .user_ops import get_user_with_id
 
-def get_export_folder(userid, folderid=None, session=None):
+def get_export_folder(orgid, folderid=None, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
     try:
-        # Query to retrieve the last folder of type 'upload' for the given user
+        # Query to retrieve the last folder of type 'upload' for the given org
         if not folderid:
             folderVal = (session.query(folder)
-                        .filter(folder.user_id == userid, folder.type == "export")
+                        .filter(folder.organization_id == orgid, folder.type == "export")
                         .order_by(folder.id.desc())
                         .first())
         else:
-            folderVal = session.query(folder).filter(folder.id == folderid, folder.user_id == userid).one()
+            folderVal = session.query(folder).filter(folder.id == folderid).one()
         return folderVal
 
     except NoResultFound:
         return None
     except MultipleResultsFound:
-        return "Multiple results found for the given user ID"
+        return "Multiple results found for the given org ID"
     except Exception as e:
         return str(e)
     finally:
         if owns_session:
             session.close()
 
-def get_upload_folder(userid, folderid=None, session=None):
+def get_upload_folder(orgid, folderid=None, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
     try:
-        # Query to retrieve the last folder of type 'upload' for the given user
+        # Query to retrieve the last folder of type 'upload' for the given org
         if not folderid:
             folderVal = (session.query(folder)
-                        .filter(folder.user_id == userid, folder.type == "upload")
+                        .filter(folder.organization_id == orgid, folder.type == "upload")
                         .order_by(folder.id.desc())
                         .first())
         else:
-            folderVal = session.query(folder).filter(folder.id == folderid, folder.user_id == userid).one()
+            folderVal = session.query(folder).filter(folder.id == folderid).one()
         return folderVal
 
     except NoResultFound:
         return None
     except MultipleResultsFound:
-        return "Multiple results found for the given user ID"
+        return "Multiple results found for the given org ID"
     except Exception as e:
         return str(e)
     finally:
         if owns_session:
             session.close()
 
-def get_folder_with_id(userid, folderid, session=None):
+def get_folder_with_id(folderid, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
     try:
-        folderVal = session.query(folder).filter(folder.id == folderid, folder.user_id == userid).one()
+        folderVal = session.query(folder).filter(folder.id == folderid).one()
         return folderVal
 
     except NoResultFound:
         return None
     except MultipleResultsFound:
-        return "Multiple results found for the given user ID"
+        return "Multiple results found for the given org ID"
     except Exception as e:
         return str(e)
     finally:
         if owns_session:
             session.close()
 
-def get_folders_by_type_for_user(userid, foldertype, session=None):
+def get_folders_by_type_for_org(orgid, foldertype, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
 
     try:
-        folders = session.query(folder).filter(folder.user_id == userid, folder.type == foldertype).all()
+        folders = session.query(folder).filter(folder.organization_id == orgid, folder.type == foldertype).all()
         return folders
 
     except NoResultFound:
@@ -96,14 +96,14 @@ def get_folders_by_type_for_user(userid, foldertype, session=None):
             session.close()
 
 
-def create_folder(foldername, userid, filingDeadline, foldertype, session=None):
+def create_folder(foldername, orgid, filingDeadline, foldertype, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
 
     try:
-        new_folder = folder(name=foldername, user_id=userid, deadline = filingDeadline, type=foldertype)
+        new_folder = folder(name=foldername, organization_id=orgid, deadline=filingDeadline, type=foldertype)
         session.add(new_folder)
         if owns_session:
             session.commit()
@@ -116,13 +116,13 @@ def create_folder(foldername, userid, filingDeadline, foldertype, session=None):
         if owns_session:
             session.close()
 
-def get_number_of_folders_for_user(userid, session=None):
+def get_number_of_folders_for_org(orgid, session=None):
     owns_session = False
     if session is None:
         session = Session()
         owns_session = True
     try:
-        count = session.query(folder).filter(folder.user_id == userid).count()
+        count = session.query(folder).filter(folder.organization_id == orgid).count()
         return count
     except Exception as e:
         return -1
@@ -155,4 +155,24 @@ def delete_folder(folderid, session=None):
         return "An error occurred while deleting the folder: " + str(e)
     finally:
         if owns_session:
-            session.close()
+            session.close()
+
+
+def folder_belongs_to_organization(folder_id, user_id, session):
+
+    # Retrieve the user based on user_id
+    userVal = session.query(user).filter(user.id == user_id).first()
+    if not userVal:
+        return False
+
+    # Get the user's organization
+    organization = userVal.organization
+    if not organization:
+        return False
+
+    # Check if the folder belongs to this organization
+    folderVal = session.query(folder).filter(folder.id == folder_id).first()
+    if not folderVal:
+        return False
+
+    return folderVal.user.organization_id == organization.id

back-end/controllers/database_controller/organization_ops.py

@@ -0,0 +1,6 @@
+from database.models import organization
+from database.sessions import Session
+from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
+from utils.logger_config import logger
+
+

back-end/database/models.py

@@ -28,10 +28,10 @@ class tower(Base):
 
     id = Column(Integer, primary_key=True)
     tower_name = Column(String, nullable=False)
-    user_id = Column(Integer, ForeignKey('user.id'), nullable=False)
+    organization_id = Column(Integer, ForeignKey('organization.id'), nullable=False)
 
     # Relationship to TowerInfo and RasterData models (assuming they exist)
-    user = relationship('user', back_populates='towers')
+    organization = relationship('organization', back_populates='towers')
     tower_info = relationship('towerinfo', back_populates='tower', uselist=False, cascade='all, delete')
     raster_data = relationship('rasterdata', back_populates='tower', uselist=False, cascade='all, delete')
 
@@ -75,16 +75,16 @@ class folder(Base): #filing, will change the name later for less confusion when
     name = Column(String, nullable=False)
     type = Column(String, default='upload') # Currently upload or export
     deadline = Column(Date) #This deadline makes it a filing for the current period
-    user_id = Column(Integer, ForeignKey('user.id', ondelete='CASCADE'))
-    user = relationship('user', back_populates='folders')
+    organization_id = Column(Integer, ForeignKey('organization.id', ondelete='CASCADE'))
+    organization = relationship('organization', back_populates='folders')
     files = relationship('file', back_populates='folder', cascade='all, delete')
     mbtiles = relationship('mbtiles', back_populates='folder', cascade='all, delete')
     editfiles = relationship('editfile', back_populates='folder', cascade='all,delete')
 
     def copy(self, session, export=True, name=None, type=None, deadline=None):
         name = name if name is not None else self.name
         type = type if type is not None else self.type
-        new_folder = folder(name=name, type=type, user_id=self.user_id, deadline = deadline)
+        new_folder = folder(name=name, type=type, organization_id=self.organization_id, deadline = deadline)
         session.add(new_folder)
         session.flush()  # To generate an ID for the new folder