acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/unity/tag_entities.py

@@ -0,0 +1,197 @@
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.source.unity.platform_resource_repository import (
+        UnityCatalogPlatformResourceRepository,
+    )
+
+from pydantic import BaseModel
+
+from datahub.api.entities.external.external_entities import (
+    ExternalEntity,
+    ExternalEntityId,
+    LinkedResourceSet,
+)
+from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
+from datahub.api.entities.platformresource.platform_resource import (
+    PlatformResource,
+    PlatformResourceKey,
+)
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.urns import TagUrn
+from datahub.utilities.urns.urn import Urn
+
+
+class UnityCatalogTagSyncContext(BaseModel):
+    # it is intentionally empty
+    platform_instance: Optional[str] = None
+
+
+logger = logging.getLogger(__name__)
+
+
+class UnityCatalogTagPlatformResourceId(ExternalEntityId):
+    """
+    A Unity Catalog tag platform resource ID.
+    """
+
+    tag_key: str
+    tag_value: Optional[str] = None
+    platform_instance: Optional[str] = None
+    exists_in_unity_catalog: bool = False
+    persisted: bool = False
+
+    # this is a hack to make sure the property is a string and not private pydantic field
+    @staticmethod
+    def _RESOURCE_TYPE() -> str:
+        return "UnityCatalogTagPlatformResource"
+
+    def to_platform_resource_key(self) -> PlatformResourceKey:
+        return PlatformResourceKey(
+            platform="databricks",
+            resource_type=str(UnityCatalogTagPlatformResourceId._RESOURCE_TYPE()),
+            primary_key=f"{self.tag_key}:{self.tag_value}",
+            platform_instance=self.platform_instance,
+        )
+
+    @classmethod
+    def get_or_create_from_tag(
+        cls,
+        tag: UnityCatalogTag,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
+        exists_in_unity_catalog: bool = False,
+    ) -> "UnityCatalogTagPlatformResourceId":
+        """
+        Creates a UnityCatalogTagPlatformResourceId from a UnityCatalogTag.
+        """
+
+        existing_platform_resource = platform_resource_repository.search_entity_by_urn(
+            tag.to_datahub_tag_urn().urn()
+        )
+        if existing_platform_resource:
+            logger.debug(
+                f"Found existing UnityCatalogTagPlatformResourceId for tag {tag.key.raw_text}: {existing_platform_resource}"
+            )
+            return existing_platform_resource
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=tag.key.raw_text,
+            tag_value=tag.value.raw_text if tag.value is not None else None,
+            platform_instance=platform_resource_repository.platform_instance,
+            exists_in_unity_catalog=exists_in_unity_catalog,
+            persisted=False,
+        )
+
+    @classmethod
+    def from_datahub_urn(
+        cls,
+        urn: str,
+        tag_sync_context: UnityCatalogTagSyncContext,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
+        graph: DataHubGraph,
+    ) -> "UnityCatalogTagPlatformResourceId":
+        """
+        Creates a UnityCatalogTagPlatformResourceId from a DataHub URN.
+        """
+        existing_platform_resource_id = (
+            platform_resource_repository.search_entity_by_urn(urn)
+        )
+        if existing_platform_resource_id:
+            return existing_platform_resource_id
+
+        new_unity_catalog_tag_id = cls.generate_tag_id(graph, tag_sync_context, urn)
+        if new_unity_catalog_tag_id:
+            resource_key = platform_resource_repository.get(
+                new_unity_catalog_tag_id.to_platform_resource_key()
+            )
+            if resource_key:
+                # Create a new ID with the correct state instead of mutating
+                return UnityCatalogTagPlatformResourceId(
+                    tag_key=new_unity_catalog_tag_id.tag_key,
+                    tag_value=new_unity_catalog_tag_id.tag_value,
+                    platform_instance=new_unity_catalog_tag_id.platform_instance,
+                    exists_in_unity_catalog=True,  # This tag exists in Unity Catalog
+                    persisted=new_unity_catalog_tag_id.persisted,
+                )
+            return new_unity_catalog_tag_id
+        raise ValueError(
+            f"Unable to create Unity Catalog tag ID from DataHub URN: {urn}"
+        )
+
+    @classmethod
+    def generate_tag_id(
+        cls, graph: DataHubGraph, tag_sync_context: UnityCatalogTagSyncContext, urn: str
+    ) -> "UnityCatalogTagPlatformResourceId":
+        parsed_urn = Urn.from_string(urn)
+        entity_type = parsed_urn.entity_type
+        if entity_type == "tag":
+            return UnityCatalogTagPlatformResourceId.from_datahub_tag(
+                TagUrn.from_string(urn), tag_sync_context
+            )
+        else:
+            raise ValueError(f"Unsupported entity type {entity_type} for URN {urn}")
+
+    @classmethod
+    def from_datahub_tag(
+        cls, tag_urn: TagUrn, tag_sync_context: UnityCatalogTagSyncContext
+    ) -> "UnityCatalogTagPlatformResourceId":
+        uc_tag = UnityCatalogTag.from_urn(tag_urn)
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=str(uc_tag.key),
+            tag_value=str(uc_tag.value) if uc_tag.value is not None else None,
+            platform_instance=tag_sync_context.platform_instance,
+            exists_in_unity_catalog=False,
+        )
+
+
+class UnityCatalogTagPlatformResource(ExternalEntity):
+    datahub_urns: LinkedResourceSet
+    managed_by_datahub: bool
+    id: UnityCatalogTagPlatformResourceId
+    allowed_values: Optional[List[str]] = None
+
+    def get_id(self) -> ExternalEntityId:
+        return self.id
+
+    def is_managed_by_datahub(self) -> bool:
+        return self.managed_by_datahub
+
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        return self.datahub_urns
+
+    def as_platform_resource(self) -> PlatformResource:
+        return PlatformResource.create(
+            key=self.id.to_platform_resource_key(),
+            secondary_keys=[u for u in self.datahub_urns.urns],
+            value=self,
+        )
+
+    @classmethod
+    def create_default(
+        cls,
+        entity_id: ExternalEntityId,
+        managed_by_datahub: bool,
+    ) -> "UnityCatalogTagPlatformResource":
+        """Create a default Unity Catalog tag entity when none found in DataHub."""
+        # Type narrowing: we know this will be a UnityCatalogTagPlatformResourceId
+        assert isinstance(entity_id, UnityCatalogTagPlatformResourceId), (
+            f"Expected UnityCatalogTagPlatformResourceId, got {type(entity_id)}"
+        )
+
+        # Create a new entity ID with correct default state instead of mutating
+        default_entity_id = UnityCatalogTagPlatformResourceId(
+            tag_key=entity_id.tag_key,
+            tag_value=entity_id.tag_value,
+            platform_instance=entity_id.platform_instance,
+            exists_in_unity_catalog=False,  # New entities don't exist in Unity Catalog yet
+            persisted=False,  # New entities are not persisted yet
+        )
+
+        return cls(
+            id=default_entity_id,
+            datahub_urns=LinkedResourceSet(urns=[]),
+            managed_by_datahub=managed_by_datahub,
+            allowed_values=None,
+        )

