Repository: gventuri/pandas-ai
File: pandasai/connectors/airtable.py (Python)

Code before:
"""
Airtable connectors are used to connect airtable records.
"""

from .base import AirtableConnectorConfig, BaseConnector, BaseConnectorConfig
from typing import Union, Optional
import requests
import pandas as pd


class AirtableConnector(BaseConnector):
    """
    Airtable connector to retrieving record data.
    """

    def __init__(
        self,
        config: Optional[Union[AirtableConnectorConfig, dict]] = None,
    ):
        if isinstance(config, dict):
            if config["token"] and config["baseID"] and config["table"]:
                config = AirtableConnectorConfig(**config)

        elif not config:
            airtable_env_vars = {
                "token": "AIRTABLE_AUTH_TOKEN",
                "baseID": "AIRTABLE_BASE_ID",
                "table": "AIRTABLE_TABLE_NAME",
            }
            config = AirtableConnectorConfig(
                **self._populate_config_from_env(config, airtable_env_vars)
            )

        self._root_url: str = "https://api.airtable.com/v0/"

        super().__init__(config)

    def _init_connection(self, config: BaseConnectorConfig):
        """
        make connection to database
        """
        config = config.dict()
        _session = requests.Session()
        _session.headers = {"Authorization": f"Bearer {config['token']}"}
        url = f"{self._root_url}{config['baseID']}/{config['table']}"
        response = _session.head(url=url)
        if response.status_code == 200:
            self._session = _session
        else:
            raise ValueError(
                f"""Failed to connect to Airtable. 
                Status code: {response.status_code}, 
                message: {response.text}"""
            )

    def execute(self):
        """
        Execute the connector and return the result.

        Returns:
            DataFrameType: The result of the connector.
        """
        url = f"{self._root_url}{self.config['baseID']}/{self.config['table']}"
        if self._session:
            _response = self._session.get(url)
            if _response.status_code == 200:
                data = _response.json()
                ## Following column selection is done
                ## to prepare output in favaourable format.
                records = [
                    {"id": record["id"], **record["fields"]}
                    for record in data["records"]
                ]
                self._response = pd.DataFrame(records)
            else:
                raise ValueError(
                    f"""Failed to connect to Airtable. 
                    Status code: {_response.status_code}, 
                    message: {_response.text}"""
                )
        return self._response

    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()


Review comment (suggestion):
Let's also add a `def fallback_name`, similarly to all the other connectors.

Example:
```
@property
    def fallback_name(self):
        """
        Return the fallback name of the connector.

        Returns:
            str: The fallback name of the connector.
        """
        return self._config.table
```

Code after:
        return response

    def _fetch_data(self):
        """
        Fetches data from the Airtable server via API and converts it to a DataFrame.
        """

        params = {
            "pageSize": 100,
            "offset": "0"
        }

        if self._config.where is not None:
            params["filterByFormula"] = self._build_formula()

        data = []
        while True:
            response = self._request_api(params=params)

            if response.status_code != 200:
                raise InvalidRequestError(
                    f"Failed to connect to Airtable. "
                    f"Status code: {response.status_code}, "
                    f"message: {response.text}"
                )

            res = response.json()
            records = res.get("records", [])
            data.extend({"id": record["id"], **record["fields"]} for record in records)

            if len(records) < 100 or "offset" not in res:
                break

            if "offset" in res:
                params["offset"] = res["offset"]

        return pd.DataFrame(data)

    @cache
    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 .
        """
        data = self._request_api(params={"maxRecords": 5})
        return pd.DataFrame(
            [