Show only tables pertaining to the user's organization

backend/databases/sql_executor.py

@@ -14,6 +14,7 @@ def append_df_to_table(self, df: pd.DataFrame, table_name: str):
         try:
             df.to_sql(table_name, self.session.bind, if_exists="append", index=False)
         except Exception as e:
+            self.session.rollback()
             print(
                 f"An error occurred while appending data to table {table_name}: {str(e)}"
             )

backend/databases/table_manager.py

@@ -1,25 +1,50 @@
+from sqlalchemy.orm import Session
+from typing import List, Optional
+
+import pandas as pd
+
 from fastapi import HTTPException
-from databases.database_manager import DatabaseManager
 from databases.sql_executor import SQLExecutor
 from databases.table_metadata_manager import TableMetadataManager
 from llms.base import BaseLLM
+from models.organization_table_map import OrganizationTableMap
 from utils.sql_string_manipulator import SQLStringManipulator
 
-import pandas as pd
-
 
 class TableManager:
     """
     Manages table operations.
     Has functions that integrate a Large Language Model (LLM) and SQLExecutor.
 
     Attributes:
+        session (Optional[str]): The session used for database operations.
         llm (BaseLLM): Instance of a Large Language Model for SQL operations.
     """
 
-    def __init__(self, llm: BaseLLM = None):
+    def __init__(self, session: Optional[Session] = None, llm: BaseLLM = None):
+        self.session = session
         self.llm = llm
 
+    def _map_table_to_org(
+        self, org_id: int, table_name: str, alias: Optional[str] = None
+    ):
+        """Maps a table to an organization."""
+        try:
+            print(f"Mapping table {table_name} to organization {org_id}")
+            if self.session:
+                print(f"Session: {self.session}")
+                alias = table_name if not alias else alias
+                self.session.add(
+                    OrganizationTableMap(
+                        organization_id=org_id, table_name=table_name, table_alias=alias
+                    )
+                )
+                self.session.commit()
+        except Exception as e:
+            self.session.rollback() if self.session else None
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+
     def create_table_with_llm(self, sample_content: str, header: str, extra_desc: str):
         """
         Creates a table using an LLM based on sample file content and a message.
@@ -32,9 +57,8 @@ def create_table_with_llm(self, sample_content: str, header: str, extra_desc: st
         - create_query: str containing the SQL create table query if successful, None otherwise.
         """
         try:
