Airtable connector

pandasai/connectors/__init__.py

@@ -9,6 +9,7 @@
 from .snowflake import SnowFlakeConnector
 from .databricks import DatabricksConnector
 from .yahoo_finance import YahooFinanceConnector
+from .airtable import AirtableConnector
 
 __all__ = [
     "BaseConnector",
@@ -18,4 +19,5 @@
     "YahooFinanceConnector",
     "SnowFlakeConnector",
     "DatabricksConnector",
+    "AirtableConnector"
 ]

pandasai/connectors/airtable.py

@@ -0,0 +1,74 @@
+"""
+Airtable connectors are used to connect airtable records.
+"""
+
+from .base import AirtableConnectorConfig , BaseConnector
+from typing import Union ,Optional
+import os
+import requests
+import pandas as pd
+
+class AirtableConnector(BaseConnector):
+    """
+    Airtable connector to retrieving record data.
+    """
+    def __init__(
+            self,
+            baseid : Optional[str] = None,
+            bearer_token : Optional[str] = None,
+            table_name : Optional[str] = None,
+            config : Union[AirtableConnectorConfig,dict] = None
+    ):
+        if not bearer_token and not config:
+            raise ValueError(
+                "You must specify bearer token for authentication or a config object."
+            )
+        if not baseid and not config.baseID:
+            raise ValueError(
+                "You must specify baseId or a proper config object."
+            )
+        if not isinstance(config,AirtableConnectorConfig):
+            if not config :
+                config = {}
+            if table_name :
+                config['table'] = table_name
+
+            airtable_config = AirtableConnectorConfig(**config)
+        else :
+            airtable_config = config
+
+        self._session = requests.Session()
+
+        self._session.headers = {
+            "Authorization" : f"Bearer {airtable_config['token']}"
+        }
+
+        self._response : str = None
+
+        super().__init__(airtable_config)
+
+    def _connect_load(self,config:AirtableConnectorConfig):
+        """
+        Authenticate and connect to the instance 
+
+        Args :
+            config (AirtableConnectorConfig) :Configuration to load table data.
+        """
+        _response = self._session.get(f"https://api.airtable.com/v0/{config['baseID']}/{config['table']}")
+        if _response.status_code == 200 :
+            self._response = pd.read_json(_response.json())
+        else :
+            raise ValueError(
+                "Please verify authentication details , baseid or table_name are correct ."
+            )
+
+
+    def head(self):
+        """
+        Return the head of the table that the connector is connected to.
+
+        Returns :
+            DatFrameType: The head of the data source that the conector is connected to . 
+        """
+        return self._response.head()
+

pandasai/connectors/base.py

@@ -74,6 +74,12 @@ class DatabricksConnectorConfig(SQLBaseConnectorConfig):
     token: str
     httpPath: str
 
+class AirtableConnectorConfig(BaseConnectorConfig):
+    """
+    Connector configuration for Airtable.
+    """
+    token : str
+    baseID : str
 
 class BaseConnector(ABC):
     """

tests/connectors/test_airtable.py

@@ -0,0 +1,42 @@
+import unittest
+import pandas 
+from unittest.mock import Mock, patch
+from pandasai.connectors.base import AirtableConnectorConfig
+from pandasai.connectors import AirtableConnector
+import requests
+
+class TestAirTableConnector(unittest.TestCase):
+    def setUp(self):
+
+        # Define your ConnectorConfig instance here
+        self.config = AirtableConnectorConfig(
+            token="your_token",
+            baseID="your_baseid",
+            table="your_table_name"
+        ).dict()
+
+        #Create an instance of Connector 
+        self.connector = AirtableConnector(config=self.config)
+
+    @patch("requests.Session")
+    def test_constructor_and_properties(self,mock_request_session):
+        # Test constructor and properties
+        self.assertEqual(self.connector._config,self.config)
+        self.assertEqual(self.connector._session,mock_request_session)
+        self.assertEqual(self.connector._session.headers, {
+            "Authorization" : f"Bearer {self.config['token']}"
+        })
+
+    @patch("pandasai.connectors.AirtableConnector._connect_load")
+    @patch("pandas.DataFrame")
+    def test_connect_load(self,_connect_load,mock_dataframe):
+        response = _connect_load(self.config)
+        self.assertEqual(type(self.connector._response),mock_dataframe())
+
+    @patch("pandasai.connectors.AirtableConnector.head")
+    @patch("pandas.DataFrame")
+    def test_head(self,mock_head,mock_dataframe):
+        response = mock_head()
+        self.assertEqual(type(response),mock_dataframe())
+
+