feat: add support for running queries with the tooling api

README.md

@@ -1306,6 +1306,8 @@ queries:
   env:
     end_date: "now()"
     start_date: "now(timedelta(minutes=-60))"
+- query: "SELECT FullName FROM EntityDefinition WHERE Label='Opportunity'"
+  api_name: tooling
 ```
 
 #### Query configuration
@@ -1332,6 +1334,32 @@ query to execute.
 The `api_ver` attribute can be used to customize the version of the Salesforce
 API that the exporter should use when executing query API calls.
 
+##### `api_name`
+
+| Description | Valid Values | Required | Default |
+| --- | --- | --- | --- |
+| The name of the Salesforce Platform API to use  | `rest` / `tooling` | N | `rest` |
+
+The `api_name` attribute can be used to specify the name of the Salesforce
+Platform API that the exporter should use when executing query API calls.
+
+When the `api_name` attribute is not set or when it is set to the value `rest`,
+the [ReST API](https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/intro_rest.htm)
+will be used to execute the
+[SOQL](https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql.htm)
+query.
+
+When the `api_name` attribute is set to `tooling`, the
+[Tooling API](https://developer.salesforce.com/docs/atlas.en-us.api_tooling.meta/api_tooling/intro_api_tooling.htm)
+will be used to execute the
+[SOQL](https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql.htm)
+query.
+
+Specifying any other value will cause an error and the exporter will terminate.
+
+**NOTE:** Not all queries can be executed with both APIs. Ensure that the query
+being used is appropriate for the specified Salesforce Platform API.
+
 ##### `id`
 
 | Description | Valid Values | Required | Default |

config_sample.yml

@@ -50,6 +50,8 @@ queries:
     start_date: "now(timedelta(minutes=-60))"
   api_ver: "58.0"
   timestamp_attr: StartDate
+- query: "SELECT FullName FROM EntityDefinition WHERE Label='Opportunity'"
+  api_name: tooling
 newrelic:
   data_format: events
   api_endpoint: US

src/newrelic_logging/__init__.py

@@ -3,7 +3,7 @@
 
 # Integration definitions
 
-VERSION         = "2.3.0"
+VERSION         = "2.4.0"
 NAME            = "salesforce-exporter"
 PROVIDER        = "newrelic-labs"
 COLLECTOR_NAME  = "newrelic-salesforce-exporter"

src/newrelic_logging/api.py

@@ -5,6 +5,10 @@
 from .auth import Authenticator
 from .telemetry import print_warn
 
+API_NAME_REST = 'rest'
+API_NAME_TOOLING = 'tooling'
+DEFAULT_API_NAME = API_NAME_REST
+
 def get(
     auth: Authenticator,
     session: Session,
@@ -72,6 +76,21 @@ def stream_lines(response: Response, chunk_size: int):
     )
 
 
+def get_query_api_path(api_ver: str, api_name: str) -> str:
+    l_api_name = api_name.lower()
+
+    if l_api_name == API_NAME_REST:
+        return f'/services/data/v{api_ver}/query'
+
+    if l_api_name == API_NAME_TOOLING:
+        return f'/services/data/v{api_ver}/tooling/query'
+
+    raise SalesforceApiException(
+        -1,
+        f'invalid query api name {api_name}',
+    )
+
+
 class Api:
     def __init__(self, authenticator: Authenticator, api_ver: str):
         self.authenticator = authenticator
@@ -80,15 +99,27 @@ def __init__(self, authenticator: Authenticator, api_ver: str):
     def authenticate(self, session: Session) -> None:
         self.authenticator.authenticate(session)
 
-    def query(self, session: Session, soql: str, api_ver: str = None) -> dict:
+    def query(
+        self,
+        session: Session,
+        soql: str,
+        api_ver: str = None,
+        api_name: str = None,
+    ) -> dict:
         ver = self.api_ver
         if not api_ver is None:
             ver = api_ver
 
+        api = DEFAULT_API_NAME
+        if not api_name is None:
+            api = api_name
+
+        url = get_query_api_path(ver, api)
+
         return get(
             self.authenticator,
             session,
-            f'/services/data/v{ver}/query?q={soql}',
+            f'{url}?q={soql}',
             lambda response : response.json()
         )
 

src/newrelic_logging/query/__init__.py

@@ -30,11 +30,13 @@ def __init__(
         query: str,
         options: Config,
         api_ver: str = None,
+        api_name: str = None,
     ):
         self.api = api
         self.query = query
         self.options = options
         self.api_ver = api_ver
+        self.api_name = api_name
 
     def get(self, key: str, default = None):
         return self.options.get(key, default)
@@ -47,7 +49,12 @@ def execute(
         session: Session,
     ):
         print_info(f'Running query {self.query}...')
-        response = self.api.query(session, self.query, self.api_ver)
+        response = self.api.query(
+            session,
+            self.query,
+            self.api_ver,
+            self.api_name,
+        )
 
         if not is_valid_records_response(response):
             print_warn(f'no records returned for query {self.query}')
@@ -124,5 +131,6 @@ def new(
                 self.get_env(qp),
             ).replace(' ', '+'),
             Config(qp),
-            qp.get('api_ver', None)
+            qp.get('api_ver', None),
+            qp.get('api_name', None),
         )

src/tests/__init__.py

@@ -41,6 +41,7 @@ def __init__(
         self.limits_result = limits_result
         self.soql = None
         self.query_api_ver = None
+        self.query_api_name = None
         self.next_records_url = None
         self.limits_api_ver = None
         self.log_file_path = None
@@ -54,9 +55,16 @@ def authenticate(self, session: Session):
 
         self.authenticator.authenticate(session)
 
-    def query(self, session: Session, soql: str, api_ver: str = None) -> dict:
+    def query(
+        self,
+        session: Session,
+        soql: str,
+        api_ver: str = None,
+        api_name: str = None,
+    ) -> dict:
         self.soql = soql
         self.query_api_ver = api_ver
+        self.query_api_name = api_name
 
         if self.raise_error:
             raise SalesforceApiException()

src/tests/test_api.py

@@ -498,6 +498,49 @@ def test_stream_lines_uses_default_encoding_and_calls_iter_lines_with_chunk_size
         self.assertTrue(response.decode_unicode)
         self.assertTrue(response.iter_lines_called)
 
+    def test_get_query_api_path_returns_rest_api_path_given_api_ver_and_rest_api_name(self):
+        '''
+        get_query_api_path() returns the ReST API query path given an api version and the rest API name
+        given: an api version
+        and given: an api name
+        when: the api name is API_NAME_REST
+        then: return the ReST API query path with the given api name
+        '''
+
+        # execute
+        url = api.get_query_api_path('55.0', api.API_NAME_REST)
+
+        # verify
+        self.assertEqual(url, '/services/data/v55.0/query')
+
+    def test_get_query_api_path_returns_tooling_api_path_given_api_ver_and_tooling_api_name(self):
+        '''
+        get_query_api_path() returns the Tooling API query path given an api version and the tooling API name
+        given: an api version
+        and given: an api name
+        when: the api name is API_NAME_TOOLING
+        then: return the tooling API query path with the given api name
+        '''
+
+        # execute
+        url = api.get_query_api_path('55.0', api.API_NAME_TOOLING)
+
+        # verify
+        self.assertEqual(url, '/services/data/v55.0/tooling/query')
+
+    def test_get_query_api_path_raises_salesforce_api_exception_given_invalid_api_name(self):
+        '''
+        get_query_api_path() raises a SalesforceApiException given an invalid API name
+        given: an api version
+        and given: an api name
+        when: the api name is invalid
+        then: raise a SalesforceApiException
+        '''
+
+        # execute / verify
+        with self.assertRaises(SalesforceApiException) as _:
+            api.get_query_api_path('55.0', 'invalid')
+
     def test_authenticate_calls_authenticator_authenticate(self):
         '''
         authenticate() calls authenticate() on the backing authenticator
@@ -637,6 +680,83 @@ def test_query_requests_correct_url_with_access_token_given_api_version_and_retu
         self.assertTrue('foo' in resp)
         self.assertEqual(resp['foo'], 'bar')
 
+    def test_query_requests_correct_url_with_access_token_given_api_name_and_returns_json_response_on_success(self):
+        '''
+        query() calls the correct query API url with the access token when a specific api name is given and returns a JSON response
+        given: an authenticator
+        and given: an api version
+        and given: a session
+        and given: a query
+        when: query() is called
+        and when: the api name parameter is specified
+        then: session.get() is called with correct URL and access token
+        and: stream is set to False
+        and when: response status code is 200
+        then: calls callback with response and returns a JSON response
+        '''
+
+        # setup
+        auth = AuthenticatorStub(
+            instance_url='https://my.salesforce.test',
+            access_token='123456',
+        )
+        session = SessionStub()
+        session.response = ResponseStub(200, 'OK', '{"foo": "bar"}', [])
+
+        # execute
+        sf_api = api.Api(auth, '55.0')
+        resp = sf_api.query(
+            session,
+            'SELECT+LogFile+FROM+EventLogFile',
+            None,
+            api.API_NAME_TOOLING,
+        )
+
+        # verify
+
+        self.assertEqual(
+            session.url,
+            f'https://my.salesforce.test/services/data/v55.0/tooling/query?q=SELECT+LogFile+FROM+EventLogFile',
+        )
+        self.assertTrue('Authorization' in session.headers)
+        self.assertEqual(session.headers['Authorization'], 'Bearer 123456')
+        self.assertFalse(session.stream)
+        self.assertIsNotNone(resp)
+        self.assertTrue(type(resp) is dict)
+        self.assertTrue('foo' in resp)
+        self.assertEqual(resp['foo'], 'bar')
+
+    def test_query_raises_salesforce_api_exception_if_get_query_api_path_does(self):
+        '''
+        query() raises a SalesforceApiException if get_query_api_path() does
+        given: an authenticator
+        and given: an api version
+        and given: a session
+        and given: a query
+        when: query() is called
+        and when: get_query_api_path() raises a SalesforceApiException
+        then: query() raises a SalesforceApiException
+        '''
+
+        # setup
+        auth = AuthenticatorStub(
+            instance_url='https://my.salesforce.test',
+            access_token='123456',
+        )
+        session = SessionStub()
+        session.response = ResponseStub(200, 'OK', '{"foo": "bar"}', [])
+
+        # execute / verify
+        with self.assertRaises(SalesforceApiException) as _:
+            sf_api = api.Api(auth, '55.0')
+            sf_api.query(
+                session,
+                'SELECT+LogFile+FROM+EventLogFile',
+                None,
+                'invalid',
+            )
+
+
     def test_query_raises_login_exception_if_get_does(self):
         '''
         query() calls the correct query API url with the access token and raises LoginException if get does