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

metadata/profiler/source/metadata.py

@@ -14,7 +14,7 @@ OpenMetadata source for the profiler
 import traceback
 from typing import Iterable, Optional, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
@@ -50,16 +50,14 @@ TAGS_FIELD = ["tags"]
 class ProfilerSourceAndEntity(BaseModel):
     """Return class for the OpenMetadata Profiler Source"""
 
-    class Config:
-        arbitrary_types_allowed = True
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
     profiler_source: ProfilerSource
     entity: Table
 
     def __str__(self):
         """Return the information of the table being profiler"""
-        return f"Table [{self.entity.name.__root__}]"
+        return f"Table [{self.entity.name.root}]"
 
 
 class OpenMetadataSource(Source):
@@ -146,7 +144,7 @@ class OpenMetadataSource(Source):
             except Exception as exc:
                 yield Either(
                     left=StackTraceError(
-                        name=database.fullyQualifiedName.__root__,
+                        name=database.fullyQualifiedName.root,
                         error=f"Error listing source and entities for database due to [{exc}]",
                         stackTrace=traceback.format_exc(),
                     )
@@ -168,15 +166,15 @@ class OpenMetadataSource(Source):
             self.metadata,
             entity_type=Database,
             service_name=self.config.source.serviceName,
-            database_name=database.name.__root__,
+            database_name=database.name.root,
         )
         if filter_by_database(
             self.source_config.databaseFilterPattern,
             database_fqn
             if self.source_config.useFqnForFiltering
-            else database.name.__root__,
+            else database.name.root,
         ):
-            self.status.filter(database.name.__root__, "Database pattern not allowed")
+            self.status.filter(database.name.root, "Database pattern not allowed")
             return None
         return database
 
@@ -203,7 +201,7 @@ class OpenMetadataSource(Source):
                     else table.databaseSchema.name,  # type: ignore
                 ):
                     self.status.filter(
-                        f"Schema pattern not allowed: {table.fullyQualifiedName.__root__}",
+                        f"Schema pattern not allowed: {table.fullyQualifiedName.root}",
                         "Schema pattern not allowed",
                     )
                     continue
@@ -213,17 +211,17 @@ class OpenMetadataSource(Source):
                     service_name=self.config.source.serviceName,
                     database_name=table.database.name,
                     schema_name=table.databaseSchema.name,
-                    table_name=table.name.__root__,
+                    table_name=table.name.root,
                 )
 
                 if filter_by_table(
                     self.source_config.tableFilterPattern,
                     table_fqn
                     if self.source_config.useFqnForFiltering
-                    else table.name.__root__,
+                    else table.name.root,
                 ):
                     self.status.filter(
-                        f"Table pattern not allowed: {table.fullyQualifiedName.__root__}",
+                        f"Table pattern not allowed: {table.fullyQualifiedName.root}",
                         "Table pattern not allowed",
                     )
                     continue
@@ -232,7 +230,7 @@ class OpenMetadataSource(Source):
                     and not self.source_config.includeViews
                 ):
                     self.status.filter(
-                        table.fullyQualifiedName.__root__,
+                        table.fullyQualifiedName.root,
                         "View filtered out",
                     )
                     continue
@@ -240,7 +238,7 @@ class OpenMetadataSource(Source):
             except Exception as exc:
                 self.status.failed(
                     StackTraceError(
-                        name=table.fullyQualifiedName.__root__,
+                        name=table.fullyQualifiedName.root,
                         error=f"Unexpected error filtering entities for table [{table}]: {exc}",
                         stackTrace=traceback.format_exc(),
                     )
@@ -293,7 +291,7 @@ class OpenMetadataSource(Source):
                     self.metadata,
                     entity_type=Database,
                     service_name=self.config.source.serviceName,
-                    database_name=database.name.__root__,
+                    database_name=database.name.root,
                 ),
             },  # type: ignore
         )

metadata/profiler/source/metadata_ext.py

@@ -79,7 +79,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
         self.test_connection()
 
         # Init and type the source config
-        self.service_connection = self.config.source.serviceConnection.__root__.config
+        self.service_connection = self.config.source.serviceConnection.root.config
         self.source_config: DatabaseServiceProfilerPipeline = cast(
             DatabaseServiceProfilerPipeline, self.config.source.sourceConfig.config
         )  # Used to satisfy type checked
@@ -91,7 +91,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
         database_source_config = DatabaseServiceMetadataPipeline()
         new_config = deepcopy(self.config.source)
         new_config.sourceConfig.config = database_source_config