datahub/ingestion/source/unity/usage.py

@@ -11,7 +11,10 @@ from databricks.sdk.service.sql import QueryStatementType
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.source_helpers import auto_empty_dataset_usage_statistics
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.unity.config import UnityCatalogSourceConfig
+from datahub.ingestion.source.unity.config import (
+    UnityCatalogSourceConfig,
+    UsageDataSource,
+)
 from datahub.ingestion.source.unity.proxy import UnityCatalogApiProxy
 from datahub.ingestion.source.unity.proxy_types import (
     OPERATION_STATEMENT_TYPES,
@@ -164,11 +167,50 @@ class UnityCatalogUsageExtractor:
                     aspect=operation_aspect,
                 ).as_workunit()
 
+    def _validate_usage_data_source_config(self) -> None:
+        """Validate usage data source configuration before execution."""
+        usage_data_source = self.config.usage_data_source
+
+        if (
+            usage_data_source == UsageDataSource.SYSTEM_TABLES
+            and not self.proxy.warehouse_id
+        ):
+            raise ValueError(
+                "usage_data_source is set to SYSTEM_TABLES but warehouse_id is not configured. "
+                "Either set warehouse_id or use AUTO/API mode."
+            )
+
     def _get_queries(self) -> Iterable[Query]:
         try:
-            yield from self.proxy.query_history(
-                self.config.start_time, self.config.end_time
-            )
+            self._validate_usage_data_source_config()
+            usage_data_source = self.config.usage_data_source
+
+            if usage_data_source == UsageDataSource.AUTO:
+                if self.proxy.warehouse_id:
+                    logger.info(
+                        "Using system tables for usage query history (AUTO mode)"
+                    )
+                    yield from self.proxy.get_query_history_via_system_tables(
+                        self.config.start_time, self.config.end_time
+                    )
+                else:
+                    logger.info(
+                        "Using API for usage query history (AUTO mode, no warehouse)"
+                    )
+                    yield from self.proxy.query_history(
+                        self.config.start_time, self.config.end_time
+                    )
+            elif usage_data_source == UsageDataSource.SYSTEM_TABLES:
+                logger.info("Using system tables for usage query history (forced)")
+                yield from self.proxy.get_query_history_via_system_tables(
+                    self.config.start_time, self.config.end_time
+                )
+            elif usage_data_source == UsageDataSource.API:
+                logger.info("Using API for usage query history (forced)")
+                yield from self.proxy.query_history(
+                    self.config.start_time, self.config.end_time
+                )
+
         except Exception as e:
             logger.warning("Error getting queries", exc_info=True)
             self.report.report_warning("get-queries", str(e))

