[DOP-32998] Extract dataset tags

dolfinus · dolfinus · commit 9091deccf208 · 2026-01-27T18:45:42.000+03:00
diff --git a/data_rentgen/consumer/extractors/batch_extraction_result.py b/data_rentgen/consumer/extractors/batch_extraction_result.py
@@ -119,6 +119,7 @@ def add_location(self, location: LocationDTO):
 
     def add_dataset(self, dataset: DatasetDTO):
         dataset.location = self.add_location(dataset.location)
+        dataset.tag_values = {self.add_tag_value(tag_value) for tag_value in dataset.tag_values}
         return self._add(self._datasets, dataset)
 
     def add_dataset_symlink(self, dataset_symlink: DatasetSymlinkDTO):
@@ -207,6 +208,7 @@ def get_tag_value(self, tag_value_key: tuple) -> TagValueDTO:
     def get_dataset(self, dataset_key: tuple) -> DatasetDTO:
         dataset = self._datasets[dataset_key]
         dataset.location = self.get_location(dataset.location.unique_key)
+        dataset.tag_values = {self.get_tag_value(tag_value.unique_key) for tag_value in dataset.tag_values}
         return dataset
 
     def get_dataset_symlink(self, dataset_symlink_key: tuple) -> DatasetSymlinkDTO:
diff --git a/data_rentgen/consumer/extractors/generic/dataset.py b/data_rentgen/consumer/extractors/generic/dataset.py
@@ -9,6 +9,8 @@
     DatasetSymlinkDTO,
     DatasetSymlinkTypeDTO,
     LocationDTO,
+    TagDTO,
+    TagValueDTO,
 )
 from data_rentgen.openlineage.dataset import (
     OpenLineageDataset,
@@ -28,7 +30,8 @@ def extract_dataset(self, dataset: OpenLineageDataset) -> DatasetDTO:
         """
         Extract DatasetDTO from input or output OpenLineageDataset
         """
-        return self._extract_dataset_ref(dataset)
+        dataset_dto = self._extract_dataset_ref(dataset)
+        return self._enrich_dataset_tags(dataset_dto, dataset)
 
     def _extract_dataset_ref(
         self,
@@ -108,3 +111,15 @@ def _connect_dataset_with_symlinks(
         )
 
         return sorted(result, key=lambda x: x.type)
+
+    def _enrich_dataset_tags(self, dataset_dto: DatasetDTO, dataset: OpenLineageDataset) -> DatasetDTO:
+        if not dataset.facets.tags:
+            return dataset_dto
+
+        for raw_tag in dataset.facets.tags.tags:
+            tag_value = TagValueDTO(
+                tag=TagDTO(name=raw_tag.key.lower().replace(" ", "_")),
+                value=raw_tag.value,
+            )
+            dataset_dto.tag_values.add(tag_value)
+        return dataset_dto
diff --git a/data_rentgen/consumer/saver.py b/data_rentgen/consumer/saver.py
@@ -62,9 +62,11 @@ async def create_datasets(self, data: BatchExtractionResult):
         self.logger.debug("Creating datasets")
         dataset_pairs = await self.unit_of_work.dataset.fetch_bulk(data.datasets())
         for dataset_dto, dataset in dataset_pairs:
-            if not dataset:
-                async with self.unit_of_work:
-                    dataset = await self.unit_of_work.dataset.create(dataset_dto)  # noqa: PLW2901
+            async with self.unit_of_work:
+                if not dataset:
+                    dataset = await self.unit_of_work.dataset.create_or_update(dataset_dto)  # noqa: PLW2901
+                else:
+                    dataset = await self.unit_of_work.dataset.update(dataset, dataset_dto)  # noqa: PLW2901
             dataset_dto.id = dataset.id
 
     async def create_dataset_symlinks(self, data: BatchExtractionResult):
diff --git a/data_rentgen/db/repositories/dataset.py b/data_rentgen/db/repositories/dataset.py
@@ -22,9 +22,11 @@
     tuple_,
     union,
 )
+from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.orm import selectinload
 
 from data_rentgen.db.models import Address, Dataset, Location, TagValue
+from data_rentgen.db.models.dataset import DatasetTagValue
 from data_rentgen.db.repositories.base import Repository
 from data_rentgen.db.utils.search import make_tsquery, ts_match, ts_rank
 from data_rentgen.dto import DatasetDTO, PaginationDTO
@@ -65,6 +67,17 @@
     .group_by(Dataset.location_id)
 )
 
