openmetadata-ingestion 1.4.0.0rc1__py3-none-any.whl → 1.4.0.0rc3__py3-none-any.whl

metadata/ingestion/source/pipeline/kafkaconnect/models.py

@@ -0,0 +1,62 @@
+#  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.
+
+"""
+KafkaConnect Source Model module
+"""
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+from metadata.generated.schema.entity.data.container import Container
+from metadata.generated.schema.entity.data.table import Table
+
+
+class KafkaConnectTasks(BaseModel):
+    id: int = Field(..., description="ID of the task")
+    state: Optional[str] = Field(
+        default="UNASSIGNED", description="State of the task (e.g., RUNNING, STOPPED)"
+    )
+    worker_id: Optional[str] = Field(
+        ..., description="ID of the worker running the task"
+    )
+
+
+class KafkaConnectTopics(BaseModel):
+    name: str = Field(..., description="Name of the topic (e.g., random-source-avro)")
+
+
+class KafkaConnectPipelineDetails(BaseModel):
+    name: str = Field(
+        ..., description="Name of the status source (e.g., random-source-json)"
+    )
+    status: Optional[str] = Field(
+        default="UNASSIGNED",
+        description="State of the connector (e.g., RUNNING, STOPPED)",
+    )
+    tasks: Optional[List[KafkaConnectTasks]]
+    topics: Optional[List[KafkaConnectTopics]]
+    conn_type: Optional[str] = Field(..., alias="type")
+
+
+class KafkaConnectDatasetDetails(BaseModel):
+    table: Optional[str]
+    database: Optional[str]
+    container_name: Optional[str]
+
+    @property
+    def dataset_type(self):
+        if self.table or self.database:
+            return Table
+        if self.container_name:
+            return Container
+        return None

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

@@ -59,8 +59,17 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
+    def test_buckets(connection: S3Connection, client: S3ObjectStoreClient):
+        if connection.bucketNames:
+            for bucket_name in connection.bucketNames:
+                client.s3_client.list_objects(Bucket=bucket_name)
+            return
+        client.s3_client.list_buckets()
+
     test_fn = {
-        "ListBuckets": client.s3_client.list_buckets,
+        "ListBuckets": partial(
+            test_buckets, client=client, connection=service_connection
+        ),
         "GetMetrics": partial(
             client.cloudwatch_client.list_metrics, Namespace="AWS/S3"
         ),

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

@@ -126,10 +126,7 @@ class S3Source(StorageServiceSource):
                         # ingest all the relevant valid paths from it
                         yield from self._generate_structured_containers(
                             bucket_response=bucket_response,
-                            entries=self._manifest_entries_to_metadata_entries_by_container(
-                                container_name=bucket_name,
-                                manifest=self.global_manifest,
-                            ),
+                            entries=manifest_entries_for_current_bucket,
                             parent=parent_entity,
                         )
                         # nothing else do to for the current bucket, skipping to the next
@@ -142,23 +139,6 @@ class S3Source(StorageServiceSource):
                         entries=metadata_config.entries,
                         parent=parent_entity,
                     )
