openmetadata-ingestion 1.4.7.2__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/ingestion/source/storage/s3/metadata.py

@@ -14,7 +14,7 @@ import secrets
 import traceback
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from pydantic import ValidationError
 
@@ -40,13 +40,18 @@ from metadata.generated.schema.metadataIngestion.storage.containerMetadataConfig
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.storage.s3.models import (
     S3BucketResponse,
     S3ContainerDetails,
+    S3Tag,
+    S3TagResponse,
 )
 from metadata.ingestion.source.storage.storage_service import (
     KEY_SEPARATOR,
@@ -58,11 +63,14 @@ from metadata.readers.file.config_source_factory import get_reader
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_container
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_label
 
 logger = ingestion_logger()
 
 S3_CLIENT_ROOT_RESPONSE = "Contents"
 
+WILD_CARD = "*"
+
 
 class S3Metric(Enum):
     NUMBER_OF_OBJECTS = "NumberOfObjects"
@@ -80,14 +88,15 @@ class S3Source(StorageServiceSource):
         self.cloudwatch_client = self.connection.cloudwatch_client
 
         self._bucket_cache: Dict[str, Container] = {}
+        self._unstructured_container_cache: Dict[str, Tuple[str, str]] = {}
         self.s3_reader = get_reader(config_source=S3Config(), client=self.s3_client)
 
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: S3Connection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: S3Connection = config.serviceConnection.root.config
         if not isinstance(connection, S3Connection):
             raise InvalidSourceException(f"Expected S3Connection, but got {connection}")
         return cls(config, metadata)
@@ -112,8 +121,12 @@ class S3Source(StorageServiceSource):
                     entity=Container, fqn=container_fqn
                 )
                 self._bucket_cache[bucket_name] = container_entity
+                self._unstructured_container_cache[container_fqn] = (
+                    container_entity.id.root,
+                    KEY_SEPARATOR,
+                )
                 parent_entity: EntityReference = EntityReference(
-                    id=self._bucket_cache[bucket_name].id.__root__, type="container"
+                    id=self._bucket_cache[bucket_name].id.root, type="container"
                 )
                 if self.global_manifest:
                     manifest_entries_for_current_bucket = (
@@ -129,6 +142,11 @@ class S3Source(StorageServiceSource):
                             entries=manifest_entries_for_current_bucket,
                             parent=parent_entity,
                         )
+                        yield from self._generate_unstructured_containers(
+                            bucket_response=bucket_response,
+                            entries=manifest_entries_for_current_bucket,
+                            parent=parent_entity,
+                        )
                         # nothing else do to for the current bucket, skipping to the next
                         continue
                 # If no global file, or no valid entries in the manifest, check for bucket level metadata file
@@ -139,6 +157,14 @@ class S3Source(StorageServiceSource):
                         entries=metadata_config.entries,
                         parent=parent_entity,
                     )
+                    yield from self._generate_unstructured_containers(
+                        bucket_response=bucket_response,
+                        entries=metadata_config.entries,
+                        parent=parent_entity,
+                    )
+
+                # clean up the cache after each bucket
+                self._unstructured_container_cache.clear()
 
             except ValidationError as err:
                 self.status.failed(
@@ -157,24 +183,100 @@ class S3Source(StorageServiceSource):
                     )
                 )
 