-            with DatabaseManager() as session:
-                sql_executor = SQLExecutor(session)
-                table_names = sql_executor.get_all_table_names_as_str()
+            sql_executor = SQLExecutor(self.session)
+            table_names = sql_executor.get_all_table_names_as_str()
 
             raw_create_query = self.llm.generate_create_statement(
                 sample_content, header, table_names, extra_desc
@@ -47,9 +71,8 @@ def create_table_with_llm(self, sample_content: str, header: str, extra_desc: st
             if SQLStringManipulator(
                 create_query
             ).is_valid_create_table_query():  # Checks if the query is valid
-                with DatabaseManager() as session:
-                    sql_executor = SQLExecutor(session)
-                    sql_executor.execute_create_query(create_query)
+                sql_executor = SQLExecutor(self.session)
+                sql_executor.execute_create_query(create_query)
                 return create_query
         except Exception as e:
             # Log the error message here
@@ -79,9 +102,8 @@ def create_table_desc_with_llm(
             ).get_table_from_create_query()
 
             # Store description in separate table
-            with DatabaseManager() as session:
-                manager = TableMetadataManager(session)
-                manager.add_table_metadata(table_name, create_query, description)
+            manager = TableMetadataManager(self.session)
+            manager.add_table_metadata(table_name, create_query, description)
 
         except Exception as e:
             # Log the error message here
@@ -99,27 +121,27 @@ def determine_table(self, sample_content: str, extra_desc: str) -> str:
         Returns:
         - table_name: str containing the table's name.
         """
-        with DatabaseManager() as session:
-            manager = TableMetadataManager(session)
-            table_metadata = manager.get_metadata()
-            formatted_table_metadata = manager.format_table_metadata_for_llm(
-                table_metadata
-            )
+        manager = TableMetadataManager(self.session)
+        table_metadata = manager.get_metadata()
+        formatted_table_metadata = manager.format_table_metadata_for_llm(table_metadata)
 
         table_name: str = self.llm.fetch_table_name_from_sample(
             sample_content, extra_desc, formatted_table_metadata
         )
         return table_name
 
-    def create_table_from_df(self, df: pd.DataFrame, table_name: str):
-        with DatabaseManager() as session:
-            executor = SQLExecutor(session)
+    def create_table_from_df(self, df: pd.DataFrame, org_id: int, table_name: str):
+        try:
+            executor = SQLExecutor(self.session)
             executor.append_df_to_table(df, table_name)
+            self._map_table_to_org(org_id, table_name)
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
 
     def create_table_from_query(self, query: str):
-        with DatabaseManager() as session:
-            executor = SQLExecutor(session)
-            executor.execute_create_query(query)
+        executor = SQLExecutor(self.session)
+        executor.execute_create_query(query)
 
     def append_to_table(self, processed_df: pd.DataFrame, table_name: str):
         """
@@ -137,9 +159,8 @@ def append_to_table(self, processed_df: pd.DataFrame, table_name: str):
         """
 
         if table_name:
-            with DatabaseManager() as session:
-                sql_executor = SQLExecutor(session)
-                sql_executor.append_df_to_table(processed_df, table_name)
+            sql_executor = SQLExecutor(self.session)
+            sql_executor.append_df_to_table(processed_df, table_name)
         else:
             raise HTTPException(
                 status_code=400, detail="Could not determine table name"
@@ -148,19 +169,35 @@ def append_to_table(self, processed_df: pd.DataFrame, table_name: str):
     def drop_table(self, table_name: str):
         # Logic to drop a table
         try:
-            with DatabaseManager() as session:
-                executor = SQLExecutor(session)
-                executor.drop_table(table_name)
+            executor = SQLExecutor(self.session)
+            executor.drop_table(table_name)
+        except Exception as e:
+            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:
+            if self.session:
+                table_names = (
+                    self.session.query(OrganizationTableMap.table_name)
+                    .filter(OrganizationTableMap.organization_id == org_id)
+                    .all()
+                )
+                return [
+                    name[0] for name in table_names
+                ]  # Extracting table_name from each tuple
+            return []
         except Exception as e:
+            self.session.rollback() if self.session else None
             print(f"An error occurred: {e}")
             raise HTTPException(status_code=400, detail=str(e))
 
     def get_table_columns(self, table_name: str):
         """Returns a list of all of the columns present within the table."""
         try:
-            with DatabaseManager() as session:
-                executor = SQLExecutor(session)
-                columns = executor.get_table_columns(table_name)
+            executor = SQLExecutor(self.session)
+            columns = executor.get_table_columns(table_name)
             return columns
         except Exception as e:
             print(f"An error occurred: {e}")

backend/envs/dev/utils.py

@@ -1,9 +1,12 @@
+import pandas as pd
+
 from databases.database_manager import DatabaseManager
 from databases.sql_executor import SQLExecutor
 from databases.table_manager import TableManager
 from databases.table_metadata_manager import TableMetadataManager
+from utils.utils import get_app_logger
 
-import pandas as pd
+logger = get_app_logger(__name__)
 
 
 def seed_db():
@@ -30,15 +33,16 @@ def seed_db():
         executor = SQLExecutor(session)
         existing_tables = executor.get_all_table_names_as_list()
 
-    # Create sample table if it doesn't exist
-    table_manager = TableManager()
-    if "sample_sales" not in existing_tables:
-        df = pd.read_csv("envs/dev/sample_data/sample_sales_data.csv")
-        df.columns = map(str.lower, df.columns)
-        table_manager.create_table_from_df(df, "sample_sales")
+        # Create sample table if it doesn't exist
+        table_manager = TableManager(session)
+        if "sample_sales" not in existing_tables:
+            df = pd.read_csv("envs/dev/sample_data/sample_sales_data.csv")
+            df.columns = map(str.lower, df.columns)
+            table_manager.create_table_from_df(
+                df=df, org_id=1, table_name="sample_sales"
+            )
 
-        # Add metadata
-        with DatabaseManager() as session:
+            # Add metadata
             metadata_manager = TableMetadataManager(session)
             metadata_manager.add_table_metadata(
                 table_name="sample_sales",
@@ -82,3 +86,5 @@ def seed_db():
                     and time-based sales performance.
                 """,
             )
+            logger.info('Created table "sample_sales".')
+        logger.info('Sample table "sample_sales" already exists.')

backend/models/organization_table_map.py

@@ -0,0 +1,22 @@
+from sqlalchemy import Column, Integer, String
+
+from .base import Base
+
+
+class OrganizationTableMap(Base):
+    """
+    Represents a mapping between an organization and a table.
+
+    Attributes:
+        id (int): The unique identifier of the mapping.
+        organization_id (int): The ID of the organization.
+        table_name (str): The name of the table.
+        table_alias (str): The alias for the table.
+    """
+
+    __tablename__ = "organization_table_map"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    organization_id = Column(Integer)
+    table_name = Column(String, unique=True)
+    table_alias = Column(String)

backend/routes/file_routes.py

@@ -2,6 +2,7 @@
 from fastapi.exceptions import HTTPException
 from fastapi.responses import JSONResponse
 
+from databases.database_manager import DatabaseManager
 from databases.table_manager import TableManager
 from llms.base import BaseLLM
 from llms.utils import get_llm_sql_object
@@ -52,23 +53,24 @@ async def upload_file(
         processed_df, sample_content, header = process_file(file, encoding)
 
         # Instantiate TableManager
-        table_manager = TableManager(llm)
-
-        # Create new table if necessary
-        if is_new_table:
-            create_table_query = table_manager.create_table_with_llm(
-                sample_content, header, extra_desc
-            )
-            table_manager.create_table_desc_with_llm(
-                create_table_query, sample_content, extra_desc
-            )
-
-        # Append file to table
-        table_name = table_manager.determine_table(sample_content, extra_desc)
-        table_manager.append_to_table(processed_df, table_name)
-
-        # Optionally, save the file to a data lake
-        save_to_data_lake(file)
+        with DatabaseManager() as session:
+            table_manager = TableManager(session, llm)
+
+            # Create new table if necessary
+            if is_new_table:
+                create_table_query = table_manager.create_table_with_llm(
+                    sample_content, header, extra_desc
+                )
+                table_manager.create_table_desc_with_llm(
+                    create_table_query, sample_content, extra_desc
+                )
+
+            # Append file to table
+            table_name = table_manager.determine_table(sample_content, extra_desc)
+            table_manager.append_to_table(processed_df, table_name)
+
+            # Optionally, save the file to a data lake
+            save_to_data_lake(file)
 
         return JSONResponse(
             content={"message": "File processed successfully"}, status_code=200

backend/routes/table_routes.py

@@ -10,12 +10,21 @@
 table_router = APIRouter()
 
 
+@table_router.get("/organization/{org_id}/tables/")
+async def get_org_tables(org_id: int, current_user: User = Depends(get_current_user)):
+    with DatabaseManager() as session:
+        manager = TableManager(session)
+        tables = manager.get_org_tables(org_id)
+    return tables
+
+
 @table_router.get("/table/columns/")
 async def get_table_columns(
     table_name: str, current_user: User = Depends(get_current_user)
 ):
-    manager = TableManager()
-    columns = manager.get_table_columns(table_name)
+    with DatabaseManager() as session:
+        manager = TableManager(session)
+        columns = manager.get_table_columns(table_name)
     return columns
 
 
@@ -47,6 +56,7 @@ async def get_all_table_metadata(current_user: User = Depends(get_current_user))
 
 @table_router.delete("/table/")
 async def drop_table(table_name: str, current_user: User = Depends(get_current_user)):
-    manager = TableManager()
-    manager.drop_table(table_name)
+    with DatabaseManager() as session:
+        manager = TableManager(session)
+        manager.drop_table(table_name)
     return {"message": f"Dropped table {table_name}"}

frontend/src/pages/Charts/Configs/ChartConfig.jsx

@@ -11,18 +11,29 @@ function ChartConfig({ onConfigChange, onRequiredSelected, chartConfig }) {
     const [selectedTable, setSelectedTable] = useState('');
     const [chartTypes, setChartTypes] = useState([]);
     const [selectedChartType, setSelectedChartType] = useState('');
+    const [organizationId, setOrganizationId] = useState(null);
 
     useEffect(() => {
-        // Fetch tables from API
-        axios.get(`${API_URL}tables/`)
-            .then(response => setTables(response.data))
-            .catch(error => console.error('Error fetching tables:', error))
+        // Fetch user data from API
+        axios.get(`${API_URL}users/me/`)
+            .then(response => {
+                // Set organizationId state
+                setOrganizationId(response.data.organization_id);
+
+                // Fetch tables from API using organizationId
+                axios.get(`${API_URL}organization/${response.data.organization_id}/tables/`)
+                    .then(response => setTables(response.data))
+                    .catch(error => console.error('Error fetching tables:', error));
+            })
+            .catch(error => console.error('Error fetching user data:', error));
+    }, []);
 
+    useEffect(() => {
         // Fetch chart types from API
         axios.get(`${API_URL}charts/types/`)
             .then(response => setChartTypes(response.data))
             .catch(error => console.error('Error fetching chart types:', error));
-    }, []);
+    }, []); // Empty dependency array means this effect runs once on mount
 
     useEffect(() => {
         onConfigChange({