datahub/ingestion/source/usage/clickhouse_usage.py

@@ -85,8 +85,11 @@ class ClickHouseUsageConfig(ClickHouseConfig, BaseUsageConfig, EnvConfigMixin):
 @platform_name("ClickHouse")
 @config_class(ClickHouseUsageConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class ClickHouseUsageSource(Source):
     """

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -15,7 +15,9 @@ from sqlalchemy.engine import Engine
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
@@ -58,7 +60,7 @@ AggregatedDataset = GenericAggregatedDataset[TrinoTableRef]
 
 class TrinoConnectorInfo(BaseModel):
     partitionIds: List[str]
-    truncated: Optional[bool]
+    truncated: Optional[bool] = None
 
 
 class TrinoAccessedMetadata(BaseModel):
@@ -78,7 +80,7 @@ class TrinoJoinedAccessEvent(BaseModel):
     table: Optional[str] = None
     accessed_metadata: List[TrinoAccessedMetadata]
     starttime: datetime = Field(alias="create_time")
-    endtime: Optional[datetime] = Field(alias="end_time")
+    endtime: Optional[datetime] = Field(None, alias="end_time")
 
 
 class EnvBasedSourceBaseConfig:
@@ -112,6 +114,7 @@ class TrinoUsageReport(SourceReport):
 @platform_name("Trino")
 @config_class(TrinoUsageConfig)
 @support_status(SupportStatus.CERTIFIED)
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class TrinoUsageSource(Source):
     """

datahub/ingestion/source/usage/usage_common.py

@@ -18,7 +18,7 @@ import pydantic
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.time_window_config import (
     BaseTimeWindowConfig,
     BucketDuration,
@@ -194,13 +194,13 @@ class GenericAggregatedDataset(Generic[ResourceType]):
 
 
 class BaseUsageConfig(BaseTimeWindowConfig):
-    queries_character_limit: int = Field(
+    queries_character_limit: HiddenFromDocs[int] = Field(
+        # Hidden since we don't want to encourage people to break elasticsearch.
         default=DEFAULT_QUERIES_CHARACTER_LIMIT,
         description=(
             "Total character limit for all queries in a single usage aspect."
             " Queries will be truncated to length `queries_character_limit / top_n_queries`."
         ),
-        hidden_from_docs=True,  # Don't want to encourage people to break elasticsearch
     )
 
     top_n_queries: pydantic.PositiveInt = Field(
@@ -268,6 +268,7 @@ class UsageAggregator(Generic[ResourceType]):
             user,
             query,
             fields,
+            user_email_pattern=self.config.user_email_pattern,
             count=count,
         )
 

datahub/ingestion/source/vertexai/vertexai.py

@@ -145,7 +145,7 @@ class PipelineMetadata:
 
 @platform_name("Vertex AI", id="vertexai")
 @config_class(VertexAIConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for Vertex AI Registered Models and Model Versions",

datahub/ingestion/source_config/pulsar.py

@@ -2,6 +2,7 @@ import re
 from typing import Dict, List, Optional, Union
 from urllib.parse import urlparse
 
+import pydantic
 from pydantic import Field, validator
 
 from datahub.configuration.common import AllowDenyPattern
@@ -121,7 +122,8 @@ class PulsarSourceConfig(
             )
         return client_secret
 
-    @validator("web_service_url")
+    @pydantic.field_validator("web_service_url", mode="after")
+    @classmethod
     def web_service_url_scheme_host_port(cls, val: str) -> str:
         # Tokenize the web url
         url = urlparse(val)

datahub/ingestion/source_report/ingestion_stage.py

@@ -1,7 +1,9 @@
 import logging
+from collections import defaultdict
 from contextlib import AbstractContextManager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from enum import Enum
 
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.stats_collections import TopKDict
@@ -20,31 +22,68 @@ QUERIES_EXTRACTION = "Queries Extraction"
 PROFILING = "Profiling"
 
 
+class IngestionHighStage(Enum):
+    """
+    The high-level stages at the framework level
+    Team to add more stages as needed
+    """
+
+    PROFILING = "Profiling"
+    _UNDEFINED = "Ingestion"
+
+
 @dataclass
 class IngestionStageReport:
+    ingestion_high_stage_seconds: dict[IngestionHighStage, float] = field(
+        default_factory=lambda: defaultdict(float)
+    )
     ingestion_stage_durations: TopKDict[str, float] = field(default_factory=TopKDict)
 
-    def new_stage(self, stage: str) -> "IngestionStageContext":
-        return IngestionStageContext(stage, self)
+    def new_stage(
+        self, stage: str, high_stage: IngestionHighStage = IngestionHighStage._UNDEFINED
+    ) -> "IngestionStageContext":
+        return IngestionStageContext(stage, self, high_stage)
+
+    def new_high_stage(self, stage: IngestionHighStage) -> "IngestionStageContext":
+        return IngestionStageContext("", self, stage)
 
 
 @dataclass
 class IngestionStageContext(AbstractContextManager):
-    def __init__(self, stage: str, report: IngestionStageReport):
-        self._ingestion_stage = f"{stage} at {datetime.now(timezone.utc)}"
+    def __init__(
+        self,
+        stage: str,
+        report: IngestionStageReport,
+        high_stage: IngestionHighStage = IngestionHighStage._UNDEFINED,
+    ):
+        self._high_stage = high_stage
+        self._ingestion_stage = (
+            f"{stage} at {datetime.now(timezone.utc)}" if stage else ""
+        )
         self._timer: PerfTimer = PerfTimer()
         self._report = report
 
     def __enter__(self) -> "IngestionStageContext":
-        logger.info(f"Stage started: {self._ingestion_stage}")
+        if self._ingestion_stage:
+            logger.info(f"Stage started: {self._ingestion_stage}")
+        else:
+            logger.info(f"High stage started: {self._high_stage.value}")
         self._timer.start()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         elapsed = self._timer.elapsed_seconds(digits=2)
-        logger.info(
-            f"Time spent in stage <{self._ingestion_stage}>: {elapsed} seconds",
-            stacklevel=2,
-        )
-        self._report.ingestion_stage_durations[self._ingestion_stage] = elapsed
-        return None
+        if self._ingestion_stage:
+            logger.info(
+                f"Time spent in stage <{self._ingestion_stage}>: {elapsed} seconds",
+                stacklevel=2,
+            )
+            # Store tuple as string to avoid serialization errors
+            key = f"({self._high_stage.value}, {self._ingestion_stage})"
+            self._report.ingestion_stage_durations[key] = elapsed
+        else:
+            logger.info(
+                f"Time spent in stage <{self._high_stage.value}>: {elapsed} seconds",
+                stacklevel=2,
+            )
+        self._report.ingestion_high_stage_seconds[self._high_stage] += elapsed

datahub/ingestion/transformer/add_dataset_ownership.py

@@ -71,8 +71,24 @@ class AddDatasetOwnership(OwnershipTransformer):
 
         server_ownership = graph.get_ownership(entity_urn=urn)
         if server_ownership:
-            owners = {owner.owner: owner for owner in server_ownership.owners}
-            owners.update({owner.owner: owner for owner in mce_ownership.owners})
+            owners = {
+                (
+                    owner.owner,
+                    owner.type,
+                    owner.typeUrn,
+                ): owner
+                for owner in server_ownership.owners
+            }
+            owners.update(
+                {
+                    (
+                        owner.owner,
+                        owner.type,
+                        owner.typeUrn,
+                    ): owner
+                    for owner in mce_ownership.owners
+                }
+            )
             mce_ownership.owners = list(owners.values())
 
         return mce_ownership

datahub/ingestion/transformer/base_transformer.py

@@ -281,11 +281,14 @@ class BaseTransformer(Transformer, metaclass=ABCMeta):
                                 )
                             )
 
-                            record_metadata = _update_work_unit_id(
-                                envelope=envelope,
-                                aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
-                                urn=mcp.entityUrn,
-                            )
+                            if mcp.entityUrn:
+                                record_metadata = _update_work_unit_id(
+                                    envelope=envelope,
+                                    aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
+                                    urn=mcp.entityUrn,
+                                )
+                            else:
+                                record_metadata = envelope.metadata.copy()
 
                             yield RecordEnvelope(
                                 record=mcp,

datahub/ingestion/transformer/set_browse_path.py

@@ -0,0 +1,112 @@
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, cast
+
+from datahub.configuration.common import (
+    TransformerSemanticsConfigModel,
+)
+from datahub.emitter.mce_builder import Aspect
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.transformer.base_transformer import (
+    BaseTransformer,
+    SingleAspectTransformer,
+)
+from datahub.metadata.schema_classes import (
+    BrowsePathEntryClass,
+    BrowsePathsV2Class,
+)
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+class SetBrowsePathTransformerConfig(TransformerSemanticsConfigModel):
+    path: List[str]
+
+
+class SetBrowsePathTransformer(BaseTransformer, SingleAspectTransformer):
+    ctx: PipelineContext
+    config: SetBrowsePathTransformerConfig
+
+    def __init__(self, config: SetBrowsePathTransformerConfig, ctx: PipelineContext):
+        super().__init__()
+        self.ctx = ctx
+        self.config = config
+
+    def aspect_name(self) -> str:
+        return "browsePathsV2"
+
+    def entity_types(self) -> List[str]:
+        # This is an arbitrary list, might be adjusted if it makes sense. It might be reasonable to make it configurable
+        return ["dataset", "dataJob", "dataFlow", "chart", "dashboard", "container"]
+
+    @classmethod
+    def create(
+        cls, config_dict: dict, ctx: PipelineContext
+    ) -> "SetBrowsePathTransformer":
+        config = SetBrowsePathTransformerConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    @staticmethod
+    def _build_model(existing_browse_paths: BrowsePathsV2Class) -> Dict[str, List[str]]:
+        template_vars: Dict[str, List[str]] = {}
+        model: Dict[str, List[str]] = defaultdict(list)
+        for entry in existing_browse_paths.path or []:
+            if entry.urn:
+                entity_type = guess_entity_type(entry.urn)
+                model[entity_type].append(entry.urn)
+
+        for entity_type, urns in model.items():
+            template_vars[f"{entity_type}[*]"] = urns
+            for i, urn in enumerate(urns):
+                template_vars[f"{entity_type}[{i}]"] = [urn]
+
+        return template_vars
+
+    @classmethod
+    def _expand_nodes(
+        cls, templates: List[str], template_vars: Dict[str, List[str]]
+    ) -> BrowsePathsV2Class:
+        expanded_nodes: List[str] = []
+        for node in templates:
+            resolved_nodes = cls._resolve_template_to_nodes(node, template_vars)
+            expanded_nodes.extend(resolved_nodes)
+
+        processed_entries: List[BrowsePathEntryClass] = []
+        for node in expanded_nodes:
+            if not node or node.isspace():
+                continue
+            processed_entries.append(
+                BrowsePathEntryClass(
+                    id=node, urn=node if node.startswith("urn:") else None
+                )
+            )
+        return BrowsePathsV2Class(path=processed_entries)
+
+    def transform_aspect(
+        self, entity_urn: str, aspect_name: str, aspect: Optional[Aspect]
+    ) -> Optional[Aspect]:
+        template_vars: Dict[str, List[str]] = {}
+        if aspect is not None:
+            assert isinstance(aspect, BrowsePathsV2Class)
+            template_vars = self._build_model(aspect)
+        new_browse_paths: BrowsePathsV2Class = self._expand_nodes(
+            self.config.path, template_vars
+        )
+        if aspect is not None and not self.config.replace_existing:
+            for node in aspect.path:
+                new_browse_paths.path.append(node)
+
+        return cast(Aspect, new_browse_paths)
+
+    @staticmethod
+    def _resolve_template_to_nodes(
+        template_str: str, template_vars: Dict[str, List[str]]
+    ) -> List[str]:
+        # This mechanism can be made simpler (match against known variables only) or more complex (e.g. by using a
+        # proper templating engine, like jinja).
+        template_str = template_str.strip()
+        var_pattern = re.findall(r"^\$([a-zA-Z]+\[[0-9*]+]$)", template_str)
+
+        if not var_pattern:
+            return [template_str]
+
+        return template_vars.get(var_pattern[0], [])

datahub/integrations/assertion/snowflake/compiler.py

@@ -84,9 +84,10 @@ class SnowflakeAssertionCompiler(AssertionCompiler):
 
         dmf_definitions_path = self.output_dir / DMF_DEFINITIONS_FILE_NAME
         dmf_associations_path = self.output_dir / DMF_ASSOCIATIONS_FILE_NAME
-        with (dmf_definitions_path).open("w") as definitions, (
-            dmf_associations_path
-        ).open("w") as associations:
+        with (
+            (dmf_definitions_path).open("w") as definitions,
+            (dmf_associations_path).open("w") as associations,
+        ):
             for assertion_spec in assertion_config_spec.assertions:
                 result.report.num_processed += 1
                 try: