openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

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

@@ -0,0 +1,86 @@
+#  Copyright 2024 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+GCS custom pydantic models
+"""
+from datetime import datetime
+from typing import List, Optional
+
+from pydantic import BaseModel, Extra, Field
+
+from metadata.generated.schema.entity.data.container import (
+    ContainerDataModel,
+    FileFormat,
+)
+from metadata.generated.schema.type import basic
+from metadata.generated.schema.type.entityReference import EntityReference
+
+
+class GCSBucketResponse(BaseModel):
+    """
+    Class modelling a response received from gcs_client.list_buckets operation
+    """
+
+    class Config:
+        extra = Extra.forbid
+
+    name: str = Field(..., description="Bucket name")
+    project_id: str = Field(..., description="Project ID")
+    creation_date: Optional[datetime] = Field(
+        None,
+        description="Timestamp of Bucket creation in ISO format",
+    )
+
+
+class GCSContainerDetails(BaseModel):
+    """
+    Class mapping container details used to create the container requests
+    """
+
+    class Config:
+        extra = Extra.forbid
+
+    name: str = Field(..., description="Bucket name")
+    prefix: str = Field(..., description="Prefix for the container")
+    description: Optional[basic.Markdown] = Field(
+        None, description="Description of the container instance."
+    )
+    number_of_objects: float = Field(
+        ...,
+        description="Total nr. of objects",
+    )
+    size: float = Field(
+        ...,
+        description="Total size in bytes of all objects",
+        title="Total size(bytes) of objects",
+    )
+    file_formats: Optional[List[FileFormat]] = Field(
+        ...,
+        description="File formats",
+    )
+    data_model: Optional[ContainerDataModel] = Field(
+        ...,
+        description="Data Model of the container",
+    )
+    creation_date: Optional[str] = Field(
+        None,
+        description="Timestamp of Bucket creation in ISO format",
+    )
+    parent: Optional[EntityReference] = Field(
+        None,
+        description="Reference to the parent container",
+    )
+    sourceUrl: Optional[basic.SourceUrl] = Field(
+        None, description="Source URL of the container."
+    )
+    fullPath: Optional[str] = Field(
+        None, description="Full path of the container/file."
+    )

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(