-        self.source = source_class.create(new_config.dict(), self.metadata)
+        self.source = source_class.create(new_config.model_dump(), self.metadata)
         self.engine = None
         self.inspector = None
         self._connection = None
@@ -253,7 +253,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                     self.metadata,
                     entity_type=Database,
                     service_name=self.config.source.serviceName,
-                    database_name=database.name.__root__,
+                    database_name=database.name.root,
                 ),
             },  # type: ignore
         )

metadata/readers/dataframe/avro.py

@@ -68,9 +68,7 @@ class AvroDataFrameReader(DataFrameReader):
         except (AssertionError, InvalidAvroBinaryEncoding):
             columns = parse_avro_schema(schema=avro_text, cls=Column)
             field_map = {
-                col.name.__root__: Series(
-                    PD_AVRO_FIELD_MAP.get(col.dataType.value, "str")
-                )
+                col.name.root: Series(PD_AVRO_FIELD_MAP.get(col.dataType.value, "str"))
                 for col in columns
             }
             return DatalakeColumnWrapper(

metadata/readers/dataframe/models.py

@@ -15,6 +15,7 @@ Module to define pydentic models related to datalake
 from typing import Any, List, Optional
 
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.data.table import Column
 
@@ -27,9 +28,20 @@ class DatalakeColumnWrapper(BaseModel):
     which can be used by both profiler and metadata ingestion
     """
 
-    columns: Optional[List[Column]]
-    dataframes: Optional[List[Any]]  # pandas.Dataframe does not have any validators
-    raw_data: Any  # in special cases like json schema, we need to store the raw data
+    columns: Annotated[
+        Optional[List[Column]], Field(None, description="List of columns")
+    ]
+    # pandas.Dataframe does not have any validators
+    dataframes: Annotated[
+        Optional[List[Any]], Field(None, description="List of dataframes")
+    ]
+    raw_data: Annotated[
+        Any,
+        Field(
+            None,
+            description="In special cases like json schema, we need to store the raw data",
+        ),
+    ]
 
 
 class DatalakeTableSchemaWrapper(BaseModel):
@@ -37,12 +49,15 @@ class DatalakeTableSchemaWrapper(BaseModel):
     Instead of sending the whole Table model from profiler, we send only key and bucket name using this model
     """
 
-    key: str
-    bucket_name: str
-    file_extension: Optional[Any]
-    separator: Optional[str] = Field(
-        None, description="Used for DSV readers to identify the separator"
-    )
+    key: Annotated[str, Field(..., description="Key of the file in the bucket")]
+    bucket_name: Annotated[str, Field(..., description="Name of the bucket")]
+    file_extension: Annotated[
+        Optional[Any], Field(None, description="File extension of the file")
+    ]
+    separator: Annotated[
+        Optional[str],
+        Field(None, description="Used for DSV readers to identify the separator"),
+    ]
 
 
 class DatalakeTableMetadata(BaseModel):
@@ -50,6 +65,8 @@ class DatalakeTableMetadata(BaseModel):
     Used to yield metadata from datalake buckets
     """
 
-    table: str
-    table_type: str
-    file_extension: Optional[Any]
+    table: Annotated[str, Field(..., description="Name of the table")]
+    table_type: Annotated[str, Field(..., description="Type of the table")]
+    file_extension: Annotated[
+        Optional[Any], Field(None, description="File extension of the file")
+    ]

metadata/readers/dataframe/parquet.py

@@ -69,9 +69,9 @@ class ParquetDataFrameReader(DataFrameReader):
 
         client_kwargs = {}
         if self.config_source.securityConfig.endPointURL:
-            client_kwargs[
-                "endpoint_url"
-            ] = self.config_source.securityConfig.endPointURL
+            client_kwargs["endpoint_url"] = str(
+                self.config_source.securityConfig.endPointURL
+            )
 
         if self.config_source.securityConfig.awsRegion:
             client_kwargs["region_name"] = self.config_source.securityConfig.awsRegion

metadata/readers/dataframe/reader_factory.py

@@ -42,6 +42,9 @@ class SupportedTypes(Enum):
     JSON = "json"
     JSONGZ = "json.gz"
     JSONZIP = "json.zip"
+    JSONL = "jsonl"
+    JSONLGZ = "jsonl.gz"
+    JSONLZIP = "jsonl.zip"
 
 
 DF_READER_MAP = {
@@ -52,6 +55,9 @@ DF_READER_MAP = {
     SupportedTypes.JSON.value: JSONDataFrameReader,
     SupportedTypes.JSONGZ.value: JSONDataFrameReader,
     SupportedTypes.JSONZIP.value: JSONDataFrameReader,
+    SupportedTypes.JSONL.value: JSONDataFrameReader,
+    SupportedTypes.JSONLGZ.value: JSONDataFrameReader,
+    SupportedTypes.JSONLZIP.value: JSONDataFrameReader,
 }
 
 

metadata/readers/file/api_reader.py

@@ -21,12 +21,15 @@ from metadata.generated.schema.security.credentials.bitbucketCredentials import
 from metadata.generated.schema.security.credentials.githubCredentials import (
     GitHubCredentials,
 )
+from metadata.generated.schema.security.credentials.gitlabCredentials import (
+    GitlabCredentials,
+)
 from metadata.readers.file.base import Reader
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
-ReadersCredentials = Union[GitHubCredentials, BitBucketCredentials]
+ReadersCredentials = Union[GitHubCredentials, BitBucketCredentials, GitlabCredentials]
 
 
 class ApiReader(Reader, ABC):
@@ -46,7 +49,7 @@ class ApiReader(Reader, ABC):
         """
         if self._auth_headers is None and self.credentials.token:
             self._auth_headers = {
-                "Authorization": f"Bearer {self.credentials.token.__root__.get_secret_value()}"
+                "Authorization": f"Bearer {self.credentials.token.root.get_secret_value()}"
             }
 
         return self._auth_headers

metadata/readers/file/bitbucket.py

@@ -56,8 +56,8 @@ class BitBucketReader(ApiReader):
                 self._build_url(
                     HOST,
                     UrlParts.REPOS.value,
-                    self.credentials.repositoryOwner.__root__,
-                    self.credentials.repositoryName.__root__,
+                    self.credentials.repositoryOwner.root,
+                    self.credentials.repositoryName.root,
                     UrlParts.SRC.value,
                     self.credentials.branch,
                     path,
@@ -114,8 +114,8 @@ class BitBucketReader(ApiReader):
         url = self._build_url(
             HOST,
             UrlParts.REPOS.value,
-            self.credentials.repositoryOwner.__root__,
-            self.credentials.repositoryName.__root__,
+            self.credentials.repositoryOwner.root,
+            self.credentials.repositoryName.root,
             UrlParts.SRC.value,
             self.credentials.branch,
         )

metadata/readers/file/credentials.py

@@ -26,8 +26,8 @@ def update_repository_name(
     Given an original set of credentials and a new repository name,
     return the updated credentials
     """
-    updated = original.copy(deep=True)
-    updated.repositoryName = RepositoryName(__root__=name)
+    updated = original.model_copy(deep=True)
+    updated.repositoryName = RepositoryName(name)
 
     return updated
 
@@ -47,16 +47,16 @@ def get_credentials_from_url(
     up information, which would still not happen since we work with a single
     token which cannot have permissions on different owners.
     """
-    if original.repositoryOwner.__root__ not in url:
+    if original.repositoryOwner.root not in url:
         logger.warning(
-            f"Default repository owner [{original.repositoryOwner.__root__}] not found in [{url}]."
+            f"Default repository owner [{original.repositoryOwner.root}] not found in [{url}]."
             " We'll use the default reader credentials."
         )
         return original
 
     # Your typical URL is git@bitbucket.org:owner/repo.git
     # or git@github.com:owner/repo.git
-    url_repository = url.split(original.repositoryOwner.__root__ + "/")[-1]
+    url_repository = url.split(original.repositoryOwner.root + "/")[-1]
     repo_name = url_repository.replace(".git", "")
 
     return update_repository_name(original=original, name=repo_name)

metadata/readers/file/github.py

@@ -66,8 +66,8 @@ class GitHubReader(ApiReader):
                 self._build_url(
                     HOST,
                     UrlParts.REPOS.value,
-                    self.credentials.repositoryOwner.__root__,
-                    self.credentials.repositoryName.__root__,
+                    self.credentials.repositoryOwner.root,
+                    self.credentials.repositoryName.root,
                     UrlParts.CONTENTS.value,
                     path,
                 ),
@@ -94,8 +94,8 @@ class GitHubReader(ApiReader):
             self._build_url(
                 HOST,
                 UrlParts.REPOS.value,
-                self.credentials.repositoryOwner.__root__,
-                self.credentials.repositoryName.__root__,
+                self.credentials.repositoryOwner.root,
+                self.credentials.repositoryName.root,
             ),
             headers=self.auth_headers,
             timeout=30,
@@ -118,8 +118,8 @@ class GitHubReader(ApiReader):
                 self._build_url(
                     HOST,
                     UrlParts.REPOS.value,
-                    self.credentials.repositoryOwner.__root__,
-                    self.credentials.repositoryName.__root__,
+                    self.credentials.repositoryOwner.root,
+                    self.credentials.repositoryName.root,
                     "git",
                     "trees",
                     f"{branch}?recursive=1",

metadata/readers/file/gitlab.py

@@ -0,0 +1,172 @@
+#  Copyright 2021 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.
+"""
+Gitlab client to read files with token auth
+"""
+import base64
+import traceback
+from enum import Enum
+from typing import Any, Dict, List, Optional
+from urllib.parse import quote_plus
+
+import requests
+
+from metadata.generated.schema.security.credentials.gitlabCredentials import (
+    GitlabCredentials,
+)
+from metadata.readers.file.api_reader import ApiReader
+from metadata.readers.file.base import ReadException
+from metadata.utils.constants import UTF_8
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+HOST = "https://gitlab.com/api/v4"
+
+
+class UrlParts(Enum):
+    FILES = "files"
+    PROJECTS = "projects"
+    REPOSITORY = "repository"
+    TREE = "tree"
+
+
+class GitlabReader(ApiReader):
+    """
+    Handle calls to the Gitlab API against a repo
+    """
+
+    credentials: GitlabCredentials
+
+    def __init__(self, credentials):
+        super().__init__(credentials)
+        self._encoded_project_path = None
+
+    @property
+    def auth_headers(self) -> Dict[str, str]:
+        """
+        Build the headers to authenticate
+        to the API
+        """
+        if self._auth_headers is None and self.credentials.token:
+            self._auth_headers = {
+                "PRIVATE-TOKEN": self.credentials.token.root.get_secret_value()
+            }
+
+        return self._auth_headers
+
+    @property
+    def encoded_project_path(self) -> str:
+        """
+        Build the URL-encoded project path for the Gitlab API
+        """
+        if (
+            self._encoded_project_path is None
+            and self.credentials.repositoryOwner.root
+            and self.credentials.repositoryName.root
+        ):
+            self._encoded_project_path = quote_plus(
+                "/".join(
+                    [
+                        self.credentials.repositoryOwner.root,
+                        self.credentials.repositoryName.root,
+                    ]
+                )
+            )
+
+        return self._encoded_project_path
+
+    @staticmethod
+    def _decode_content(json_response: Dict[str, Any]) -> str:
+        """
+        Return the content of the response
+
+        If no `content` there, throw the KeyError
+        """
+        return base64.b64decode(json_response["content"]).decode(UTF_8)
+
+    def read(self, path: str, **__) -> str:
+        """
+        Read a file from a Gitlab Repo and return its
+        contents as a string
+        https://docs.gitlab.com/ee/api/repository_files.html
+        """
+        encoded_file_path = quote_plus(path)
+        branch = self._get_default_branch()
+        try:
+            res = requests.get(
+                self._build_url(
+                    HOST,
+                    UrlParts.PROJECTS.value,
+                    self.encoded_project_path,
+                    UrlParts.REPOSITORY.value,
+                    UrlParts.FILES.value,
+                    f"{encoded_file_path}?ref={branch}",
+                ),
+                headers=self.auth_headers,
+                timeout=30,
+            )
+            if res.status_code == 200:
+                return self._decode_content(res.json())
+
+            # If we don't get a 200, raise
+            res.raise_for_status()
+
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            raise ReadException(f"Error fetching file [{path}] from repo: {err}")
+
+        raise ReadException(f"Could not fetch file [{path}] from repo")
+
+    def _get_default_branch(self) -> str:
+        """
+        Get repo default branch
+        """
+        res = requests.get(
+            self._build_url(HOST, UrlParts.PROJECTS.value, self.encoded_project_path),
+            headers=self.auth_headers,
+            timeout=30,
+        )
+        if res.status_code == 200:
+            return res.json().get("default_branch")
+
+        # If we don't get a 200, raise
+        res.raise_for_status()
+        raise RuntimeError("Could not fetch the default branch")
+
+    def _get_tree(self, url: str = None) -> Optional[List[str]]:
+        """
+        Use the Gitlab Repository Tree API to iterate over tree pages recursively
+        """
+        if url is None:
+            url = self._build_url(
+                HOST,
+                UrlParts.PROJECTS.value,
+                self.encoded_project_path,
+                UrlParts.REPOSITORY.value,
+                f"{UrlParts.TREE.value}?recursive=true&pagination=keyset&per_page=100&order_by=path&sort=desc",
+            )
+        res = requests.get(
+            url,
+            headers=self.auth_headers,
+            timeout=30,
+        )
+        if res.status_code == 200:
+            paths = [elem.get("path") for elem in res.json()]
+            if res.links.get("next"):
+                paths.extend(self._get_tree(res.links["next"]["url"]))
+            return paths
+
+        # If we don't get a 200, raise
+        res.raise_for_status()
+
+        return None

metadata/utils/class_helper.py

@@ -103,16 +103,14 @@ def _clean(source_type: str):
     return source_type
 
 
-def get_pipeline_type_from_source_config(
-    source_config_type: SourceConfig.__fields__["config"].type_,
-) -> PipelineType:
+def get_pipeline_type_from_source_config(source_config: SourceConfig) -> PipelineType:
     """From the YAML serviceType, get the Ingestion Pipeline Type"""
     pipeline_type = SOURCE_CONFIG_TYPE_INGESTION.get(
-        source_config_type.__class__.__name__
+        source_config.config.__class__.__name__
     )
     if not pipeline_type:
         raise ValueError(
-            f"Cannot find Pipeline Type for SourceConfig {source_config_type}"
+            f"Cannot find Pipeline Type for SourceConfig {source_config.config}"
         )
     return pipeline_type
 

metadata/utils/credentials.py

@@ -105,7 +105,7 @@ def build_google_credentials_dict(
 
         return {
             "type": gcp_values.type,
-            "project_id": gcp_values.projectId.__root__,
+            "project_id": gcp_values.projectId.root,
             "private_key_id": gcp_values.privateKeyId,
             "private_key": private_key_str,
             "client_email": gcp_values.clientEmail,
@@ -134,7 +134,7 @@ def set_google_credentials(gcp_credentials: GCPCredentials) -> None:
     :param gcp_credentials: GCPCredentials
     """
     if isinstance(gcp_credentials.gcpConfig, GcpCredentialsPath):
-        os.environ[GOOGLE_CREDENTIALS] = str(gcp_credentials.gcpConfig.__root__)
+        os.environ[GOOGLE_CREDENTIALS] = str(gcp_credentials.gcpConfig.path)
         return
 
     if (

metadata/utils/datalake/datalake_utils.py

@@ -452,11 +452,13 @@ class ParquetDataFrameColumnParser:
 
             if parsed_column["dataType"] == DataType.BINARY:
                 try:
-                    data_length = type(column.type).byte_width
-                except AttributeError:
+                    # Either we an int number or -1
+                    data_length = int(type(column.type).byte_width)
+                except Exception as exc:
                     # if the byte width is not specified, we will set it to -1
                     # following pyarrow convention
                     data_length = -1
+                    logger.debug("Could not extract binary field length due to %s", exc)
                 parsed_column["dataLength"] = data_length
 
             if parsed_column["dataType"] == DataType.STRUCT:

metadata/utils/entity_link.py

@@ -57,6 +57,14 @@ def split(str_: str) -> List[str]:
 def get_decoded_column(entity_link: str) -> str:
     """From an URL encoded entity link get the decoded column name
 
+    Examples:
+        >>> get_decoded_column("<#E::table::rds.dev.dbt_jaffle.column_w_space::columns::first_name>")
+        'first name'
+        >>> get_decoded_column("<#E::table::rds.dev.dbt_jaffle.column_w_space::columns::随机的>")
+        '随机的'
+        >>> get_decoded_column("<#E::table::rds.dev.dbt_jaffle.table_w_space>")
+        ''
+
     Args:
         entity_link: entity link
     """

metadata/utils/execution_time_tracker.py

@@ -17,7 +17,7 @@ import threading
 from copy import deepcopy
 from functools import wraps
 from time import perf_counter
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel
 
@@ -81,7 +81,7 @@ class ExecutionTimeTrackerState(metaclass=Singleton):
 
     def __init__(self):
         """Initializes the state and the lock."""
-        self.state = {}
+        self.state: Dict[str, float] = {}
         self.lock = threading.Lock()
 
     def add(self, context: ExecutionTimeTrackerContext, elapsed: float):