openmetadata-ingestion 1.9.2.0__py3-none-any.whl → 1.9.3.0__py3-none-any.whl

metadata/ingestion/models/custom_basemodel_validation.py

@@ -13,19 +13,103 @@ Validation logic for Custom Pydantic BaseModel
 """
 
 import logging
+from enum import Enum
+from typing import Any, Callable, Dict, Optional
 
 logger = logging.getLogger("metadata")
 
-
 RESTRICTED_KEYWORDS = ["::", ">"]
 RESERVED_COLON_KEYWORD = "__reserved__colon__"
 RESERVED_ARROW_KEYWORD = "__reserved__arrow__"
 RESERVED_QUOTE_KEYWORD = "__reserved__quote__"
 
-CREATE_ADJACENT_MODELS = {"ProfilerResponse", "SampleData"}
-NAME_FIELDS = {"EntityName", "str", "ColumnName", "TableData"}
-FETCH_MODELS = {"Table", "CustomColumnName", "DashboardDataModel"}
-FIELD_NAMES = {"name", "columns", "root"}
+
+class TransformDirection(Enum):
+    """Direction of name transformation"""
+
+    ENCODE = "encode"  # For storage (Create operations) - replace separators
+    DECODE = "decode"  # For display (Fetch operations) - revert separators
+
+
+def is_service_level_create_model(model_name: str) -> bool:
+    """
+    Check if a model is a Service-level Create model that should NOT be transformed.
+    Service-level models follow the pattern: Create*ServiceRequest where * is the service name
+    This is scalable and requires no maintenance for new services.
+    """
+    if not model_name.startswith("Create") or not model_name.endswith("ServiceRequest"):
+        return False
+
+    # Extract the middle part (service name) - must not be empty
+    # "CreateServiceRequest" -> middle = "" (invalid)
+    # "CreateDatabaseServiceRequest" -> middle = "Database" (valid)
+    middle = model_name[
+        6:-14
+    ]  # Remove "Create" (6 chars) and "ServiceRequest" (14 chars)
+    return len(middle) > 0
+
+
+# Explicit configuration for entity name transformations
+# This dictionary will be populated lazily to avoid circular imports
+TRANSFORMABLE_ENTITIES: Dict[Any, Dict[str, Any]] = {}
+
+
+def _initialize_transformable_entities():
+    """Initialize the transformable entities dictionary lazily to avoid circular imports"""
+    # Import all model classes here to avoid circular dependency at module load time
+    from metadata.generated.schema.api.data.createDashboardDataModel import (
+        CreateDashboardDataModelRequest,
+    )
+    from metadata.generated.schema.api.data.createTable import CreateTableRequest
+    from metadata.generated.schema.entity.data.dashboardDataModel import (
+        DashboardDataModel,
+    )
+    from metadata.generated.schema.entity.data.table import (
+        ColumnName,
+        ColumnProfile,
+        Table,
+        TableData,
+    )
+    from metadata.profiler.api.models import ProfilerResponse
+    from metadata.utils.entity_link import CustomColumnName
+
+    # Now populate the dictionary with the imported classes
+    TRANSFORMABLE_ENTITIES.update(
+        {
+            # Fetch models - decode reserved keywords back to original characters
+            Table: {
+                "fields": {"name", "columns", "children", "tableConstraints"},
+                "direction": TransformDirection.DECODE,
+            },
+            DashboardDataModel: {
+                "fields": {"name", "columns", "children"},
+                "direction": TransformDirection.DECODE,
+            },
+            CustomColumnName: {
+                "fields": {"root"},
+                "direction": TransformDirection.DECODE,
+            },
+            # Create/Store models - encode special characters to reserved keywords
+            ProfilerResponse: {
+                "fields": {"name", "profile"},
+                "direction": TransformDirection.ENCODE,
+            },
+            TableData: {"fields": {"columns"}, "direction": TransformDirection.ENCODE},
+            ColumnName: {"fields": {"root"}, "direction": TransformDirection.ENCODE},
+            CreateTableRequest: {
+                "fields": {"name", "columns", "children", "tableConstraints"},
+                "direction": TransformDirection.ENCODE,
+            },
+            CreateDashboardDataModelRequest: {
+                "fields": {"name", "columns", "children"},
+                "direction": TransformDirection.ENCODE,
+            },
+            ColumnProfile: {
+                "fields": {"name"},
+                "direction": TransformDirection.ENCODE,
+            },
+        }
+    )
 
 
 def revert_separators(value):
@@ -44,42 +128,84 @@ def replace_separators(value):
     )
 
 
-def validate_name_and_transform(values, modification_method, field_name: str = None):
-    """
-    Validate the name and transform it if needed.
-    """
-    if isinstance(values, str) and field_name in FIELD_NAMES:
-        values = modification_method(values)
-    elif (
-        hasattr(values, "root")
-        and isinstance(values.root, str)
-        and field_name in FIELD_NAMES
+def get_entity_config(model: Optional[Any]) -> Optional[Dict[str, Any]]:
+    """Get transformation configuration for entity"""
+    _initialize_transformable_entities()  # Ensure entities are loaded
+    return TRANSFORMABLE_ENTITIES.get(model)
+
+
+def get_transformer(model: Optional[Any]) -> Optional[Callable]:
+    """Get the appropriate transformer function for model"""
+    config = get_entity_config(model)
+    if not config:
+        return None
+
+    direction = config.get("direction")
+    if direction == TransformDirection.ENCODE:
+        return replace_separators
+    elif direction == TransformDirection.DECODE:
+        return revert_separators
+    return None
+
+
+def transform_all_names(obj, transformer):
+    """Transform all name fields recursively"""
+    if not obj:
+        return
+
+    # Transform name field if it exists (supports both obj.name.root and obj.root)
+    name = getattr(obj, "name", None)
+    if name and hasattr(name, "root") and name.root is not None:
+        name.root = transformer(name.root)
+    elif hasattr(obj, "root") and obj.root is not None:
+        obj.root = transformer(obj.root)
+
+    # Transform nested collections in a single loop each
+    for attr_name in ["columns", "children"]:
+        if hasattr(obj, attr_name):
+            attr_value = getattr(obj, attr_name)
+            if attr_value is not None:
+                for item in attr_value:
+                    transform_all_names(item, transformer)
+
+    # Transform table constraints
+    if hasattr(obj, "tableConstraints"):
+        table_constraints = getattr(obj, "tableConstraints")
+        if table_constraints is not None:
+            for constraint in table_constraints:
+                if hasattr(constraint, "columns"):
+                    constraint.columns = [
+                        transformer(col) for col in constraint.columns
+                    ]
+
+    if transformer == replace_separators and type(name) == str:
+        obj.name = transformer(name)
+
+
+def transform_entity_names(entity: Any, model: Optional[Any]) -> Any:
+    """Transform entity names"""
+    model_name = model.__name__
+    if not entity or (
+        model_name.startswith("Create") and is_service_level_create_model(model_name)
     ):
-        values.root = modification_method(values.root)
-    elif hasattr(type(values), "model_fields"):
-        for key in type(values).model_fields.keys():
-            if getattr(values, key):
-                if getattr(values, key).__class__.__name__ in NAME_FIELDS:
-                    setattr(
-                        values,
-                        key,
-                        validate_name_and_transform(
-                            getattr(values, key),
-                            modification_method=modification_method,
-                            field_name=key,
-                        ),
-                    )
-                elif isinstance(getattr(values, key), list):
-                    setattr(
-                        values,
-                        key,
-                        [
-                            validate_name_and_transform(
-                                item,
-                                modification_method=modification_method,
-                                field_name=key,
-                            )
-                            for item in getattr(values, key)
-                        ],
-                    )
-    return values
+        return entity
+
+    # Root attribute handling
+    if hasattr(entity, "root") and entity.root is not None:
+        entity.root = (
+            replace_separators(entity.root)
+            if model_name.startswith("Create")
+            else revert_separators(entity.root)
+        )
+        return entity
+
+    # Get model-specific transformer
+    transformer = get_transformer(model)
+    if not transformer:
+        # Fallback to original logic for backward compatibility
+        transformer = (
+            replace_separators if model_name.startswith("Create") else revert_separators
+        )
+
+    transform_all_names(entity, transformer)
+    return entity

metadata/ingestion/models/custom_pydantic.py

@@ -26,13 +26,7 @@ from pydantic.types import SecretStr
 from pydantic_core.core_schema import SerializationInfo
 from typing_extensions import Annotated
 
-from metadata.ingestion.models.custom_basemodel_validation import (
-    CREATE_ADJACENT_MODELS,
-    FETCH_MODELS,
-    replace_separators,
-    revert_separators,
-    validate_name_and_transform,
-)
+from metadata.ingestion.models.custom_basemodel_validation import transform_entity_names
 
 logger = logging.getLogger("metadata")
 
@@ -75,25 +69,18 @@ class BaseModel(PydanticBaseModel):
     @classmethod
     def parse_name(cls, values):  # pylint: disable=inconsistent-return-statements
         """
