diff --git a/backend/databases/sql_executor.py b/backend/databases/sql_executor.py index 198ba3b..5dc7d5e 100644 --- a/backend/databases/sql_executor.py +++ b/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)}" ) diff --git a/backend/databases/table_manager.py b/backend/databases/table_manager.py index 20715e9..bae4911 100644 --- a/backend/databases/table_manager.py +++ b/backend/databases/table_manager.py @@ -1,12 +1,15 @@ +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: """ @@ -14,12 +17,34 @@ class TableManager: 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}") diff --git a/backend/envs/dev/utils.py b/backend/envs/dev/utils.py index 3cba575..4f55251 100644 --- a/backend/envs/dev/utils.py +++ b/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.') diff --git a/backend/models/organization_table_map.py b/backend/models/organization_table_map.py new file mode 100644 index 0000000..0512500 --- /dev/null +++ b/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) diff --git a/backend/routes/file_routes.py b/backend/routes/file_routes.py index 2013e42..b80e02a 100644 --- a/backend/routes/file_routes.py +++ b/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 diff --git a/backend/routes/table_routes.py b/backend/routes/table_routes.py index 922758f..3df2f9d 100644 --- a/backend/routes/table_routes.py +++ b/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}"} diff --git a/frontend/src/pages/Charts/Configs/ChartConfig.jsx b/frontend/src/pages/Charts/Configs/ChartConfig.jsx index dfd1985..3a88600 100644 --- a/frontend/src/pages/Charts/Configs/ChartConfig.jsx +++ b/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({