BRONTO-1812 refactor and add CI automation of tests

.github/workflows/build.yml

@@ -0,0 +1,30 @@
+name: Build
+
+on:
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  id-token: write
+  contents: write
+  packages: write
+  deployments: write
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    env:
+      VERSION: "1.0.${{ github.run_number }}-${{ github.run_attempt }}"
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up environment
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Test
+        run: ./test.sh

.gitignore

@@ -51,6 +51,10 @@ terraform.rc
 
 # Python virtual envs
 **/env/
+test_env/
+
+# Coverage
+**/.coverage
 
 # Minikube
 **/minikube-darwin-arm64

requirements.txt

@@ -1,4 +1 @@
-mcp[cli]
-zstd
-
-pytest
+mcp[cli]==1.13.1

src/main/brmcpserver/clients.py

@@ -1,5 +1,3 @@
-import time
-
 import urllib.request
 import logging
 import os
@@ -178,20 +176,6 @@ def search_post(self, timestamp_start: int, timestamp_end: int, log_ids: list[st
             logger.error('Cannot interact with Bronto', exc_info=True)
             raise Exception('Cannot interact with Bronto. Please check endpoint configuration.')
 
-    def get_recent_keys(self, log_id) -> Dict[str, List[str]]:
-        now = int(time.time()) * 1000
-        ten_minutes_ago = now - 10 * 60 * 1000
-        log_events: List[LogEvent] = self.search(ten_minutes_ago, now, [log_id], _select = ['*', '@raw'])
-        keys_and_values: Dict[str, List[str]] = {}
-        for event in log_events:
-            for key in event.attributes:
-                if key in keys_and_values:
-                    keys_and_values[key].append(event.attributes[key])
-                else:
-                    keys_and_values[key] = [event.attributes[key]]
-        keys_and_unique_values = {key: list(set(keys_and_values[key])) for key in keys_and_values}
-        return keys_and_unique_values
-
     def get_top_keys(self, log_id) -> Dict[str, List[str]]:
         url_path = f'top-keys?log_id={log_id}'
         request = urllib.request.Request(os.path.join(self.api_endpoint, url_path), headers=self.headers)
@@ -234,28 +218,103 @@ def get_top_keys(self, log_id) -> Dict[str, List[str]]:
             logger.error('Cannot interact with Bronto', exc_info=True)
             raise Exception('Cannot interact with Bronto. Please check endpoint configuration.')
 
+    def get_all_datasets_top_keys(self) -> Dict[str, List[str]]:
+        url_path = f'top-keys'
+        request = urllib.request.Request(os.path.join(self.api_endpoint, url_path), headers=self.headers)
+        try:
+            with urllib.request.urlopen(request) as resp:
+                if resp.status != 200 and resp.status != 201:
+                    logger.error('Keys retrieval failed, status=%s, reason=%s',resp.status, resp.reason)
+                    raise FailedBrontoRequestException(f'Cannot retrieve top keys from Bronto. status={resp.status}, '
+                                                       f'reason="{resp.reason}"')
+                try:
+                    body = json.loads(resp.read())
+                except json.decoder.JSONDecodeError as _:
+                    logger.error('Cannot decode search response', exc_info=True)
+                    raise BrontoResponseDecodingException('Unexpected format for retrieved data')
+
+                log_ids_and_keys: Dict[str, List[str]]  = {}
+                for log_id in body:
+                    if log_id not in log_ids_and_keys:
+                        log_ids_and_keys[log_id] = []
+                    log_ids_and_keys[log_id].extend(body[log_id].keys())
+                logging.info('log_ids_and_keys=%s', log_ids_and_keys)
+                return log_ids_and_keys
+        except (FailedBrontoRequestException, BrontoResponseDecodingException) as e:
+            raise e
+        except HTTPError as e:
+            if e.code == 400:
+                raise BrontoResponseException('One of the search parameters is unsuitable. Check the filter syntax as '
+                                              'well as the names of the keys used in the "where", "_select" and '
+                                              '"group_by_keys" parameters.')
+            if e.code == 403:
+                raise BrontoResponseException('You are not allowed to perform this Bronto search. Please check your '
+                                              'Bronto API key')
+            if e.code == 401:
+                raise BrontoResponseException('You are not authorised to perform this Bronto search. Please check your '
+                                              'Bronto API key, as well as the Bronto endpoint, to make sure that they '
+                                              'match')
+        except Exception as _:
+            logger.error('Cannot interact with Bronto', exc_info=True)
+            raise Exception('Cannot interact with Bronto. Please check endpoint configuration.')
+
+    def get_all_datasets_top_keys_and_values(self) -> Dict[str, Dict[str, List[str]]]:
+        url_path = f'top-keys'
+        request = urllib.request.Request(os.path.join(self.api_endpoint, url_path), headers=self.headers)
+        try:
+            with urllib.request.urlopen(request) as resp:
+                if resp.status != 200 and resp.status != 201:
+                    logger.error('Keys retrieval failed, status=%s, reason=%s',resp.status, resp.reason)
+                    raise FailedBrontoRequestException(f'Cannot retrieve top keys from Bronto. status={resp.status}, '
+                                                       f'reason="{resp.reason}"')
+                try:
+                    body = json.loads(resp.read())
+                except json.decoder.JSONDecodeError as _:
+                    logger.error('Cannot decode search response', exc_info=True)
+                    raise BrontoResponseDecodingException('Unexpected format for retrieved data')
+
+                log_ids_and_keys_and_values: Dict[str, Dict[str, List[str]]]  = {}
+                for log_id in body:
+                    if log_id not in log_ids_and_keys_and_values:
+                        log_ids_and_keys_and_values[log_id] = {}
+                    log_ids_and_keys_and_values[log_id].update({key: [value for value in body[log_id][key]['values'].keys()] for key in body[log_id]})
+                logging.info('log_ids_and_keys_and_values=%s', log_ids_and_keys_and_values)
+                return log_ids_and_keys_and_values
+        except (FailedBrontoRequestException, BrontoResponseDecodingException) as e:
+            raise e
+        except HTTPError as e:
+            if e.code == 400:
+                raise BrontoResponseException('One of the search parameters is unsuitable. Check the filter syntax as '
+                                              'well as the names of the keys used in the "where", "_select" and '
+                                              '"group_by_keys" parameters.')
+            if e.code == 403:
+                raise BrontoResponseException('You are not allowed to perform this Bronto search. Please check your '
+                                              'Bronto API key')
+            if e.code == 401:
+                raise BrontoResponseException('You are not authorised to perform this Bronto search. Please check your '
+                                              'Bronto API key, as well as the Bronto endpoint, to make sure that they '
+                                              'match')
+        except Exception as _:
+            logger.error('Cannot interact with Bronto', exc_info=True)
+            raise Exception('Cannot interact with Bronto. Please check endpoint configuration.')
+
+
     @staticmethod
-    def _get_dataset_key(key_name: str, dataset_keys: List[DatasetKey]) -> Optional[DatasetKey]:
+    def get_dataset_key(key_name: str, dataset_keys: List[DatasetKey]) -> Optional[DatasetKey]:
         for dataset_key in dataset_keys:
             if dataset_key.name == key_name:
                 return dataset_key
         return None
 
     def get_keys(self, log_id) -> List[DatasetKey]:
-        recent_keys = self.get_recent_keys(log_id)
         top_keys = self.get_top_keys(log_id)
         result = []
         processed_keys = set()
-        for key in recent_keys:
-            if key in processed_keys:
-                dataset = BrontoClient._get_dataset_key(key, result)
-                dataset.add_values(recent_keys[key])
-            else:
-                result.append(DatasetKey(name=key, values=recent_keys[key]))
         for key in top_keys:
             if key in processed_keys:
-                dataset = BrontoClient._get_dataset_key(key, result)
+                dataset = BrontoClient.get_dataset_key(key, result)
                 dataset.add_values(top_keys[key])
             else:
                 result.append(DatasetKey(name=key, values=top_keys[key]))
+                processed_keys.add(key)
         return result

src/main/brmcpserver/main.py

@@ -21,8 +21,23 @@
         streamable_http_path="/",
         json_response=True,
         instructions='Use this MCP server to interact with log data stored in Bronto as well as its metadata, datasets, '
-                     'collections, keys, etc',
-        dependencies=['pydantic', 'zstd']
+                     'collections, keys, etc. The tools provided by this server help '
+                     '- selecting datasets based on their tags or the keys that their data contains For instance, '
+                     '  - when looking to select datasets based on the key they contain, the following process should be used:'
+                     '    - list the datasets and check their tags. Select datasets whose tags match the keys under interest.'
+                     '    - also list the keys of all datasets and select datasets that contain the key under interest.'
+                     '  - when looking to select datasets based on the value of a key that they contain, then the '
+                     '    following process should be used:'
+                     '    - list the datasets and check their tags. Select datasets whose tags and values match the keys '
+                     '      and values under interest.'
+                     '    - also always list the keys of all datasets to look for keys that is relevant to the one under '
+                     '      interest. Then retrieve the dataset IDs for datasets that contain at least one of those keys '
+                     '      and so that the value of the key matches the provided value.'
+                     '- searching log events in datasets, based on some filter. Filter are typically based on some of '
+                     'keys and values that the dataset contains.'
+                     '- computing metrics present in datasets, based on some filter, and grouping the results according '
+                     'to some of the keys present in the dataset',
+        dependencies=['pydantic']
     )
     config = Config()
     bronto_client = BrontoClient(config.bronto_api_key, config.bronto_api_endpoint)

src/main/brmcpserver/tools.py

@@ -5,7 +5,7 @@
 from typing_extensions import Annotated
 from datetime import datetime, timezone
 from typing import List, Optional, Dict
-from models import Dataset, DatasetKey, LogEvent, Datapoint, Timeseries
+from models import Dataset, LogEvent, Datapoint, Timeseries
 
 from clients import BrontoClient
 
@@ -70,7 +70,24 @@ def register(self, mcp):
                 which provides details on datasets.
                 This tool returns a list strings. Each string provides the name of a key present in the provided dataset
                 """
-        )(self.get_keys)
+        )(self.get_dataset_keys)
+
+        mcp.tool(
+            name='get_all_datasets_keys',
+            description="""Fetches all keys present in all datasets.
+                This tool returns a list of strings. Each string provides the name of a key present in the provided 
+                dataset. This tool is useful in cases such as:
+                - to select datasets the contain certain keys
+                - to identify the exact key name based on some description 
+                """
+        )(self.get_all_datasets_keys)
+
+        mcp.tool(
+            name='get_key_values',
+            description="""Fetches the values of the provided key and dataset ID.
+                This tool returns a list of strings. Each string provides the value of the key provided as input, for 
+                the dataset provided as input."""
+        )(self.get_key_values)
 
         mcp.tool(
             name='get_current_time',
@@ -183,13 +200,12 @@ def get_timestamp_as_unix_epoch(
             input_time: Annotated[
                 str,
                 BeforeValidator(_validate_input_time),
-                Field(description='Time represented in the "%Y-%m-%d %H:%M:%S" format')]
+                Field(description='Time represented in the "%Y-%m-%d %H:%M:%S" format. Timezone is assumed to be UTC')]
     ) -> Annotated[
         int,
         Field(description='A unix timestamp (in milliseconds) since epoch, representing the `input_time` parameter')
     ]:
-        return int(datetime.strptime(input_time, '%Y-%m-%d %H:%M:%S').astimezone(timezone.utc).timestamp()) * 1000
-
+        return int(datetime.strptime(input_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp()) * 1000
 
     def get_datasets(self) -> Annotated[
         List[Dataset],
@@ -208,7 +224,6 @@ def get_datasets(self) -> Annotated[
                                   tags=dataset["tags"]))
         return result
 
-
     def get_datasets_by_name(
             self,
             dataset_name: Annotated[str, Field(description="The dataset name", min_length=1)],
@@ -233,17 +248,34 @@ def get_datasets_by_name(
             return []
         return result
 
-
-    def get_keys(
+    def get_dataset_keys(
             self,
             log_id: Annotated[str, Field(description='The dataset ID, also named log ID', min_length=36, max_length=36)]
     ) -> Annotated[
-        List[DatasetKey],
+        List[str],
         Field(description='list key names for keys present in the provided dataset referenced with the `log_id` parameter')
     ]:
-        keys = self.bronto_client.get_keys(log_id)
+        keys = [dataset.name for dataset in self.bronto_client.get_keys(log_id)]
+        return keys
+
+    def get_all_datasets_keys(self) -> Annotated[
+        Dict[str, List[str]],
+        Field(description='Map from dataset IDs to the list of key names, for keys present in each dataset')
+    ]:
+        keys = self.bronto_client.get_all_datasets_top_keys()
         return keys
 
+    def get_key_values(
+            self,
+            key: Annotated[str, Field(description='The name of a key')],
+            log_id: Annotated[str, Field(description='A string representing a dataset ID')]
+    ) -> Annotated[List[str], Field(description='The list of values of the provided key, present in the provided '
+                                                'dataset.')]:
+        datasets_top_keys_and_values = self.bronto_client.get_all_datasets_top_keys_and_values()
+        keys_and_values = datasets_top_keys_and_values.get(log_id, {})
+        key_and_values = keys_and_values.get(key, {})
+        return key_and_values.get('values', {}).get(key, [])
+
     @staticmethod
     def get_current_time() -> Annotated[str, Field(description='the current time in the YYYY-MM-DD HH:mm:ss format')]:
         return datetime.strftime(datetime.now(), '%Y-%m-%d %H:%M:%S')