add ability to preview table schema of selected data profile

backend/database/data_profile_manager.py

@@ -1,4 +1,4 @@
-from models.data_profile import DataProfile, DataProfileCreateRequest
+from models.data_profile import DataProfile
 
 
 class DataProfileManager:
@@ -28,17 +28,11 @@ def get_all_data_profile_names_by_org_id(self, org_id):
         data_profile_names = [name for (name,) in result]
         return data_profile_names
 
-    def create_dataprofile(self, data_profile_data: DataProfileCreateRequest):
+    def create_dataprofile(self, data_profile: DataProfile):
         """Create a new DataProfile."""
-        new_data_profile = DataProfile(
-            name=data_profile_data.name,
-            file_type=data_profile_data.file_type,
-            organization_id=data_profile_data.organization_id,
-            extract_instructions=data_profile_data.extract_instructions,
-        )
-        self.session.add(new_data_profile)
+        self.session.add(data_profile)
         self.session.commit()
-        return new_data_profile
+        return data_profile
 
     def get_dataprofile_by_id(self, data_profile_id: int):
         """Retrieve a DataProfile by its ID."""

backend/database/sql_executor.py

@@ -101,7 +101,7 @@ def get_all_table_names_as_list(self) -> List:
             print(f"An error occurred: {e}")
             raise
 
-    def get_table_columns(self, table_name: str) -> List:
+    def get_table_column_names(self, table_name: str) -> List:
         try:
             engine = self.session.bind
             inspector = inspect(engine)

backend/database/table_manager.py

@@ -42,16 +42,30 @@ def _map_table_to_org(
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
+    def _unmap_table_from_org(self, table_name: str):
+        """Unmaps a table from an organization."""
+        try:
+            if self.session:
+                table_map_manager = TableMapManager(self.session)
+                table_map_manager.delete_table_map(table_name)
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+
     def create_table_for_data_profile(
-        self, org_id: int, table_name: str, column_names_and_types: dict
+        self,
+        org_id: int,
+        table_name: str,
+        table_alias: str,
+        column_names_and_types: dict,
     ):
         """Creates a table for a data profile."""
         try:
             executor = SQLExecutor(self.session)
             executor.create_table_for_data_profile(
                 org_id, table_name, column_names_and_types
             )
-            self._map_table_to_org(org_id, table_name)
+            self._map_table_to_org(org_id, table_name, table_alias)
         except Exception as e:
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
@@ -180,6 +194,7 @@ def drop_table(self, table_name: str):
         try:
             executor = SQLExecutor(self.session)
             executor.drop_table(table_name)
+            self._unmap_table_from_org(table_name)
         except Exception as e:
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
@@ -193,11 +208,11 @@ def execute_select_query(self, query: str, format_as_dict: bool = True):
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
-    def get_table_columns(self, table_name: str):
+    def get_table_column_names(self, table_name: str):
         """Returns a list of all of the columns present within the table."""
         try:
             executor = SQLExecutor(self.session)
-            columns = executor.get_table_columns(table_name)
+            columns = executor.get_table_column_names(table_name)
             return columns
         except Exception as e:
             print(f"An error occurred: {e}")

backend/database/table_map_manager.py

@@ -42,6 +42,24 @@ def create_table_map(self, table_map: TableMap):
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
+    def delete_table_map(self, table_name: str):
+        """
+        Delete a table map from the database.
+
+        Args:
+            table_name (str): The name of the table to be deleted.
+        """
+        try:
+            if self.db_session:
+                self.db_session.query(TableMap).filter(
+                    TableMap.table_name == table_name
+                ).delete()
+                self.db_session.commit()
+        except Exception as e:
+            self.db_session.rollback() if self.db_session else None
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+
     def get_org_tables(self, org_id: int) -> List:
         """Returns a list of names of all of the tables associated with an organization."""
         try:

backend/envs/dev/initialization/setup_dev_environment.py

@@ -174,6 +174,7 @@ def create_sample_dataprofile():
         file_type="pdf",
         organization_id=1,
         extract_instructions="Sample extract instructions",
+        table_name="sample_sales",
     )
     # Using DatabaseManager to manage the database session
     with DatabaseManager() as session:

backend/routes/data_profile_routes.py

