|
| 1 | +import logging |
| 2 | + |
| 3 | +from typing import Optional, Dict, Any, Union, Iterator # noqa: F401 |
| 4 | + |
| 5 | +from databuilder.models.dashboard.dashboard_metadata import DashboardMetadata |
| 6 | +from databuilder.models.neo4j_csv_serde import ( |
| 7 | + Neo4jCsvSerializable, RELATION_START_KEY, RELATION_END_KEY, RELATION_START_LABEL, |
| 8 | + RELATION_END_LABEL, RELATION_TYPE, RELATION_REVERSE_TYPE) |
| 9 | +from databuilder.models.usage.usage_constants import ( |
| 10 | + READ_RELATION_TYPE, READ_REVERSE_RELATION_TYPE, READ_RELATION_COUNT_PROPERTY |
| 11 | +) |
| 12 | +from databuilder.models.user import User |
| 13 | + |
| 14 | +LOGGER = logging.getLogger(__name__) |
| 15 | + |
| 16 | + |
| 17 | +class DashboardUsage(Neo4jCsvSerializable): |
| 18 | + """ |
| 19 | + A model that encapsulate Dashboard usage between Dashboard and User |
| 20 | + """ |
| 21 | + |
| 22 | + def __init__(self, |
| 23 | + dashboard_group_id, # type: Optional[str] |
| 24 | + dashboard_id, # type: Optional[str] |
| 25 | + email, # type: str |
| 26 | + view_count, # type: int |
| 27 | + should_create_user_node=False, # type: Optional[bool] |
| 28 | + product='', # type: Optional[str] |
| 29 | + cluster='gold', # type: Optional[str] |
| 30 | + **kwargs |
| 31 | + ): |
| 32 | + # type: () -> None |
| 33 | + """ |
| 34 | +
|
| 35 | + :param dashboard_group_id: |
| 36 | + :param dashboard_id: |
| 37 | + :param email: |
| 38 | + :param view_count: |
| 39 | + :param should_create_user_node: Enable this if it is fine to create/update User node with only with email |
| 40 | + address. Please be advised that other fields will be emptied. Current use case is to create anonymous user. |
| 41 | + For example, Mode dashboard does not provide which user viewed the dashboard and anonymous user can be used |
| 42 | + to show the usage. |
| 43 | + :param product: |
| 44 | + :param cluster: |
| 45 | + :param kwargs: |
| 46 | + """ |
| 47 | + self._dashboard_group_id = dashboard_group_id |
| 48 | + self._dashboard_id = dashboard_id |
| 49 | + self._email = email |
| 50 | + self._view_count = view_count |
| 51 | + self._product = product |
| 52 | + self._cluster = cluster |
| 53 | + self._user_model = User(email=email) |
| 54 | + self._should_create_user_node = bool(should_create_user_node) |
| 55 | + self._relation_iterator = self._create_relation_iterator() |
| 56 | + |
| 57 | + def create_next_node(self): |
| 58 | + # type: () -> Union[Dict[str, Any], None] |
| 59 | + if self._should_create_user_node: |
| 60 | + return self._user_model.create_next_node() |
| 61 | + |
| 62 | + def create_next_relation(self): |
| 63 | + # type: () -> Union[Dict[str, Any], None] |
| 64 | + try: |
| 65 | + return next(self._relation_iterator) |
| 66 | + except StopIteration: |
| 67 | + return None |
| 68 | + |
| 69 | + def _create_relation_iterator(self): |
| 70 | + # type: () -> Iterator[[Dict[str, Any]]] |
| 71 | + |
| 72 | + yield { |
| 73 | + RELATION_START_LABEL: DashboardMetadata.DASHBOARD_NODE_LABEL, |
| 74 | + RELATION_END_LABEL: User.USER_NODE_LABEL, |
| 75 | + RELATION_START_KEY: DashboardMetadata.DASHBOARD_KEY_FORMAT.format( |
| 76 | + product=self._product, |
| 77 | + cluster=self._cluster, |
| 78 | + dashboard_group=self._dashboard_group_id, |
| 79 | + dashboard_name=self._dashboard_id |
| 80 | + ), |
| 81 | + RELATION_END_KEY: User.get_user_model_key(email=self._email), |
| 82 | + RELATION_TYPE: READ_REVERSE_RELATION_TYPE, |
| 83 | + RELATION_REVERSE_TYPE: READ_RELATION_TYPE, |
| 84 | + READ_RELATION_COUNT_PROPERTY: self._view_count |
| 85 | + } |
| 86 | + |
| 87 | + def __repr__(self): |
| 88 | + return 'DashboardUsage({!r}, {!r}, {!r}, {!r}, {!r}, {!r}, {!r})'.format( |
| 89 | + self._dashboard_group_id, |
| 90 | + self._dashboard_id, |
| 91 | + self._email, |
| 92 | + self._view_count, |
| 93 | + self._should_create_user_node, |
| 94 | + self._product, |
| 95 | + self._cluster |
| 96 | + ) |
0 commit comments