acryl-datahub-cloud 0.3.8.3__py3-none-any.whl → 0.3.9__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.8.3",
+  "version": "0.3.9",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

acryl_datahub_cloud/acryl_cs_issues/acryl_customer.py

@@ -662,7 +662,7 @@ def generate_markdown_with_logo(
     logo_url = get_platform_logo_url(platform)
     formatted_date = date.strftime("%B %d, %Y")
     if logo_url:
-        markdown = f'<img src="{logo_url}" alt="Logo" width="{logo_size}" height="{logo_size}" style="vertical-align: middle;"> by {author.name} on {formatted_date} {""if is_public else "(Internal)"}'
+        markdown = f'<img src="{logo_url}" alt="Logo" width="{logo_size}" height="{logo_size}" style="vertical-align: middle;"> by {author.name} on {formatted_date} {"" if is_public else "(Internal)"}'
     else:
         markdown = f":{platform}: by {author.name} on {formatted_date} {'' if is_public else '(Internal)'}"
     return markdown

acryl_datahub_cloud/action_request/action_request_owner_source.py

@@ -0,0 +1,174 @@
+import logging
+from typing import Dict, Iterable, List, Optional
+
+from datahub.configuration import ConfigModel
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import Source, SourceReport
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import (
+    ActionRequestInfoClass,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+logger = logging.getLogger(__name__)
+
+
+class ActionRequestOwnerSourceConfig(ConfigModel):
+    batch_size: int = 20
+
+
+class ActionRequestOwnerSourceReport(SourceReport):
+    total_requests: int = 0
+    processed_proposals = 0
+    correct_assignees_not_found = 0
+    correct_proposal_owners = 0
+    incorrect_proposal_owners = 0
+    action_request_info_not_found = 0
+
+
+ACTION_REQUESTS = """
+query listActionRequests($input: ListActionRequestsInput!) {
+  listActionRequests(input: $input) {
+    start
+    count
+    total
+    actionRequests {
+      urn
+      type
+      entity {
+        urn
+      }
+      subResource
+      subResourceType
+      assignedUsers
+      assignedGroups
+      assignedRoles
+    }
+  }
+}
+"""
+
+ACTION_REQUEST_ASSIGNEES = """
+query getActionRequestAssignee($input: GetActionRequestAssigneeInput!) {
+  getActionRequestAssignee(input: $input)
+}
+"""
+
+
+@platform_name(id="datahub", platform_name="DataHub")
+@config_class(ActionRequestOwnerSourceConfig)
+@support_status(SupportStatus.INCUBATING)
+class ActionRequestOwnerSource(Source):
+    def __init__(self, config: ActionRequestOwnerSourceConfig, ctx: PipelineContext):
+        super().__init__(ctx)
+        self.config: ActionRequestOwnerSourceConfig = config
+        self.report = ActionRequestOwnerSourceReport()
+        self.graph = ctx.require_graph("Proposal Owner source")
+        self.event_not_produced_warn = False
+
+    def _process_action_request(
+        self, action_request: Dict
+    ) -> Optional[MetadataChangeProposalWrapper]:
+        self.report.processed_proposals += 1
+        action_request_urn = action_request.get("urn")
+        action_type = action_request.get("type")
+        action_request_entity = action_request.get("entity")
+        assert action_request_entity is not None
+        resource_urn = action_request_entity.get("urn")
+        sub_resource = action_request.get("subResource")
+        sub_resource_type = action_request.get("subResourceType")
+        assigned_users = action_request.get("assignedUsers")
+        assigned_groups = action_request.get("assignedGroups")
+        assigned_roles = action_request.get("assignedRoles")
+
+        correct_assignees = self.graph.execute_graphql(
+            query=ACTION_REQUEST_ASSIGNEES,
+            variables={
+                "input": {
+                    "resourceUrn": resource_urn,
+                    "actionRequestType": action_type,
+                    "subResource": sub_resource,
+                    "subResourceType": sub_resource_type,
+                }
+            },
+        ).get("getActionRequestAssignee")
+        if correct_assignees is None:
+            self.report.correct_assignees_not_found += 1
+            logger.error(
+                f"Correct assignees not found for action request {action_request_urn}"
+            )
+            return None
+        correct_users = [
+            x for x in correct_assignees if guess_entity_type(x) == "corpuser"
+        ]
+        correct_groups = [
+            x for x in correct_assignees if guess_entity_type(x) == "corpGroup"
+        ]
+        correct_roles = [
+            x for x in correct_assignees if guess_entity_type(x) == "dataHubRole"
+        ]
+        if (
+            assigned_users == correct_users
+            and assigned_groups == correct_groups
+            and assigned_roles == correct_roles
+        ):
+            self.report.correct_proposal_owners += 1
+            return None
+        action_request_info = self.graph.get_aspect_v2(
+            entity_urn=str(action_request_urn),
+            aspect="actionRequestInfo",
+            aspect_type=ActionRequestInfoClass,
+        )
+        if action_request_info is None:
+            self.report.action_request_info_not_found += 1
+            logger.error(
+                f"Action request info not found for action request {action_request_urn}"
+            )
+            return None
+        action_request_info.assignedUsers = correct_users
+        action_request_info.assignedGroups = correct_groups
+        action_request_info.assignedRoles = correct_roles
+        return MetadataChangeProposalWrapper(
+            entityUrn=action_request_urn, aspect=action_request_info
+        )
+
+    def _get_action_requests(self, start: int) -> List:
+        list_action_requests = self.graph.execute_graphql(
+            query=ACTION_REQUESTS,
+            variables={
+                "input": {
+                    "status": "PENDING",
+                    "allActionRequests": True,
+                    "start": start,
+                    "count": self.config.batch_size,
+                }
+            },
+        )
+        assert list_action_requests is not None
+        listActionRequests = list_action_requests.get("listActionRequests")
+        assert listActionRequests is not None
+        self.report.total_requests = listActionRequests.get("total", 0)
+        return listActionRequests.get("actionRequests", [])
+
+    def get_workunits(self) -> Iterable[MetadataWorkUnit]:
+        start = 0
+        while True:
+            action_requests = self._get_action_requests(start)
+            if len(action_requests) == 0:
+                break
+            for action_request in action_requests:
+                result = self._process_action_request(action_request)
+                if result is not None:
+                    yield result.as_workunit()
+            start += self.config.batch_size
+
+    def get_report(self) -> SourceReport:
+        return self.report

acryl_datahub_cloud/api/__init__.py

@@ -1 +1 @@
-from acryl_datahub_cloud.api.client import AcrylGraph
+from acryl_datahub_cloud.api.client import AcrylGraph as AcrylGraph

acryl_datahub_cloud/api/client.py

@@ -1,6 +1,6 @@
-from acryl_datahub_cloud.api.entity_versioning import EntityVersioningAPI
 from datahub.ingestion.graph.client import DataHubGraph
 
 
-class AcrylGraph(EntityVersioningAPI, DataHubGraph):
+# Add other graph mixins here when applicable
+class AcrylGraph(DataHubGraph):
     pass

acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -66,9 +66,9 @@ class DatasetRegistrationSpec(BaseModel):
 
 class FileStoreBackedDatasetConfig(ConfigModel):
     dataset_name: str
-    dataset_urn: Optional[
-        str
-    ] = None  # If not set, it will be generated from the dataset_name
+    dataset_urn: Optional[str] = (
+        None  # If not set, it will be generated from the dataset_name
+    )
     bucket_prefix: str
     store_platform: str = "s3"
     file_name: str = "data"
@@ -79,9 +79,9 @@ class FileStoreBackedDatasetConfig(ConfigModel):
     generate_presigned_url: bool = True
     presigned_url_expiry_days: int = 7
     dataset_registration_spec: DatasetRegistrationSpec = DatasetRegistrationSpec()
-    file: Optional[
-        str
-    ] = None  # This is the file to be registered. When set, the file will be registered as a dataset immediately
+    file: Optional[str] = (
+        None  # This is the file to be registered. When set, the file will be registered as a dataset immediately
+    )
 
     datahub_platform: str = "acryl"
 

acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -10,11 +10,13 @@ from pydantic import BaseModel
 from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
-from datahub.ingestion.graph.filters import SearchFilterRule
+from datahub.ingestion.graph.filters import RawSearchFilterRule
 from datahub.metadata.schema_classes import (
     DomainPropertiesClass,
+    FormAssociationClass,
     FormInfoClass,
     FormsClass,
+    FormStateClass,
     FormTypeClass,
 )
 
@@ -40,7 +42,7 @@ class FormType(str, Enum):
 class FormReportingRow(BaseModelRow):
     form_urn: str
     form_type: FormType
-    form_assigned_date: date
+    form_assigned_date: Optional[date]
     form_completed_date: Optional[date]
     form_status: FormStatus
     question_id: str
@@ -147,7 +149,7 @@ class DataHubFormReportingData(FormData):
             )
         )
 