+    def _get_bucket_name_and_key(self, full_path: str) -> Tuple[str, str]:
+        """
+        Method to get the bucket name and key from the full path
+        """
+        if full_path:
+            parts = full_path.removeprefix("s3://").split(KEY_SEPARATOR)
+            if len(parts) > 2:
+                return parts[0], KEY_SEPARATOR.join(parts[1:])
+        return None, None
+
+    def get_tag_by_fqn(self, entity_fqn: str) -> Optional[List[TagLabel]]:
+        """
+        Pick up the tags registered in the context
+        searching by entity FQN
+        """
+        try:
+            tag_labels = []
+            for tag_and_category in self.context.get().tags or []:
+                if tag_and_category.fqn and tag_and_category.fqn.root == entity_fqn:
+                    tag_label = get_tag_label(
+                        metadata=self.metadata,
+                        tag_name=tag_and_category.tag_request.name.root,
+                        classification_name=tag_and_category.classification_request.name.root,
+                    )
+                    if tag_label:
+                        tag_labels.append(tag_label)
+            return tag_labels or None
+        except Exception as exc:
+            logger.debug(f"Failed to ingest tags due to: {exc}")
+            logger.debug(traceback.format_exc())
+
+        return None
+
+    def yield_container_tags(
+        self, container_details: S3ContainerDetails
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each container
+        """
+        try:
+            bucket_name, key = self._get_bucket_name_and_key(container_details.fullPath)
+            if (
+                container_details.leaf_container
+                and container_details.container_fqn
+                and bucket_name
+                and key
+            ):
+                tags = self.s3_client.get_object_tagging(Bucket=bucket_name, Key=key)
+                tags_list: List[S3Tag] = S3TagResponse.model_validate(tags).TagSet
+                for tag in tags_list:
+                    yield from get_ometa_tag_and_classification(
+                        tag_fqn=FullyQualifiedEntityName(
+                            container_details.container_fqn
+                        ),
+                        tags=[tag.Value],
+                        classification_name=tag.Key,
+                        tag_description="S3 TAG VALUE",
+                        classification_description="S3 TAG KEY",
+                    )
+        except Exception as exc:
+            logger.debug(f"Failed to ingest tags due to: {exc}")
+            logger.debug(traceback.format_exc())
+
     def yield_create_container_requests(
         self, container_details: S3ContainerDetails
     ) -> Iterable[Either[CreateContainerRequest]]:
         container_request = CreateContainerRequest(
-            name=container_details.name,
+            name=EntityName(container_details.name),
             prefix=container_details.prefix,
             numberOfObjects=container_details.number_of_objects,
             size=container_details.size,
             dataModel=container_details.data_model,
-            service=self.context.get().objectstore_service,
+            service=FullyQualifiedEntityName(self.context.get().objectstore_service),
             parent=container_details.parent,
             sourceUrl=container_details.sourceUrl,
             fileFormats=container_details.file_formats,
             fullPath=container_details.fullPath,
+            tags=self.get_tag_by_fqn(container_details.container_fqn),
         )
         yield Either(right=container_request)
         self.register_record(container_request=container_request)
 
+    def get_size(self, bucket_name: str, file_path: str) -> Optional[float]:
+        """
+        Method to get the size of the file
+        """
+        try:
+            file_obj = self.s3_client.head_object(Bucket=bucket_name, Key=file_path)
+            return file_obj["ContentLength"]
+        except Exception as exc:
+            logger.debug(f"Failed to get size of file due to {exc}")
+            logger.debug(traceback.format_exc())
+        return None
+
     def _generate_container_details(
         self,
         bucket_response: S3BucketResponse,
@@ -182,6 +284,10 @@ class S3Source(StorageServiceSource):
         parent: Optional[EntityReference] = None,
     ) -> Optional[S3ContainerDetails]:
         bucket_name = bucket_response.name
+
+        if not metadata_entry.structureFormat:
+            return None
+
         sample_key = self._get_sample_file_path(
             bucket_name=bucket_name, metadata_entry=metadata_entry
         )
@@ -206,12 +312,6 @@ class S3Source(StorageServiceSource):
                     creation_date=bucket_response.creation_date.isoformat()
                     if bucket_response.creation_date
                     else None,
-                    number_of_objects=self._fetch_metric(
-                        bucket_name=bucket_name, metric=S3Metric.NUMBER_OF_OBJECTS
-                    ),
-                    size=self._fetch_metric(
-                        bucket_name=bucket_name, metric=S3Metric.BUCKET_SIZE_BYTES
-                    ),
                     file_formats=[container.FileFormat(metadata_entry.structureFormat)],
                     data_model=ContainerDataModel(
                         isPartitioned=metadata_entry.isPartitioned, columns=columns
@@ -230,8 +330,7 @@ class S3Source(StorageServiceSource):
         bucket_response: S3BucketResponse,
         entries: List[MetadataEntry],
         parent: Optional[EntityReference] = None,
-    ) -> List[S3ContainerDetails]:
-        result: List[S3ContainerDetails] = []
+    ) -> Iterable[S3ContainerDetails]:
         for metadata_entry in entries:
             logger.info(
                 f"Extracting metadata from path {metadata_entry.dataPath.strip(KEY_SEPARATOR)} "
@@ -245,9 +344,160 @@ class S3Source(StorageServiceSource):
                 parent=parent,
             )
             if structured_container:
-                result.append(structured_container)
+                yield structured_container
+
+    def is_valid_unstructured_file(self, accepted_extensions: List, key: str) -> bool:
+        # Split the string into a list of values
+        if WILD_CARD in accepted_extensions:
+            return True
+
+        for ext in accepted_extensions:
+            if key.endswith(ext):
+                return True
+
+        return False
 
-        return result
+    def _yield_parents_of_unstructured_container(
+        self,
+        bucket_name: str,
+        list_of_parent: List[str],
+        parent: Optional[EntityReference] = None,
+    ):
+        full_path = self._get_full_path(bucket_name)
+        sub_parent = parent
+        for i in range(len(list_of_parent) - 1):
+            container_fqn = fqn._build(  # pylint: disable=protected-access
+                *(
+                    self.context.get().objectstore_service,
+                    bucket_name,
+                    *list_of_parent[: i + 1],
+                )
+            )
+            if container_fqn in self._unstructured_container_cache:
+                parent_id, full_path = self._unstructured_container_cache[container_fqn]
+                sub_parent = EntityReference(id=parent_id, type="container")
+                continue
+            yield S3ContainerDetails(
+                name=list_of_parent[i],
+                prefix=full_path,
+                file_formats=[],
+                parent=sub_parent,
+                fullPath=full_path + KEY_SEPARATOR + list_of_parent[i],
+                sourceUrl=self._get_object_source_url(
+                    bucket_name=bucket_name,
+                    prefix=self._clean_path(
+                        full_path + KEY_SEPARATOR + list_of_parent[i]
+                    ),
+                ),
+            )
+            container_entity = self.metadata.get_by_name(
+                entity=Container, fqn=container_fqn
+            )
+            full_path += KEY_SEPARATOR + list_of_parent[i]
+            self._unstructured_container_cache[container_fqn] = (
+                container_entity.id.root,
+                full_path,
+            )
+            sub_parent = EntityReference(id=container_entity.id.root, type="container")
+
+    def _yield_nested_unstructured_containers(
+        self,
+        bucket_response: S3BucketResponse,
+        metadata_entry: MetadataEntry,
+        parent: Optional[EntityReference] = None,
+    ):
+        bucket_name = bucket_response.name
+        response = self.s3_client.list_objects_v2(
+            Bucket=bucket_name, Prefix=metadata_entry.dataPath
+        )
+        candidate_keys = [
+            entry["Key"]
+            for entry in response[S3_CLIENT_ROOT_RESPONSE]
+            if entry and entry.get("Key")
+        ]
+        for key in candidate_keys:
+            if self.is_valid_unstructured_file(metadata_entry.unstructuredFormats, key):
+                logger.info(
+                    f"Extracting metadata from path {key.strip(KEY_SEPARATOR)} "
+                    f"and generating unstructured container"
+                )
+                list_of_parent = key.strip(KEY_SEPARATOR).split(KEY_SEPARATOR)
+                yield from self._yield_parents_of_unstructured_container(
+                    bucket_name, list_of_parent, parent
+                )
+                parent_fqn = fqn._build(  # pylint: disable=protected-access
+                    *(
+                        self.context.get().objectstore_service,
+                        bucket_name,
+                        *list_of_parent[:-1],
+                    )
+                )
+                parent_id, parent_path = self._unstructured_container_cache[parent_fqn]
+                container_fqn = fqn._build(  # pylint: disable=protected-access
+                    *(
+                        self.context.get().objectstore_service,
+                        bucket_name,
+                        *list_of_parent,
+                    )
+                )
+                size = self.get_size(bucket_name, key)
+                yield S3ContainerDetails(
+                    name=list_of_parent[-1],
+                    prefix=self._clean_path(parent_path + KEY_SEPARATOR),
+                    file_formats=[],
+                    size=size,
+                    container_fqn=container_fqn,
+                    leaf_container=True,
+                    parent=EntityReference(id=parent_id, type="container"),
+                    fullPath=self._get_full_path(bucket_name, key),
+                    sourceUrl=self._get_object_source_url(
+                        bucket_name=bucket_name,
+                        prefix=self._clean_path(
+                            parent_path + KEY_SEPARATOR + list_of_parent[-1]
+                        ),
+                    ),
+                )
+
+    def _generate_unstructured_containers(
+        self,
+        bucket_response: S3BucketResponse,
+        entries: List[MetadataEntry],
+        parent: Optional[EntityReference] = None,
+    ) -> Iterable[S3ContainerDetails]:
+        bucket_name = bucket_response.name
+        for metadata_entry in entries:
+            if metadata_entry.structureFormat:
+                continue
+            if metadata_entry.unstructuredFormats:
+                yield from self._yield_nested_unstructured_containers(
+                    bucket_response=bucket_response,
+                    metadata_entry=metadata_entry,
+                    parent=parent,
+                )
+            else:
+                logger.info(
+                    f"Extracting metadata from path {metadata_entry.dataPath.strip(KEY_SEPARATOR)} "
+                    f"and generating unstructured container"
+                )
+                prefix = (
+                    f"{KEY_SEPARATOR}{metadata_entry.dataPath.strip(KEY_SEPARATOR)}"
+                )
+                yield S3ContainerDetails(
+                    name=metadata_entry.dataPath.strip(KEY_SEPARATOR),
+                    prefix=prefix,
+                    file_formats=[],
+                    data_model=None,
+                    parent=parent,
+                    size=self.get_size(
+                        bucket_name=bucket_name,
+                        file_path=metadata_entry.dataPath.strip(KEY_SEPARATOR),
+                    ),
+                    fullPath=self._get_full_path(bucket_name, prefix),
+                    sourceUrl=self._get_object_source_url(
+                        bucket_name=bucket_name,
+                        prefix=metadata_entry.dataPath.strip(KEY_SEPARATOR),
+                    ),
+                )
 
     def fetch_buckets(self) -> List[S3BucketResponse]:
         results: List[S3BucketResponse] = []
@@ -265,7 +515,7 @@ class S3Source(StorageServiceSource):
                 ):
                     self.status.filter(bucket["Name"], "Bucket Filtered Out")
                 else:
-                    results.append(S3BucketResponse.parse_obj(bucket))
+                    results.append(S3BucketResponse.model_validate(bucket))
         except Exception as err:
             logger.debug(traceback.format_exc())
             logger.error(f"Failed to fetch buckets list - {err}")
@@ -450,7 +700,7 @@ class S3Source(StorageServiceSource):
                 verbose=False,
             )
             content = json.loads(response_object)
-            metadata_config = StorageContainerConfig.parse_obj(content)
+            metadata_config = StorageContainerConfig.model_validate(content)
             return metadata_config
         except ReadException:
             logger.warning(

metadata/ingestion/source/storage/s3/models.py

@@ -14,7 +14,7 @@ S3 custom pydantic models
 from datetime import datetime
 from typing import List, Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from metadata.generated.schema.entity.data.container import (
     ContainerDataModel,
@@ -29,8 +29,9 @@ class S3BucketResponse(BaseModel):
     Class modelling a response received from s3_client.list_buckets operation
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     name: str = Field(..., description="Bucket name", alias="Name")
     creation_date: Optional[datetime] = Field(
@@ -40,31 +41,49 @@ class S3BucketResponse(BaseModel):
     )
 
 
+class S3Tag(BaseModel):
+    Key: str
+    Value: str
+
+
+class S3TagResponse(BaseModel):
+    """
+    Class modelling a response received from s3_client.get_bucket_tagging operation
+    """
+
+    TagSet: List[S3Tag] = Field([], description="List of tags")
+
+
 class S3ContainerDetails(BaseModel):
     """
     Class mapping container details used to create the container requests
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
+    leaf_container: bool = Field(False, description="Leaf container")
+    container_fqn: Optional[str] = Field(
+        None, description="Fully qualified name of the container"
+    )
     name: str = Field(..., description="Bucket name")
     prefix: str = Field(..., description="Prefix for the container")
-    number_of_objects: float = Field(
-        ...,
+    number_of_objects: Optional[float] = Field(
+        None,
         description="Total nr. of objects",
     )
-    size: float = Field(
-        ...,
+    size: Optional[float] = Field(
+        None,
         description="Total size in bytes of all objects",
         title="Total size(bytes) of objects",
     )
     file_formats: Optional[List[FileFormat]] = Field(
-        ...,
+        None,
         description="File formats",
     )
     data_model: Optional[ContainerDataModel] = Field(
-        ...,
+        None,
         description="Data Model of the container",
     )
     creation_date: Optional[str] = Field(

metadata/ingestion/source/storage/storage_service.py

@@ -14,6 +14,9 @@ Base class for ingesting Object Storage services
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Set
 
+from pydantic import Field
+from typing_extensions import Annotated
+
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.services.storageService import (
@@ -38,6 +41,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
@@ -68,7 +72,14 @@ OPENMETADATA_TEMPLATE_FILE_NAME = "openmetadata.json"
 
 
 class StorageServiceTopology(ServiceTopology):
-    root = TopologyNode(
+    """
+    Defines the hierarchy in Messaging Services.
+    service -> container -> container -> container...
+    """
+
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
         producer="get_services",
         stages=[
             NodeStage(
@@ -84,9 +95,18 @@ class StorageServiceTopology(ServiceTopology):
         post_process=["mark_containers_as_deleted"],
     )
 
-    container = TopologyNode(
+    container: Annotated[
+        TopologyNode, Field(description="Container Processing Node")
+    ] = TopologyNode(
         producer="get_containers",
         stages=[
+            NodeStage(
+                type_=OMetaTagAndClassification,
+                context="tags",
+                processor="yield_tag_details",
+                nullable=True,
+                store_all_in_context=True,
+            ),
             NodeStage(
                 type_=Container,
                 context="container",
@@ -94,7 +114,7 @@ class StorageServiceTopology(ServiceTopology):
                 consumer=["objectstore_service"],
                 nullable=True,
                 use_cache=True,
-            )
+            ),
         ],
     )
 
@@ -109,7 +129,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     config: WorkflowSource
     metadata: OpenMetadata
     # Big union of types we want to fetch dynamically
-    service_connection: StorageConnection.__fields__["config"].type_
+    service_connection: StorageConnection.model_fields["config"].annotation
 
     topology = StorageServiceTopology()
     context = TopologyContextManager(topology)
@@ -125,7 +145,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.source_config: StorageServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
@@ -176,6 +196,22 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     def prepare(self):
         """By default, nothing needs to be taken care of when loading the source"""
 
+    def yield_container_tags(
+        self, container_details: Any
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each container
+        """
+
+    def yield_tag_details(
+        self, container_details: Any
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each container
+        """
+        if self.source_config.includeTags:
+            yield from self.yield_container_tags(container_details) or []
+
     def register_record(self, container_request: CreateContainerRequest) -> None:
         """
         Mark the container record as scanned and update
@@ -184,7 +220,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         parent_container = (
             self.metadata.get_by_id(
                 entity=Container, entity_id=container_request.parent.id
-            ).fullyQualifiedName.__root__
+            ).fullyQualifiedName.root
             if container_request.parent
             else None
         )
@@ -193,7 +229,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             entity_type=Container,
             service_name=self.context.get().objectstore_service,
             parent_container=parent_container,
-            container_name=container_request.name.__root__,
+            container_name=fqn.quote_name(container_request.name.root),
         )
 
         self.container_source_state.add(container_fqn)

metadata/ingestion/stage/table_usage.py

@@ -76,7 +76,7 @@ class TableUsageStage(Stage):
         metadata: OpenMetadata,
         pipeline_name: Optional[str] = None,
     ):
-        config = TableStageConfig.parse_obj(config_dict)
+        config = TableStageConfig.model_validate(config_dict)
         return cls(config, metadata)
 
     def init_location(self) -> None:
@@ -100,7 +100,7 @@ class TableUsageStage(Stage):
         if username:
             user = self.metadata.get_by_name(entity=User, fqn=username)
             if user:
-                return [user.fullyQualifiedName.__root__], [username]
+                return [user.fullyQualifiedName.root], [username]
             return None, [username]
         return None, None
 
@@ -191,7 +191,7 @@ class TableUsageStage(Stage):
         for key, value in self.table_usage.items():
             if value:
                 value.sqlQueries = self.table_queries.get(key, [])
-                data = value.json()
+                data = value.model_dump_json()
                 with open(
                     os.path.join(self.config.filename, f"{value.serviceName}_{key[1]}"),
                     "a+",

metadata/mixins/pandas/pandas_mixin.py

@@ -91,7 +91,7 @@ class PandasInterfaceMixin:
             config_source=service_connection_config.configSource,
             client=client,
             file_fqn=DatalakeTableSchemaWrapper(
-                key=table.name.__root__,
+                key=table.name.root,
                 bucket_name=table.databaseSchema.name,
                 file_extension=table.fileFormat,
             ),
@@ -125,4 +125,4 @@ class PandasInterfaceMixin:
                         for df in data
                     ]
             return data
-        raise TypeError(f"Couldn't fetch {table.name.__root__}")
+        raise TypeError(f"Couldn't fetch {table.name.root}")

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -112,5 +112,5 @@ class SQAInterfaceMixin:
         return [
             column.name
             for column in self.table.__table__.columns
-            if column.name in {col.name.__root__ for col in self.table_entity.columns}
+            if column.name in {col.name.root for col in self.table_entity.columns}
         ]