+insert_tag_value_query = (
+    insert(DatasetTagValue)
+    .values(
+        {
+            "dataset_id": bindparam("dataset_id"),
+            "tag_value_id": bindparam("tag_value_id"),
+        }
+    )
+    .on_conflict_do_nothing(index_elements=["dataset_id", "tag_value_id"])
+)
+
 
 class DatasetRepository(Repository[Dataset]):
     async def fetch_bulk(self, datasets_dto: list[DatasetDTO]) -> list[tuple[DatasetDTO, Dataset | None]]:
@@ -87,10 +100,51 @@ async def fetch_bulk(self, datasets_dto: list[DatasetDTO]) -> list[tuple[Dataset
             for dto in datasets_dto
         ]
 
-    async def create(self, dataset: DatasetDTO) -> Dataset:
-        # if another worker already created the same row, just use it. if not - create with holding the lock.
-        await self._lock(dataset.location.id, dataset.name.lower())
-        return await self._get(dataset) or await self._create(dataset)
+    async def create_or_update(self, dataset: DatasetDTO) -> Dataset:
+        result = await self._get(dataset)
+        if not result:
+            # try one more time, but with lock acquired.
+            # if another worker already created the same row, just use it. if not - create with holding the lock.
+            await self._lock(dataset.location.id, dataset.name.lower())
+            result = await self._get(dataset)
+
+        if not result:
+            result = await self._create(dataset)
+        return await self.update(result, dataset)
+
+    async def _get(self, dataset: DatasetDTO) -> Dataset | None:
+        return await self._session.scalar(
+            get_one_query,
+            {
+                "location_id": dataset.location.id,
+                "name_lower": dataset.name.lower(),
+            },
+        )
+
+    async def _create(self, dataset: DatasetDTO) -> Dataset:
+        result = Dataset(location_id=dataset.location.id, name=dataset.name)
+        self._session.add(result)
+        await self._session.flush([result])
+        return result
+
+    async def update(self, existing: Dataset, new: DatasetDTO) -> Dataset:
+        if not new.tag_values:
+            # in most cases datasets have no tag values, so we can avoid INSERT statements
+            return existing
+
+        # Lock to prevent inserting the same rows from multiple workers
+        await self._lock(existing.location_id, existing.name)
+        await self._session.execute(
+            insert_tag_value_query,
+            [
+                {
+                    "dataset_id": existing.id,
+                    "tag_value_id": tag_value_dto.id,
+                }
+                for tag_value_dto in new.tag_values
+            ],
+        )
+        return existing
 
     async def paginate(
         self,
@@ -184,18 +238,3 @@ async def get_stats_by_location_ids(self, location_ids: Collection[int]) -> dict
 
         query_result = await self._session.execute(get_stats_query, {"location_ids": list(location_ids)})
         return {row.location_id: row for row in query_result.all()}
-
-    async def _get(self, dataset: DatasetDTO) -> Dataset | None:
-        return await self._session.scalar(
-            get_one_query,
-            {
-                "location_id": dataset.location.id,
-                "name_lower": dataset.name.lower(),
-            },
-        )
-
-    async def _create(self, dataset: DatasetDTO) -> Dataset:
-        result = Dataset(location_id=dataset.location.id, name=dataset.name)
-        self._session.add(result)
-        await self._session.flush([result])
-        return result
diff --git a/data_rentgen/dto/dataset.py b/data_rentgen/dto/dataset.py
@@ -6,12 +6,14 @@
 from dataclasses import dataclass, field
 
 from data_rentgen.dto.location import LocationDTO
+from data_rentgen.dto.tag import TagValueDTO
 
 
 @dataclass(slots=True)
 class DatasetDTO:
     location: LocationDTO
     name: str
+    tag_values: set[TagValueDTO] = field(default_factory=set)
     id: int | None = field(default=None, compare=False)
 
     @property
@@ -21,4 +23,5 @@ def unique_key(self) -> tuple:
     def merge(self, new: DatasetDTO) -> DatasetDTO:
         self.location.merge(new.location)
         self.id = new.id or self.id
+        self.tag_values.update(new.tag_values)
         return self
diff --git a/data_rentgen/openlineage/dataset_facets/__init__.py b/data_rentgen/openlineage/dataset_facets/__init__.py
@@ -13,6 +13,10 @@
     OpenLineageColumnLineageDatasetFacetFieldRef,
     OpenLineageColumnLineageDatasetFacetFieldTransformation,
 )
+from data_rentgen.openlineage.dataset_facets.dataset_tags import (
+    OpenLineageDatasetTagsFacet,
+    OpenLineageDatasetTagsFacetField,
+)
 from data_rentgen.openlineage.dataset_facets.documentation import (
     OpenLineageDocumentationDatasetFacet,
 )
@@ -46,6 +50,8 @@
     "OpenLineageDatasetFacets",
     "OpenLineageDatasetLifecycleStateChange",
     "OpenLineageDatasetPreviousIdentifier",
+    "OpenLineageDatasetTagsFacet",
+    "OpenLineageDatasetTagsFacetField",
     "OpenLineageDocumentationDatasetFacet",
     "OpenLineageInputDatasetFacets",
     "OpenLineageInputStatisticsInputDatasetFacet",
@@ -70,6 +76,7 @@ class OpenLineageDatasetFacets(OpenLineageBase):
     datasetSchema: OpenLineageSchemaDatasetFacet | None = Field(default=None, alias="schema")
     symlinks: OpenLineageSymlinksDatasetFacet | None = None
     columnLineage: OpenLineageColumnLineageDatasetFacet | None = None
+    tags: OpenLineageDatasetTagsFacet | None = None
 
 
 class OpenLineageInputDatasetFacets(OpenLineageBase):
diff --git a/data_rentgen/openlineage/dataset_facets/dataset_tags.py b/data_rentgen/openlineage/dataset_facets/dataset_tags.py
@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: 2024-present MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import BaseModel, Field
+
+from data_rentgen.openlineage.dataset_facets.base import OpenLineageDatasetFacet
+
+
+class OpenLineageDatasetTagsFacetField(BaseModel):
+    """Dataset tags field type.
+    See [TagsDatasetFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/TagsDatasetFacet.json).
+    """
+
+    key: str = Field(description="Key that identifies the tag")
+    value: str = Field(description="The value of the field")
+    source: str | None = Field(default=None, description="The source of the tag. INTEGRATION|USER|DBT CORE|SPARK|etc.")
+    field: str | None = Field(default=None, description="Identifies the field in a dataset if a tag applies to one")
+
+
+class OpenLineageDatasetTagsFacet(OpenLineageDatasetFacet):
+    """Dataset facet describing dataset tags.
+    See [TagsDatasetFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/DatasetTypeDatasetFacet.json).
+    """
+
+    tags: list[OpenLineageDatasetTagsFacetField] = Field(
+        default_factory=list,
+        description="The tags applied to the dataset facet",
+    )
diff --git a/docs/changelog/next_release/372.feature.rst b/docs/changelog/next_release/372.feature.rst
@@ -0,0 +1 @@
+Extract dataset tags provided by OpenLineage integrations.
diff --git a/tests/test_consumer/test_extractors/test_extractors_dataset.py b/tests/test_consumer/test_extractors/test_extractors_dataset.py
@@ -5,12 +5,16 @@
     DatasetSymlinkDTO,
     DatasetSymlinkTypeDTO,
     LocationDTO,
+    TagDTO,
+    TagValueDTO,
 )
 from data_rentgen.openlineage.dataset import (
     OpenLineageDataset,
 )
 from data_rentgen.openlineage.dataset_facets import (
     OpenLineageDatasetFacets,
+    OpenLineageDatasetTagsFacet,
+    OpenLineageDatasetTagsFacetField,
     OpenLineageSymlinkIdentifier,
     OpenLineageSymlinksDatasetFacet,
     OpenLineageSymlinkType,
@@ -293,3 +297,35 @@ def test_extractors_extract_dataset_unknown():
         name="some.name",
     )
     assert symlinks_dto == []