-    def get_form_existence_or_filters(self) -> List[SearchFilterRule]:
+    def get_form_existence_or_filters(self) -> List[RawSearchFilterRule]:
         """
         Datasets must either have completedForms or incompleteForms assigned to
         them
@@ -177,24 +179,64 @@ class DataHubFormReportingData(FormData):
                 },
             ]
 
+    def is_published(self, form_urn: str) -> bool:
+        form_info = self.form_registry.get_form(form_urn)
+        return (
+            form_info.status.state == FormStateClass.PUBLISHED if form_info else False
+        )
+
+    def form_published_time(self, form_urn: str) -> float:
+        form_info = self.form_registry.get_form(form_urn)
+        is_published = (
+            form_info.status.state == FormStateClass.PUBLISHED if form_info else False
+        )
+        return (
+            form_info.status.lastModified.time / 1000
+            if form_info and form_info.status.lastModified and is_published
+            else 0
+        )
+
+    def assigned_to_asset_time(self, form_association: FormAssociationClass) -> float:
+        return form_association.created.time / 1000 if form_association.created else 0
+
+    def assignment_time(self, form_association: FormAssociationClass) -> float:
+        published_time = self.form_published_time(form_association.urn)
+        assigned_to_asset_time = self.assigned_to_asset_time(form_association)
+        return max(published_time, assigned_to_asset_time)
+
+    # For a given asset, the assigned date is the more recent of the published date and the time this was actually assigned.
+    # Assets can be assigned before publishing, but we don't want to show that date because it's not open to the public yet.
+    # Assets can also be assigned after publishing, so we should show that date for those assets.
     def form_assigned_date(
         self, search_row: DataHubDatasetSearchRow
-    ) -> Dict[str, date]:
-        form_assigned_dates: Dict[str, date] = {}
+    ) -> Dict[str, Optional[date]]:
+        form_assigned_dates: Dict[str, Optional[date]] = {}
         forms = self.graph.get_aspect(search_row.urn, FormsClass)
         if not forms:
             return form_assigned_dates
         assert forms, f"Forms aspect not found for {search_row.urn}"
         for incomplete_form in forms.incompleteForms:
-            form_assigned_dates[incomplete_form.urn] = datetime.fromtimestamp(
-                incomplete_form.created.time / 1000 if incomplete_form.created else 0,
-                tz=timezone.utc,
-            ).date()
+            is_published = self.is_published(incomplete_form.urn)
+            assignment_time = self.assignment_time(incomplete_form)
+            form_assigned_dates[incomplete_form.urn] = (
+                datetime.fromtimestamp(
+                    assignment_time,
+                    tz=timezone.utc,
+                ).date()
+                if is_published and assignment_time != 0
+                else None
+            )
         for completed_form in forms.completedForms:
-            form_assigned_dates[completed_form.urn] = datetime.fromtimestamp(
-                completed_form.created.time / 1000 if completed_form.created else 0,
-                tz=timezone.utc,
-            ).date()
+            is_published = self.is_published(completed_form.urn)
+            assignment_time = self.assignment_time(completed_form)
+            form_assigned_dates[completed_form.urn] = (
+                datetime.fromtimestamp(
+                    assignment_time,
+                    tz=timezone.utc,
+                ).date()
+                if is_published and assignment_time != 0
+                else None
+            )
         return form_assigned_dates
 
     def form_completed_date(
@@ -206,7 +248,9 @@ class DataHubFormReportingData(FormData):
         form_completion_dates = {}
         for form in search_row.completedForms:
             form_info = self.form_registry.get_form(form)
-            assert form_info, f"Form {form} not found"
+            if not form_info:
+                logger.warning(f"Found form attached that does not exist: {form}")
+                continue
             form_prompts = [x.id for x in form_info.prompts]
             completed_prompts_map = {
                 prompt_id: response_time
@@ -291,7 +335,11 @@ class DataHubFormReportingData(FormData):
                         on_form_scanned(form_id)
                     forms_scanned.add(form_id)
                 form_info = self.form_registry.get_form(form_id)
-                assert form_info, f"Form {form_id} not found"
+                if not form_info:
+                    logger.warning(
+                        f"Found form attached that does not exist: {form_id}"
+                    )
+                    continue
                 form_prompts = [x.id for x in form_info.prompts]
                 form_incomplete_prompts = [
                     p
@@ -315,11 +363,6 @@ class DataHubFormReportingData(FormData):
                     if p in form_prompts
                 ]:
                     for owner in assignees:
-                        if form_id not in form_assigned_dates:
-                            logger.warning(
-                                f"Form {form_id} not found in form_assigned_dates"
-                            )
-                            continue
                         yield FormReportingRow(
                             form_urn=form_id,
                             form_assigned_date=form_assigned_dates[form_id],
@@ -349,11 +392,6 @@ class DataHubFormReportingData(FormData):
                     if p in form_prompts
                 ]:
                     for owner in assignees:
-                        if form_id not in form_assigned_dates:
-                            logger.warning(
-                                f"Form {form_id} not found in form_assigned_dates"
-                            )
-                            continue
                         yield FormReportingRow(
                             form_urn=form_id,
                             form_assigned_date=form_assigned_dates[form_id],
@@ -393,7 +431,11 @@ class DataHubFormReportingData(FormData):
                         on_form_scanned(form_id)
                     forms_scanned.add(form_id)
                 form_info = self.form_registry.get_form(form_id)
-                assert form_info, f"Form {form_id} not found"
+                if not form_info:
+                    logger.warning(
+                        f"Found form attached that does not exist: {form_id}"
+                    )
+                    continue
                 form_type = (
                     FormType.DOCUMENTATION
                     if form_info.type == FormTypeClass.COMPLETION
@@ -411,11 +453,8 @@ class DataHubFormReportingData(FormData):
                 for prompt_id in [
                     p
                     for p in search_row.completedFormsIncompletePromptIds
-                    for p in form_prompts
+                    if p in form_prompts
                 ]:
-                    logger.warning(
-                        f"Unexpected incomplete prompt {prompt_id} in completed form {form_id}"
-                    )
                     for owner in assignees:
                         yield FormReportingRow(
                             form_urn=form_id,
@@ -450,11 +489,6 @@ class DataHubFormReportingData(FormData):
                     if p in form_prompts
                 ]:
                     for owner in assignees:
-                        if form_id not in form_assigned_dates:
-                            logger.warning(
-                                f"Form {form_id} not found in form_assigned_dates"
-                            )
-                            continue
                         yield FormReportingRow(
                             form_urn=form_id,
                             form_assigned_date=form_assigned_dates[form_id],

acryl_datahub_cloud/datahub_reporting/extract_sql.py

@@ -54,9 +54,9 @@ class DataHubReportingExtractSQLSourceConfig(ConfigModel):
 
             if "file" not in v:
                 default_config = FileStoreBackedDatasetConfig.dummy()
-                v[
-                    "file"
-                ] = f"{default_config.file_name}.{default_config.file_extension}"
+                v["file"] = (
+                    f"{default_config.file_name}.{default_config.file_extension}"
+                )
             else:
                 v["file_name"] = v["file"].split(".")[0]
                 v["file_extension"] = v["file"].split(".")[-1]
@@ -133,7 +133,7 @@ class DataHubReportingExtractSQLSource(Source):
         tmp_dir_aux = (
             self.ctx.pipeline_name if self.ctx.pipeline_name else "sql_default_dir"
         )
-        tmp_dir = f'/tmp/{tmp_dir_aux.replace(":", "_")}'
+        tmp_dir = f"/tmp/{tmp_dir_aux.replace(':', '_')}"
 
         output_file = (
             self.datahub_based_s3_dataset.config.file

acryl_datahub_cloud/datahub_usage_reporting/usage_feature_patch_builder.py

@@ -37,7 +37,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/queryCountLast30Days",
+            path=("queryCountLast30Days",),
             value=count,
         )
         return self
@@ -56,7 +56,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/viewCountLast30Days",
+            path=("viewCountLast30Days",),
             value=count,
         )
         return self
@@ -73,7 +73,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/viewCountTotal",
+            path=("viewCountTotal",),
             value=count,
         )
         return self
@@ -92,7 +92,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/viewCountPercentileLast30Days",
+            path=("viewCountPercentileLast30Days",),
             value=count,
         )
         return self
@@ -111,7 +111,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/usageCountLast30Days",
+            path=("usageCountLast30Days",),
             value=count,
         )
         return self
@@ -130,7 +130,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/uniqueUserCountLast30Days",
+            path=("uniqueUserCountLast30Days",),
             value=count,
         )
         return self
@@ -149,7 +149,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/writeCountLast30Days",
+            path=("writeCountLast30Days",),
             value=count,
         )
         return self
@@ -168,7 +168,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/queryCountPercentileLast30Days",
+            path=("queryCountPercentileLast30Days",),
             value=count,
         )
         return self
@@ -187,7 +187,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/queryCountRankLast30Days",
+            path=("queryCountRankLast30Days",),
             value=count,
         )
         return self
@@ -206,7 +206,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/uniqueUserPercentileLast30Days",
+            path=("uniqueUserPercentileLast30Days",),
             value=count,
         )
         return self
@@ -225,7 +225,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/uniqueUserRankLast30Days",
+            path=("uniqueUserRankLast30Days",),
             value=count,
         )
         return self
@@ -244,7 +244,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/writeCountPercentileLast30Days",
+            path=("writeCountPercentileLast30Days",),
             value=count,
         )
         return self
@@ -263,7 +263,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/writeCountRankLast30Days",
+            path=("writeCountRankLast30Days",),
             value=count,
         )
         return self
@@ -282,7 +282,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/topUsersLast30Days",
+            path=("topUsersLast30Days",),
             value=users,
         )
         return self
@@ -301,7 +301,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/sizeInBytesPercentile",
+            path=("sizeInBytesPercentile",),
             value=percentile,
         )
         return self
@@ -318,7 +318,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/sizeInBytesRank",
+            path=("sizeInBytesRank",),
             value=rank,
         )
         return self
@@ -337,7 +337,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/rowCountPercentile",
+            path=("rowCountPercentile",),
             value=percentile,
         )
         return self
@@ -356,7 +356,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/usageSearchScoreMultiplier",
+            path=("usageSearchScoreMultiplier",),
             value=multiplier,
         )
         return self
@@ -375,7 +375,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/usageFreshnessScoreMultiplier",
+            path=("usageFreshnessScoreMultiplier",),
             value=multiplier,
         )
         return self
@@ -394,7 +394,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/customDatahubScoreMultiplier",
+            path=("customDatahubScoreMultiplier",),
             value=multiplier,
         )
         return self
@@ -413,7 +413,7 @@ class UsageFeaturePatchBuilder(MetadataPatchProposal):
         self._add_patch(
             UsageFeaturesClass.ASPECT_NAME,
             "add",
-            path="/combinedSearchRankingMultiplier",
+            path=("combinedSearchRankingMultiplier",),
             value=multiplier,
         )
         return self