MindsDB 25.5.4.1__py3-none-any.whl → 25.6.2.0__py3-none-any.whl

mindsdb/integrations/handlers/salesforce_handler/salesforce_handler.py

@@ -1,14 +1,14 @@
-from typing import Any, Dict, Text
+from typing import Any, Dict, List, Optional, Text
 
 import pandas as pd
 import salesforce_api
 from salesforce_api.exceptions import AuthenticationError, RestRequestCouldNotBeUnderstoodError
 
-from mindsdb.integrations.libs.api_handler import APIHandler
+from mindsdb.integrations.libs.api_handler import MetaAPIHandler
 from mindsdb.integrations.libs.response import (
     HandlerResponse as Response,
     HandlerStatusResponse as StatusResponse,
-    RESPONSE_TYPE
+    RESPONSE_TYPE,
 )
 from mindsdb.integrations.handlers.salesforce_handler.salesforce_tables import create_table_class
 from mindsdb.utilities import log
@@ -17,12 +17,12 @@ from mindsdb.utilities import log
 logger = log.getLogger(__name__)
 
 
-class SalesforceHandler(APIHandler):
+class SalesforceHandler(MetaAPIHandler):
     """
     This handler handles the connection and execution of SQL statements on Salesforce.
     """
 
-    name = 'salesforce'
+    name = "salesforce"
 
     def __init__(self, name: Text, connection_data: Dict, **kwargs: Any) -> None:
         """
@@ -57,23 +57,23 @@ class SalesforceHandler(APIHandler):
             return self.connection
 
         # Mandatory connection parameters.
-        if not all(key in self.connection_data for key in ['username', 'password', 'client_id', 'client_secret']):
+        if not all(key in self.connection_data for key in ["username", "password", "client_id", "client_secret"]):
             raise ValueError("Required parameters (username, password, client_id, client_secret) must be provided.")
 
         try:
             self.connection = salesforce_api.Salesforce(
-                username=self.connection_data['username'],
-                password=self.connection_data['password'],
-                client_id=self.connection_data['client_id'],
-                client_secret=self.connection_data['client_secret'],
-                is_sandbox=self.connection_data.get('is_sandbox', False)
+                username=self.connection_data["username"],
+                password=self.connection_data["password"],
+                client_id=self.connection_data["client_id"],
+                client_secret=self.connection_data["client_secret"],
+                is_sandbox=self.connection_data.get("is_sandbox", False),
             )
             self.is_connected = True
 
             # Register Salesforce tables.
             for resource_name in self._get_resource_names():
                 table_class = create_table_class(resource_name)
-                self._register_table(resource_name.lower(), table_class(self))
+                self._register_table(resource_name, table_class(self))
 
             return self.connection
         except AuthenticationError as auth_error:
@@ -96,10 +96,10 @@ class SalesforceHandler(APIHandler):
             self.connect()
             response.success = True
         except (AuthenticationError, ValueError) as known_error:
-            logger.error(f'Connection check to Salesforce failed, {known_error}!')
+            logger.error(f"Connection check to Salesforce failed, {known_error}!")
             response.error_message = str(known_error)
         except Exception as unknown_error:
-            logger.error(f'Connection check to Salesforce failed due to an unknown error, {unknown_error}!')
+            logger.error(f"Connection check to Salesforce failed due to an unknown error, {unknown_error}!")
             response.error_message = str(unknown_error)
 
         self.is_connected = response.success
@@ -123,7 +123,7 @@ class SalesforceHandler(APIHandler):
 
             parsed_results = []
             for result in results:
-                del result['attributes']
+                del result["attributes"]
 
                 # Check if the result contains any of the other Salesforce resources.
                 if any(key in self.resource_names for key in result.keys()):
@@ -131,8 +131,10 @@ class SalesforceHandler(APIHandler):
                     parsed_result = {}
                     for key, value in result.items():
                         if key in self.resource_names:
-                            del value['attributes']
-                            parsed_result.update({f'{key}_{sub_key}': sub_value for sub_key, sub_value in value.items()})
+                            del value["attributes"]
+                            parsed_result.update(
+                                {f"{key}_{sub_key}": sub_value for sub_key, sub_value in value.items()}
+                            )
 
                         else:
                             parsed_result[key] = value
@@ -142,24 +144,13 @@ class SalesforceHandler(APIHandler):
                 else:
                     parsed_results.append(result)
 
-            response = Response(
-                RESPONSE_TYPE.TABLE,
-                pd.DataFrame(parsed_results)
-            )
+            response = Response(RESPONSE_TYPE.TABLE, pd.DataFrame(parsed_results))
         except RestRequestCouldNotBeUnderstoodError as rest_error:
-            logger.error(f'Error running query: {query} on Salesforce, {rest_error}!')
-            response = Response(
-                RESPONSE_TYPE.ERROR,
-                error_code=0,
-                error_message=str(rest_error)
-            )
+            logger.error(f"Error running query: {query} on Salesforce, {rest_error}!")
+            response = Response(RESPONSE_TYPE.ERROR, error_code=0, error_message=str(rest_error))
         except Exception as unknown_error:
-            logger.error(f'Error running query: {query} on Salesforce, {unknown_error}!')
-            response = Response(
-                RESPONSE_TYPE.ERROR,
-                error_code=0,
-                error_message=str(unknown_error)
-            )
+            logger.error(f"Error running query: {query} on Salesforce, {unknown_error}!")
+            response = Response(RESPONSE_TYPE.ERROR, error_code=0, error_message=str(unknown_error))
 
         return response
 
@@ -171,6 +162,34 @@ class SalesforceHandler(APIHandler):
             None
         """
         if not self.resource_names:
-            self.resource_names = [resource['name'] for resource in self.connection.sobjects.describe()['sobjects']]
+            # Fetch the queryable list of Salesforce resources (sobjects).
+            self.resource_names = [
+                resource["name"]
+                for resource in self.connection.sobjects.describe()["sobjects"]
+                if resource.get("queryable", False)
+            ]
 
         return self.resource_names
+
+    def meta_get_tables(self, table_names: Optional[List[str]] = None) -> Response:
+        """
+        Retrieves metadata for the specified tables (or all tables if no list is provided).
+
+        Args:
+            table_names (List): A list of table names for which to retrieve metadata.
+
+        Returns:
+            Response: A response object containing the table metadata.
+        """
+        connection = self.connect()
+
+        # Retrieve the metadata for all Salesforce resources.
+        main_metadata = connection.sobjects.describe()
+
+        if table_names:
+            # Filter the metadata for the specified tables.
+            main_metadata = [resource for resource in main_metadata["sobjects"] if resource["name"] in table_names]
+        else:
+            main_metadata = main_metadata["sobjects"]
+
+        return super().meta_get_tables(table_names=table_names, main_metadata=main_metadata)

mindsdb/integrations/handlers/salesforce_handler/salesforce_tables.py

@@ -2,8 +2,9 @@ from typing import Dict, List, Text
 
 from mindsdb_sql_parser.ast import Select, Star, Identifier
 import pandas as pd
+from salesforce_api.exceptions import RestRequestCouldNotBeUnderstoodError
 
-from mindsdb.integrations.libs.api_handler import APIResource
+from mindsdb.integrations.libs.api_handler import MetaAPIResource
 from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
 from mindsdb.utilities import log
 
@@ -11,16 +12,28 @@ from mindsdb.utilities import log
 logger = log.getLogger(__name__)
 
 
-def create_table_class(resource_name: Text) -> APIResource:
+def create_table_class(resource_name: Text) -> MetaAPIResource:
     """
     Creates a table class for the given Salesforce resource.
     """
 
-    class AnyTable(APIResource):
+    class AnyTable(MetaAPIResource):
         """
         This is the table abstraction for any resource of the Salesforce API.
         """
 
+        def __init__(self, *args, table_name=None, **kwargs):
+            """
+            Initializes the AnyTable class.
+
+            Args:
+                *args: Variable length argument list.
+                table_name (str): The name of the table that represents the Salesforce resource.
+                **kwargs: Arbitrary keyword arguments.
+            """
+            super().__init__(*args, table_name=table_name, **kwargs)
+            self.resource_metadata = None
+
         def select(self, query: Select) -> pd.DataFrame:
             """
             Executes a SELECT SQL query represented by an ASTNode object on the Salesforce resource and retrieves the data (if any).
@@ -53,7 +66,7 @@ def create_table_class(resource_name: Text) -> APIResource:
             results = client.sobjects.query(query_str)
 
             for result in results:
-                del result['attributes']
+                del result["attributes"]
 
             df = pd.DataFrame(results)
             df.rename(columns=column_aliases, inplace=True)
@@ -107,7 +120,7 @@ def create_table_class(resource_name: Text) -> APIResource:
                 conditions (List[FilterCondition]): The conditions to be validated.
             """
             # Salesforce API does not support filtering items based on attributes other than 'Id'. Raise an error if any other column is used.
-            if len(conditions) != 1 or conditions[0].column != 'Id':
+            if len(conditions) != 1 or conditions[0].column != "Id":
                 raise ValueError("Only the 'Id' column can be used to filter items.")
 
             # Only the 'equals' and 'in' operators can be used on the 'Id' column for deletion. Raise an error if any other operator is used.
@@ -116,6 +129,19 @@ def create_table_class(resource_name: Text) -> APIResource:
 
             return conditions[0].value if isinstance(conditions[0].value, list) else [conditions[0].value]
 
+        def _get_resource_metadata(self) -> Dict:
+            """
+            Retrieves metadata about the Salesforce resource.
+
+            Returns:
+                Dict: A dictionary containing metadata about the Salesforce resource.
+            """
+            if self.resource_metadata:
+                return self.resource_metadata
+
+            client = self.handler.connect()
+            return getattr(client.sobjects, resource_name).describe()
+
         def get_columns(self) -> List[Text]:
             """
             Retrieves the attributes (columns) of the Salesforce resource.
@@ -123,7 +149,111 @@ def create_table_class(resource_name: Text) -> APIResource:
             Returns:
                 List[Text]: A list of Attributes (columns) of the Salesforce resource.
             """
+            return [field["name"] for field in self._get_resource_metadata()["fields"]]
+
+        def meta_get_tables(self, table_name: str, main_metadata: Dict) -> Dict:
+            """
+            Retrieves table metadata for the Salesforce resource.
+
+            Args:
+                table_name (str): The name given to the table that represents the Salesforce resource.
+                main_metadata (Dict): The main metadata dictionary containing information about all Salesforce resources.
+
+            Returns:
+                Dict: A dictionary containing table metadata for the Salesforce resource.
+            """
             client = self.handler.connect()
-            return [field['name'] for field in getattr(client.sobjects, resource_name).describe()['fields']]
+
+            resource_metadata = next(
+                (resource for resource in main_metadata if resource["name"] == resource_name),
+            )
+
+            # Get row count if Id column is aggregatable.
+            row_count = None
+            # if next(field for field in resource_metadata['fields'] if field['name'] == 'Id').get('aggregatable', False):
+            try:
+                row_count = client.sobjects.query(f"SELECT COUNT(Id) FROM {resource_name}")[0]["expr0"]
+            except RestRequestCouldNotBeUnderstoodError as request_error:
+                logger.warning(f"Failed to get row count for {resource_name}: {request_error}")
+
+            return {
+                "table_name": table_name,
+                "table_type": "BASE TABLE",
+                "table_description": resource_metadata.get("label", ""),
+                "row_count": row_count,
+            }
+
+        def meta_get_columns(self, table_name: str) -> List[Dict]:
+            """
+            Retrieves column metadata for the Salesforce resource.
+
+            Args:
+                table_name (str): The name given to the table that represents the Salesforce resource.
+
+            Returns:
+                List[Dict]: A list of dictionaries containing column metadata for the Salesforce resource.
+            """
+            resource_metadata = self._get_resource_metadata()
+
+            column_metadata = []
+            for field in resource_metadata["fields"]:
+                column_metadata.append(
+                    {
+                        "table_name": table_name,
+                        "column_name": field["name"],
+                        "data_type": field["type"],
+                        "is_nullable": field.get("nillable", False),
+                    }
+                )
+
+            return column_metadata
+
+        def meta_get_primary_keys(self, table_name: str) -> List[Dict]:
+            """
+            Retrieves the primary keys for the Salesforce resource.
+
+            Args:
+                table_name (str): The name given to the table that represents the Salesforce resource.
+
+            Returns:
+                List[Dict]: A list of dictionaries containing primary key metadata for the Salesforce resource.
+            """
+            return [
+                {
+                    "table_name": table_name,
+                    "column_name": "Id",
+                }
+            ]
+
+        def meta_get_foreign_keys(self, table_name: str, all_tables: List[str]) -> List[Dict]:
+            """
+            Retrieves the foreign keys for the Salesforce resource.
+
+            Args:
+                table_name (str): The name given to the table that represents the Salesforce resource.
+                all_tables (List[str]): A list of all table names in the Salesforce database.
+
+            Returns:
+                List[Dict]: A list of dictionaries containing foreign key metadata for the Salesforce resource.
+            """
+            resource_metadata = self._get_resource_metadata()
+
+            foreign_key_metadata = []
+            for child_relationship in resource_metadata.get("childRelationships", []):
+                # Skip if the child relationship is not one of the supported tables.
+                child_table_name = child_relationship["childSObject"]
+                if child_table_name not in all_tables:
+                    continue
+
+                foreign_key_metadata.append(
+                    {
+                        "parent_table_name": table_name,
+                        "parent_column_name": "Id",
+                        "child_table_name": child_table_name,
+                        "child_column_name": child_relationship["field"],
+                    }
+                )
+
+            return foreign_key_metadata
 
     return AnyTable