+
+
+def test_extractors_extract_dataset_with_tags():
+    dataset = OpenLineageDataset(
+        namespace="postgres://192.168.1.1:5432",
+        name="mydb.myschema.mytable",
+        facets=OpenLineageDatasetFacets(
+            tags=OpenLineageDatasetTagsFacet(
+                tags=[
+                    OpenLineageDatasetTagsFacetField(key="somekey", value="somevalue"),
+                    OpenLineageDatasetTagsFacetField(key="somekey", value="othervalue", source="OTHER"),
+                    OpenLineageDatasetTagsFacetField(key="anotherkey", value="anothervalue", source="ABC", field="abc"),
+                ],
+            ),
+        ),
+    )
+
+    dataset_dto, symlinks_dto = GenericExtractor().extract_dataset_and_symlinks(dataset)
+    assert dataset_dto == DatasetDTO(
+        location=LocationDTO(
+            type="postgres",
+            name="192.168.1.1:5432",
+            addresses={"postgres://192.168.1.1:5432"},
+        ),
+        name="mydb.myschema.mytable",
+        tag_values={
+            TagValueDTO(tag=TagDTO(name="somekey"), value="somevalue"),
+            TagValueDTO(tag=TagDTO(name="somekey"), value="othervalue"),
+            TagValueDTO(tag=TagDTO(name="anotherkey"), value="anothervalue"),
+        },
+    )
+    assert symlinks_dto == []

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Extract dataset tags provided by OpenLineage integrations.`