Skip to content

Commit

Permalink
fix: handle invalid data source type
Browse files Browse the repository at this point in the history
  • Loading branch information
ArslanSaleem committed Nov 19, 2024
1 parent 237c67b commit d759679
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 5 deletions.
8 changes: 3 additions & 5 deletions pandasai/dataframe/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from datetime import datetime, timedelta
import hashlib

from pandasai.exceptions import InvalidDataSourceType
from pandasai.helpers.path import find_project_root
from .base import DataFrame
import importlib
Expand Down Expand Up @@ -101,11 +102,8 @@ def _load_from_source(self) -> pd.DataFrame:
load_function = getattr(module, f"load_from_{source_type}")
return load_function(connection_info, query)
else:
connector_class = getattr(
module, f"{source_type.capitalize()}Connector"
)
connector = connector_class(config=connection_info)
return connector.execute_query(query)
raise InvalidDataSourceType("Invalid data source type")

except ImportError as e:
raise ImportError(
f"{source_type.capitalize()} connector not found. "
Expand Down
6 changes: 6 additions & 0 deletions pandasai/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -250,3 +250,9 @@ class LazyLoadError(Exception):
"""Raised when trying to access data that hasn't been loaded in lazy load mode."""

pass


class InvalidDataSourceType(Exception):
"""Raised error with invalid data source provided"""

pass

0 comments on commit d759679

Please sign in to comment.