sift_py.report_templates.service

  1from __future__ import annotations
  2
  3from datetime import datetime
  4from pathlib import Path
  5from typing import Any, Dict, List, Optional, cast
  6
  7from google.protobuf.field_mask_pb2 import FieldMask
  8from sift.report_templates.v1.report_templates_pb2 import (
  9    CreateReportTemplateRequest,
 10    CreateReportTemplateRequestClientKeys,
 11    GetReportTemplateRequest,
 12    GetReportTemplateResponse,
 13    ReportTemplate,
 14    ReportTemplateRule,
 15    ReportTemplateTag,
 16    UpdateReportTemplateRequest,
 17)
 18from sift.report_templates.v1.report_templates_pb2_grpc import ReportTemplateServiceStub
 19
 20from sift_py._internal.time import to_timestamp_pb
 21from sift_py.grpc.transport import SiftChannel
 22from sift_py.report_templates.config import ReportTemplateConfig, ReportTemplateUpdate
 23from sift_py.yaml.report_templates import load_report_templates
 24
 25
 26class ReportTemplateService:
 27    """
 28    A service for managing report templates. Allows for creating, updating, and retrieving report
 29    templates.
 30    """
 31
 32    _report_template_service_stub: ReportTemplateServiceStub
 33
 34    def __init__(self, channel: SiftChannel):
 35        self._report_template_service_stub = ReportTemplateServiceStub(channel)
 36
 37    def create_or_update_report_template(
 38        self, config: ReportTemplateConfig, updates: Optional[ReportTemplateUpdate] = None
 39    ):
 40        """
 41        Create or update a report template via a ReportTemplateConfig. The config must contain a
 42        template_client_key, otherwise an exception will be raised. If a report template with the
 43        same client key exists, it will be updated. Otherwise, a new report template will be created.
 44        See `sift_py.report_templates.config.ReportTemplateConfig` for more information on available
 45        fields to configure.
 46        """
 47        if updates:
 48            if not config.template_client_key and not config.template_id:
 49                raise Exception(
 50                    f"Report template {config.name} requires either a template_client_key or report_template_id to update."
 51                )
 52            self._update_report_template(config, updates)
 53        else:
 54            self._create_report_template(config)
 55
 56    def get_report_template(
 57        self, client_key: Optional[str] = None, id: Optional[str] = None
 58    ) -> Optional[ReportTemplateConfig]:
 59        """
 60        Retrieve a report template by client key or id. Only one of client_key or id should be
 61        provided, otherwise an exception will be raised. If a report template is found, it will be
 62        returned as a ReportTemplateConfig object.
 63        """
 64        report_template = None
 65        if (not client_key and not id) or (client_key and id):
 66            raise ValueError("One of client_key or id must be provided")
 67
 68        if id:
 69            report_template = self._get_report_template_by_id(id)
 70        elif client_key:
 71            report_template = self._get_report_template_by_client_key(client_key)
 72
 73        return (
 74            ReportTemplateConfig(
 75                name=report_template.name,
 76                template_id=report_template.report_template_id,
 77                template_client_key=report_template.client_key,
 78                tags=[tag.tag_name for tag in report_template.tags],
 79                description=report_template.description,
 80                rule_client_keys=[rule.client_key for rule in report_template.rules],
 81                archived_date=report_template.archived_date.ToDatetime()
 82                if report_template.archived_date
 83                else None,
 84                archived=True if report_template.archived_date else False,
 85            )
 86            if report_template
 87            else None
 88        )
 89
 90    def load_report_templates_from_yaml(self, paths: List[Path]) -> List[ReportTemplateConfig]:
 91        """
 92        Load report templates from YAML definitions. The YAML defined report template must have
 93        a client key. If the report template with the given client key exists, it will be updated,
 94        otherwise a new report template will be created.
 95        See `sift_py.yaml.report_templates.load_report_templates` for more information on the YAML
 96        spec for report templates.
 97        """
 98        report_templates = load_report_templates(paths)
 99        [
100            self.create_or_update_report_template(report_template)
101            for report_template in report_templates
102        ]
103        return report_templates
104
105    def _get_report_template_by_id(self, report_template_id: str) -> Optional[ReportTemplate]:
106        req = GetReportTemplateRequest(report_template_id=report_template_id)
107        try:
108            res = cast(
109                GetReportTemplateResponse, self._report_template_service_stub.GetReportTemplate(req)
110            )
111            return cast(ReportTemplate, res.report_template) or None
112        except:
113            return None
114
115    def _get_report_template_by_client_key(self, client_key: str) -> Optional[ReportTemplate]:
116        req = GetReportTemplateRequest(client_key=client_key)
117        try:
118            res = cast(
119                GetReportTemplateResponse, self._report_template_service_stub.GetReportTemplate(req)
120            )
121            return res.report_template or None
122        except:
123            return None
124
125    def _create_report_template(self, config: ReportTemplateConfig):
126        client_keys_req = CreateReportTemplateRequestClientKeys(
127            rule_client_keys=config.rule_client_keys
128        )
129        req = CreateReportTemplateRequest(
130            name=config.name,
131            client_key=config.template_client_key,
132            description=config.description,
133            tag_names=config.tags,
134            rule_client_keys=client_keys_req,
135        )
136        self._report_template_service_stub.CreateReportTemplate(req)
137
138    def _update_report_template(self, config: ReportTemplateConfig, updates: ReportTemplateUpdate):
139        """
140        Updates a report template using the provided updates. Only the fields specified in the updates
141        will be modified.
142        """
143        update_map: Dict[str, Any] = {}
144        if "name" in updates:
145            update_map["name"] = updates["name"]
146        if "description" in updates:
147            update_map["description"] = updates["description"]
148        if "tags" in updates:
149            update_map["tags"] = [ReportTemplateTag(tag_name=tag) for tag in updates["tags"]]
150        if "rule_client_keys" in updates:
151            update_map["rules"] = [
152                ReportTemplateRule(client_key=client_key)
153                for client_key in updates["rule_client_keys"]
154            ]
155        if "archived" in updates:
156            update_map["archived_date"] = (
157                to_timestamp_pb(datetime.now()) if updates["archived"] else None
158            )
159
160        updated_report_template = ReportTemplate(
161            report_template_id=config.template_id or "",
162            client_key=config.template_client_key,
163            name=update_map.get("name", config.name),
164            description=update_map.get("description", config.description),
165            tags=update_map.get("tags", config.tags),
166            rules=update_map.get(
167                "rules", [ReportTemplateRule(client_key=key) for key in config.rule_client_keys]
168            ),
169            archived_date=update_map.get("archived_date", config.archived_date),
170        )
171
172        field_mask = FieldMask(paths=list(update_map.keys()))
173        self._report_template_service_stub.UpdateReportTemplate(
174            UpdateReportTemplateRequest(
175                report_template=updated_report_template, update_mask=field_mask
176            )
177        )