-        Primary entry point to process values based on their class.
+        Transform entity names using hybrid configuration system.
         """
 
         if not values:
-            return
+            return values
 
         try:
-
-            if cls.__name__ in CREATE_ADJACENT_MODELS or cls.__name__.startswith(
-                "Create"
-            ):
-                values = validate_name_and_transform(values, replace_separators)
-            elif cls.__name__ in FETCH_MODELS:
-                values = validate_name_and_transform(values, revert_separators)
-
+            # Try new hybrid system first
+            return transform_entity_names(entity=values, model=cls)
         except Exception as exc:
             logger.warning("Exception while parsing Basemodel: %s", exc)
-            raise exc
-        return values
+            return values
 
     def model_dump_json(  # pylint: disable=too-many-arguments
         self,

metadata/ingestion/source/database/databricks/connection.py

@@ -56,11 +56,17 @@ class DatabricksEngineWrapper:
         self.inspector = inspect(engine)
         self.schemas = None
         self.first_schema = None
+        self.first_catalog = None
 
-    def get_schemas(self):
+    def get_schemas(self, schema_name: Optional[str] = None):
         """Get schemas and cache them"""
+        if schema_name is not None:
+            with self.engine.connect() as connection:
+                connection.execute(f"USE CATALOG `{self.first_catalog}`")
+            self.first_schema = schema_name
+            return [schema_name]
         if self.schemas is None:
-            self.schemas = self.inspector.get_schema_names()
+            self.schemas = self.inspector.get_schema_names(database=self.first_catalog)
             if self.schemas:
                 # Find the first schema that's not a system schema
                 for schema in self.schemas:
@@ -81,7 +87,11 @@ class DatabricksEngineWrapper:
         if self.first_schema is None:
             self.get_schemas()  # This will set first_schema
         if self.first_schema:
-            return self.inspector.get_table_names(self.first_schema)
+            with self.engine.connect() as connection:
+                tables = connection.execute(
+                    f"SHOW TABLES IN `{self.first_catalog}`.`{self.first_schema}`"
+                )
+            return tables
         return []
 
     def get_views(self):
@@ -89,9 +99,27 @@ class DatabricksEngineWrapper:
         if self.first_schema is None:
             self.get_schemas()  # This will set first_schema
         if self.first_schema:
-            return self.inspector.get_view_names(self.first_schema)
+            with self.engine.connect() as connection:
+                views = connection.execute(
+                    f"SHOW VIEWS IN `{self.first_catalog}`.`{self.first_schema}`"
+                )
+            return views
         return []
 
+    def get_catalogs(self, catalog_name: Optional[str] = None):
+        """Get catalogs"""
+        catalogs = []
+        if catalog_name is not None:
+            self.first_catalog = catalog_name
+            return [catalog_name]
+        with self.engine.connect() as connection:
+            catalogs = connection.execute(DATABRICKS_GET_CATALOGS).fetchall()
+            for catalog in catalogs:
+                if catalog[0] != "__databricks_internal":
+                    self.first_catalog = catalog[0]
+                    break
+        return catalogs
+
 
 def get_connection_url(connection: DatabricksConnection) -> str:
     url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
@@ -144,13 +172,13 @@ def test_connection(
 
     test_fn = {
         "CheckAccess": partial(test_connection_engine_step, connection),
-        "GetSchemas": engine_wrapper.get_schemas,
+        "GetSchemas": partial(
+            engine_wrapper.get_schemas, schema_name=service_connection.databaseSchema
+        ),
         "GetTables": engine_wrapper.get_tables,
         "GetViews": engine_wrapper.get_views,
         "GetDatabases": partial(
-            test_database_query,
-            engine=connection,
-            statement=DATABRICKS_GET_CATALOGS,
+            engine_wrapper.get_catalogs, catalog_name=service_connection.catalog
         ),
         "GetQueries": partial(
             test_database_query,

metadata/ingestion/source/database/iceberg/models.py

@@ -51,24 +51,33 @@ class IcebergTable(BaseModel):
         """Responsible for parsing the needed information from a PyIceberg Table."""
         iceberg_columns = table.schema().fields
 
+        partition_columns = []
+        for partition in table.spec().fields:
+            partition_column_name = get_column_from_partition(
+                iceberg_columns, partition
+            )
+            column_partition_type = get_column_partition_type(
+                iceberg_columns, partition
+            )
+
+            if not (partition_column_name and column_partition_type):
+                continue
+
+            partition_columns.append(
+                PartitionColumnDetails(
+                    columnName=partition_column_name,
+                    intervalType=column_partition_type,
+                    interval=None,
+                )
+            )
+
         return IcebergTable(
             name=name,
             tableType=table_type,
             description=table.properties.get("comment"),
             owners=owners,
             columns=[IcebergColumnParser.parse(column) for column in iceberg_columns],
-            tablePartition=TablePartition(
-                columns=[
-                    PartitionColumnDetails(
-                        columnName=get_column_from_partition(
-                            iceberg_columns, partition
-                        ),
-                        intervalType=get_column_partition_type(
-                            iceberg_columns, partition
-                        ),
-                        interval=None,
-                    )
-                    for partition in table.spec().fields
-                ]
+            tablePartition=(
+                TablePartition(columns=partition_columns) if partition_columns else None
             ),
         )

metadata/ingestion/source/database/snowflake/metadata.py

@@ -512,13 +512,8 @@ class SnowflakeSource(
                         )
                     )
                 except Exception as inner_exc:
-                    yield Either(
-                        left=StackTraceError(
-                            name="Tags and Classifications",
-                            error=f"Failed to fetch tags due to [{inner_exc}]",
-                            stackTrace=traceback.format_exc(),
-                        )
-                    )
+                    logger.debug(traceback.format_exc())
+                    logger.error(f"Failed to fetch tags due to [{inner_exc}]")
 
             for res in result:
                 row = list(res)

metadata/ingestion/source/database/unitycatalog/connection.py

@@ -111,16 +111,15 @@ def test_connection(
 
     def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
         for catalog in connection.catalogs.list():
-            table_obj.catalog_name = catalog.name
-            break
+            if catalog.name != "__databricks_internal":
+                table_obj.catalog_name = catalog.name
+                return
 
     def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
-        for catalog in connection.catalogs.list():
-            for schema in connection.schemas.list(catalog_name=catalog.name):
-                if schema.name:
-                    table_obj.schema_name = schema.name
-                    table_obj.catalog_name = catalog.name
-                    return
+        for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
+            if schema.name:
+                table_obj.schema_name = schema.name
+                return
 
     def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
         if table_obj.catalog_name and table_obj.schema_name:

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -92,7 +92,7 @@ logger = ingestion_logger()
 UNITY_CATALOG_TAG = "UNITY CATALOG TAG"
 UNITY_CATALOG_TAG_CLASSIFICATION = "UNITY CATALOG TAG CLASSIFICATION"
 
-
+# pylint: disable=protected-access
 class UnitycatalogSource(
     ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
 ):
@@ -306,6 +306,8 @@ class UnitycatalogSource(
                 if table.table_type:
                     if table.table_type.value.lower() == TableType.View.value.lower():
                         table_type: TableType = TableType.View
+                    if table.table_type.value.lower() == "materialized_view":
+                        table_type: TableType = TableType.MaterializedView
                     elif (
                         table.table_type.value.lower()
                         == TableType.External.value.lower()
@@ -427,7 +429,6 @@ class UnitycatalogSource(
             )
             if referred_table_fqn:
                 for parent_column in column.parent_columns:
-                    # pylint: disable=protected-access
                     col_fqn = fqn._build(referred_table_fqn, parent_column, quote=False)
                     if col_fqn:
                         referred_column_fqns.append(FullyQualifiedEntityName(col_fqn))
@@ -563,7 +564,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,
@@ -617,7 +618,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -371,7 +371,11 @@ class AirflowSource(PipelineServiceSource):
         offset = 0  # Start
 
         while True:
-            paginated_query = session_query.limit(limit).offset(offset)
+            paginated_query = (
+                session_query.order_by(SerializedDagModel.dag_id.asc())
+                .limit(limit)
+                .offset(offset)
+            )
             results = paginated_query.all()
             if not results:
                 break

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

@@ -23,6 +23,7 @@ from metadata.generated.schema.entity.data.table import (
     ColumnName,
     DataType,
     SystemProfile,
+    TableType,
 )
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
@@ -51,6 +52,11 @@ class DatabricksProfilerInterface(SQAProfilerInterface):
         *args,
         **kwargs,
     ) -> List[SystemProfile]:
+        if self.table_entity.tableType in (TableType.View, TableType.MaterializedView):
+            logger.debug(
+                f"Skipping {metrics.name()} metric for view {runner.table_name}"
+            )
+            return []
         logger.debug(f"Computing {metrics.name()} metric for {runner.table_name}")
         self.system_metrics_class = cast(
             Type[DatabricksSystemMetricsComputer], self.system_metrics_class

metadata/utils/fqn.py

@@ -15,6 +15,7 @@ ES indexes definitions
 """
 import hashlib
 import re
+import traceback
 from typing import Dict, List, Optional, Type, TypeVar, Union
 
 from antlr4.CommonTokenStream import CommonTokenStream
@@ -53,6 +54,9 @@ from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.dispatch import class_register
 from metadata.utils.elasticsearch import get_entity_from_es_result
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -144,12 +148,32 @@ def build(
     :param kwargs: required to build the FQN
     :return: FQN as a string
     """
+    # Transform table_name and column_name if they exist and contain special characters
+    if kwargs.get("table_name") or kwargs.get("column_name"):
+        from metadata.ingestion.models.custom_basemodel_validation import (  # pylint: disable=import-outside-toplevel
+            replace_separators,
+        )
+
+        table_name = kwargs.get("table_name")
+        if table_name and isinstance(table_name, str):
+            kwargs["table_name"] = replace_separators(table_name)
+
+        column_name = kwargs.get("column_name")
+        if column_name and isinstance(column_name, str):
+            kwargs["column_name"] = replace_separators(column_name)
+
     func = fqn_build_registry.registry.get(entity_type.__name__)
-    if not func:
+    try:
+        if not func:
+            raise FQNBuildingException(
+                f"Invalid Entity Type {entity_type.__name__}. FQN builder not implemented."
+            )
+        return func(metadata, **kwargs)
+    except Exception as e:
+        logger.debug(traceback.format_exc())
         raise FQNBuildingException(
-            f"Invalid Entity Type {entity_type.__name__}. FQN builder not implemented."
+            f"Error building FQN for {entity_type.__name__}: {e}"
         )
-    return func(metadata, **kwargs)
 
 
 @fqn_build_registry.add(Table)