@@ -72,15 +72,17 @@ async def save_data_profile(
         table_manager.create_table_for_data_profile(
             org_id=current_user.organization_id,
             table_name=table_name,
+            table_alias=request.name,
             column_names_and_types=request.column_names_and_types,
         )
 
+        print(table_name)
         # Create the data profile
         new_data_profile = DataProfile(
             name=request.name,
             extract_instructions=request.extract_instructions,
             organization_id=current_user.organization_id,
-            table_name=table_name,  # TODO: To be further implemented
+            table_name=table_name,
         )
         created_data_profile = data_profile_manager.create_dataprofile(new_data_profile)
 
@@ -175,6 +177,24 @@ async def generate_suggested_column_types(
     return suggested_column_types
 
 
+@data_profile_router.get("/data-profiles/{data_profile_name}/table/column-names/")
+async def get_data_profile_table_column_names(
+    data_profile_name: str, current_user: User = Depends(get_current_user)
+):
+    with DatabaseManager() as session:
+        data_profile_manager = DataProfileManager(session)
+        data_profile = data_profile_manager.get_dataprofile_by_name_and_org(
+            data_profile_name, current_user.organization_id
+        )
+        if data_profile is None:
+            raise HTTPException(status_code=404, detail="Data Profile not found")
+        print("data_profile.table_name", data_profile.table_name)
+        table_manager = TableManager(session)
+        column_names = table_manager.get_table_column_names(data_profile.table_name)
+        print("column_names", column_names)
+        return column_names
+
+
 @data_profile_router.post("/data-profiles/{data_profile_name}/preview/")
 async def preview_data_profile_upload(
     data_profile_name: str,

backend/routes/table_routes.py

@@ -17,12 +17,12 @@ async def get_org_tables(org_id: int, current_user: User = Depends(get_current_u
 
 
 @table_router.get("/table/columns/")
-async def get_table_columns(
+async def get_table_column_names(
     table_name: str, current_user: User = Depends(get_current_user)
 ):
     with DatabaseManager() as session:
         manager = TableManager(session)
-        columns = manager.get_table_columns(table_name)
+        columns = manager.get_table_column_names(table_name)
     return columns
 
 

backend/utils/sql_string_manager.py

@@ -100,7 +100,7 @@ def is_valid_create_table_query(self) -> bool:
         pattern = r"^CREATE TABLE .+;\s*$"
         return bool(re.match(pattern, clean_query))
 
-    def is_valid_pg_table_name(self, table_name) -> bool:
+    def is_valid_table_name(self, table_name) -> bool:
         # Check if the table name matches the allowed pattern
         # Pattern explanation:
         # ^[_a-z]       : Must start with an underscore or a lowercase letter

frontend/src/pages/upload/PreviewTable.jsx

@@ -9,32 +9,29 @@ import {
   TableRow,
 } from "@mui/material";
 
-function PreviewTable({ previewData }) {
-  const data = Array.isArray(previewData) ? previewData : [previewData];
-
-  const generateTableHeaders = (data) => {
-    if (data && data.length > 0) {
-      return Object.keys(data[0]).map((key) => (
-        <TableCell key={key}>{key.replace(/_/g, " ").toUpperCase()}</TableCell>
-      ));
-    }
-    return null;
-  };
-
+function PreviewTable({ columnNames, previewData }) {
   return (
     <TableContainer component={Paper}>
       <Table>
         <TableHead>
-          <TableRow>{generateTableHeaders(data)}</TableRow>
+          <TableRow>
+            {columnNames &&
+              columnNames.map((columnName) => (
+                <TableCell key={columnName}>
+                  {columnName.replace(/_/g, " ").toUpperCase()}
+                </TableCell>
+              ))}
+          </TableRow>
         </TableHead>
         <TableBody>
-          {data.map((row, index) => (
-            <TableRow key={index}>
-              {Object.values(row).map((value, idx) => (
-                <TableCell key={idx}>{value}</TableCell>
-              ))}
-            </TableRow>
-          ))}
+          {previewData &&
+            previewData.map((row, index) => (
+              <TableRow key={index}>
+                {Object.values(row).map((value, idx) => (
+                  <TableCell key={idx}>{value}</TableCell>
+                ))}
+              </TableRow>
+            ))}
         </TableBody>
       </Table>
     </TableContainer>

frontend/src/pages/upload/UploadPage.jsx

@@ -24,6 +24,7 @@ function UploadPage() {
     severity: "info",
   });
   const [showCreateDataProfile, setShowCreateDataProfile] = useState(false);
+  const [columnNames, setColumnNames] = useState([]);
   const [previewData, setPreviewData] = useState(null);
   const [isPreviewLoading, setIsPreviewLoading] = useState(false);
   const [isPreviewTableOpen, setIsPreviewTableOpen] = useState(false);
@@ -37,6 +38,18 @@ function UploadPage() {
       .catch((error) => console.error("Error fetching data profiles:", error));
   }, []);
 
+  useEffect(() => {
+    if (dataProfile) {
+      axios
+        .get(`${API_URL}data-profiles/${dataProfile}/table/column-names/`)
+        .then((response) => {
+          setColumnNames(response.data);
+          setPreviewData(null); // Reset preview data
+        })
+        .catch((error) => console.error("Error fetching column names:", error));
+    }
+  }, [dataProfile]);
+
   const handleCreateDataProfile = (
     name,
     extractInstructions,
@@ -154,7 +167,9 @@ function UploadPage() {
       </Box>
 
       <Box mt={2}>
-        {previewData && <PreviewTable previewData={previewData} />}
+        {(columnNames.length > 0 || previewData) && (
+          <PreviewTable columnNames={columnNames} previewData={previewData} />
+        )}
       </Box>
       <Box display="flex" justifyContent="center" mt={2}>
         {isPreviewLoading && <CircularProgress />}