-                    for metadata_entry in metadata_config.entries:
-                        logger.info(
-                            f"Extracting metadata from path {metadata_entry.dataPath.strip(KEY_SEPARATOR)} "
-                            f"and generating structured container"
-                        )
-                        structured_container: Optional[
-                            S3ContainerDetails
-                        ] = self._generate_container_details(
-                            bucket_response=bucket_response,
-                            metadata_entry=metadata_entry,
-                            parent=EntityReference(
-                                id=self._bucket_cache[bucket_response.name].id.__root__,
-                                type="container",
-                            ),
-                        )
-                        if structured_container:
-                            yield structured_container
 
             except ValidationError as err:
                 self.status.failed(
@@ -223,7 +203,9 @@ class S3Source(StorageServiceSource):
                 return S3ContainerDetails(
                     name=metadata_entry.dataPath.strip(KEY_SEPARATOR),
                     prefix=prefix,
-                    creation_date=bucket_response.creation_date.isoformat(),
+                    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
                     ),
@@ -270,6 +252,11 @@ class S3Source(StorageServiceSource):
     def fetch_buckets(self) -> List[S3BucketResponse]:
         results: List[S3BucketResponse] = []
         try:
+            if self.service_connection.bucketNames:
+                return [
+                    S3BucketResponse(Name=bucket_name)
+                    for bucket_name in self.service_connection.bucketNames
+                ]
             # No pagination required, as there is a hard 1000 limit on nr of buckets per aws account
             for bucket in self.s3_client.list_buckets().get("Buckets") or []:
                 if filter_by_container(
@@ -335,7 +322,9 @@ class S3Source(StorageServiceSource):
         return S3ContainerDetails(
             name=bucket_response.name,
             prefix=KEY_SEPARATOR,
-            creation_date=bucket_response.creation_date.isoformat(),
+            creation_date=bucket_response.creation_date.isoformat()
+            if bucket_response.creation_date
+            else None,
             number_of_objects=self._fetch_metric(
                 bucket_name=bucket_response.name, metric=S3Metric.NUMBER_OF_OBJECTS
             ),

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

@@ -33,8 +33,8 @@ class S3BucketResponse(BaseModel):
         extra = Extra.forbid
 
     name: str = Field(..., description="Bucket name", alias="Name")
-    creation_date: datetime = Field(
-        ...,
+    creation_date: Optional[datetime] = Field(
+        None,
         description="Timestamp of Bucket creation in ISO format",
         alias="CreationDate",
     )
@@ -67,8 +67,8 @@ class S3ContainerDetails(BaseModel):
         ...,
         description="Data Model of the container",
     )
-    creation_date: str = Field(
-        ...,
+    creation_date: Optional[str] = Field(
+        None,
         description="Timestamp of Bucket creation in ISO format",
     )
     parent: Optional[EntityReference] = Field(

metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py

@@ -45,14 +45,12 @@ class DatabricksProfilerInterface(SQAProfilerInterface):
             splitted_result = result.split(".", 2)[-1].split(".")
             result = ".".join(result.split(".", 2)[:-1])
             result += "." + "`.`".join(splitted_result)
-
         return result
 
-    HiveCompiler.visit_column = visit_column
-
     def __init__(self, service_connection_config, **kwargs):
         super().__init__(service_connection_config=service_connection_config, **kwargs)
         self.set_catalog(self.session)
+        HiveCompiler.visit_column = DatabricksProfilerInterface.visit_column
 
     def _get_struct_columns(self, columns: List[OMColumn], parent: str):
         """Get struct columns"""

metadata/profiler/processor/metric_filter.py

@@ -168,7 +168,9 @@ class MetricFilter:
         if not isinstance(column, SQALikeColumn):
             mapper = converter_registry[service_type]
             sqa_to_om_types = mapper.map_sqa_to_om_types()
-            om_data_types: Optional[Set] = sqa_to_om_types.get(column.type, None)
+            om_data_types: Optional[Set] = sqa_to_om_types.get(
+                column.type.__class__, None
+            )
         else:
             om_data_types = {column.type}
 

metadata/utils/elasticsearch.py

@@ -18,8 +18,10 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.analytics.reportData import ReportData
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
+from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.glossary import Glossary
@@ -27,6 +29,8 @@ from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.searchIndex import SearchIndex
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
@@ -42,9 +46,13 @@ ES_INDEX_MAP = {
     Database.__name__: "database_search_index",
     DatabaseSchema.__name__: "database_schema_search_index",
     DatabaseService.__name__: "database_service_search_index",
+    StoredProcedure.__name__: "stored_procedure_search_index",
+    DashboardDataModel.__name__: "dashboard_data_model_search_index",
+    SearchIndex.__name__: "search_entity_search_index",
     Team.__name__: "team_search_index",
     User.__name__: "user_search_index",
     Dashboard.__name__: "dashboard_search_index",
+    Chart.__name__: "chart_search_index",
     Topic.__name__: "topic_search_index",
     Pipeline.__name__: "pipeline_search_index",
     Glossary.__name__: "glossary_search_index",

metadata/utils/messaging_utils.py

@@ -0,0 +1,43 @@
+#  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.
+
+"""
+Helpers module for messaging sources
+"""
+
+import re
+import traceback
+from typing import Optional
+
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
+
+
+def merge_and_clean_protobuf_schema(schema_text: Optional[str]) -> Optional[str]:
+    """
+    Remove the import and extra syntax lines for a schema with references
+    """
+    try:
+        lines = schema_text.splitlines() if schema_text else []
+        new_lines = []
+        for i, line in enumerate(lines):
+            if not re.search(r'import ".*";', line) and not re.search(
+                r"option .*;", line
+            ):
+                if re.search(r'\s*syntax\s*=\s*"proto\d+";\s*', line) and i != 0:
+                    continue
+                new_lines.append(line)
+        return "\n".join(new_lines)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Failed to merge and clean protobuf schema: {exc}")
+    return None

metadata/utils/sqlalchemy_utils.py

@@ -12,10 +12,15 @@
 """
 Module for sqlalchemy dialect utils
 """
-
+import traceback
 from typing import Dict, Optional, Tuple
 
 from sqlalchemy.engine import Engine, reflection
+from sqlalchemy.schema import CreateTable, MetaData
+
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 @reflection.cache
@@ -125,3 +130,42 @@ def convert_numpy_to_list(data):
     if isinstance(data, dict):
         return {key: convert_numpy_to_list(value) for key, value in data.items()}
     return data
+
+
+@reflection.cache
+def get_all_table_ddls(
+    self, connection, query, schema_name, **kw
+):  # pylint: disable=unused-argument
+    """
+    Method to fetch ddl of all available tables
+    """
+    try:
+        self.all_table_ddls: Dict[Tuple[str, str], str] = {}
+        self.current_db: str = schema_name
+        meta = MetaData()
+        meta.reflect(bind=connection.engine, schema=schema_name)
+        for table in meta.sorted_tables or []:
+            self.all_table_ddls[(table.schema, table.name)] = str(CreateTable(table))
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Failed to get table ddls for {schema_name}: {exc}")
+
+
+def get_table_ddl_wrapper(
+    self, connection, query, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    if not hasattr(self, "all_table_ddls") or self.current_db != schema:
+        self.get_all_table_ddls(connection, query, schema)
+    return self.all_table_ddls.get((schema, table_name), "")
+
+
+def get_table_ddl(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    return get_table_ddl_wrapper(
+        self,
+        connection=connection,
+        query=None,
+        table_name=table_name,
+        schema=schema,
+    )

metadata/utils/tag_utils.py

@@ -85,7 +85,7 @@ def get_ometa_tag_and_classification(
 def get_tag_label(
     metadata: OpenMetadata,
     tag_name: str,
-    classification_name: str,
+    classification_name: Optional[str],
     tag_type: Union[Tag, GlossaryTerm] = Tag,
 ) -> Optional[TagLabel]:
     """
@@ -127,7 +127,7 @@ def get_tag_label(
 def get_tag_labels(
     metadata: OpenMetadata,
     tags: List[str],
-    classification_name: Optional[str],
+    classification_name: Optional[str] = None,
     include_tags: bool = True,
     tag_type: Union[Tag, GlossaryTerm] = Tag,
 ) -> Optional[List[TagLabel]]: