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

metadata/great_expectations/utils/ometa_config_handler.py

@@ -86,4 +86,4 @@ def render_template(environment: Environment, template_file: str = "config.yml")
 
 def create_ometa_connection_obj(config: str) -> OpenMetadataConnection:
     """Create OpenMetadata connection"""
-    return OpenMetadataConnection.parse_obj(yaml.safe_load(config))
+    return OpenMetadataConnection.model_validate(yaml.safe_load(config))

metadata/ingestion/api/common.py

@@ -13,7 +13,7 @@ Common definitions for configuration management
 """
 from typing import Any, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from metadata.utils.logger import ingestion_logger
 
@@ -27,13 +27,12 @@ Entity = TypeVar("Entity", bound=BaseModel)
 
 
 class ConfigModel(BaseModel):
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class DynamicTypedConfig(ConfigModel):
     type: str
-    config: Optional[Any]
+    config: Optional[Any] = None
 
 
 class WorkflowExecutionError(Exception):

metadata/ingestion/api/delete.py

@@ -43,7 +43,7 @@ def delete_entity_from_source(
     try:
         entity_state = metadata.list_all_entities(entity=entity_type, params=params)
         for entity in entity_state:
-            if str(entity.fullyQualifiedName.__root__) not in entity_source_state:
+            if str(entity.fullyQualifiedName.root) not in entity_source_state:
                 yield Either(
                     right=DeleteEntity(
                         entity=entity,

metadata/ingestion/api/models.py

@@ -13,7 +13,7 @@ Generic models
 """
 from typing import Generic, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -25,11 +25,9 @@ T = TypeVar("T")
 
 
 class Either(BaseModel, Generic[T]):
-    """
-    Any execution should return us Either an Entity of an error for us to handle
-    - left: Optional error we encounter
-    - right: Correct instance of an Entity
-    """
+    """Any execution should return us Either an Entity of an error for us to handle"""
 
-    left: Optional[StackTraceError]
-    right: Optional[T]
+    left: Optional[StackTraceError] = Field(
+        None, description="Error encountered during execution"
+    )
+    right: Optional[T] = Field(None, description="Correct instance of an Entity")

metadata/ingestion/api/parser.py

@@ -261,7 +261,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
         if len(err.get("loc")) == 1
         else f"Extra parameter in {err.get('loc')}"
         for err in validation_error.errors()
-        if err.get("type") == "value_error.extra"
+        if err.get("type") == "extra_forbidden"
     ]
 
     extra_fields = [
@@ -269,7 +269,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
         if len(err.get("loc")) == 1
         else f"Missing parameter in {err.get('loc')}"
         for err in validation_error.errors()
-        if err.get("type") == "value_error.missing"
+        if err.get("type") == "missing"
     ]
 
     invalid_fields = [
@@ -277,7 +277,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
         if len(err.get("loc")) == 1
         else f"Invalid parameter value for {err.get('loc')}"
         for err in validation_error.errors()
-        if err.get("type") not in ("value_error.missing", "value_error.extra")
+        if err.get("type") not in ("missing", "extra")
     ]
 
     return "\t - " + "\n\t - ".join(missing_fields + extra_fields + invalid_fields)
@@ -291,10 +291,10 @@ def _unsafe_parse_config(config: dict, cls: Type[T], message: str) -> None:
     logger.debug(f"Parsing message: [{message}]")
     # Parse the service connection dictionary with the scoped class
     try:
-        cls.parse_obj(config)
+        cls.model_validate(config)
     except ValidationError as err:
         logger.debug(
-            f"The supported properties for {cls.__name__} are {list(cls.__fields__.keys())}"
+            f"The supported properties for {cls.__name__} are {list(cls.model_fields.keys())}"
         )
         raise err
 
@@ -309,13 +309,13 @@ def _unsafe_parse_dbt_config(config: dict, cls: Type[T], message: str) -> None:
         # Parse the oneOf config types of dbt to check
         dbt_config_type = config["dbtConfigSource"]["dbtConfigType"]
         dbt_config_class = DBT_CONFIG_TYPE_MAP.get(dbt_config_type)
-        dbt_config_class.parse_obj(config["dbtConfigSource"])
+        dbt_config_class.model_validate(config["dbtConfigSource"])
 
         # Parse the entire dbtPipeline object
-        cls.parse_obj(config)
+        cls.model_validate(config)
     except ValidationError as err:
         logger.debug(
-            f"The supported properties for {cls.__name__} are {list(cls.__fields__.keys())}"
+            f"The supported properties for {cls.__name__} are {list(cls.model_fields.keys())}"
         )
         raise err
 
@@ -437,21 +437,17 @@ def parse_workflow_config_gracefully(
     """
 
     try:
-        workflow_config = OpenMetadataWorkflowConfig.parse_obj(config_dict)
+        workflow_config = OpenMetadataWorkflowConfig.model_validate(config_dict)
         return workflow_config
 
     except ValidationError as original_error:
         try:
             parse_workflow_source(config_dict)
-            WorkflowConfig.parse_obj(config_dict["workflowConfig"])
+            WorkflowConfig.model_validate(config_dict["workflowConfig"])
         except (ValidationError, InvalidWorkflowException) as scoped_error:
             if isinstance(scoped_error, ValidationError):
                 # Let's catch validations of internal Workflow models, not the Workflow itself
-                object_error = (
-                    scoped_error.model.__name__
-                    if scoped_error.model is not None
-                    else "workflow"
-                )
+                object_error = scoped_error.title or "workflow"
                 raise ParsingConfigurationError(
                     f"We encountered an error parsing the configuration of your {object_error}.\n"
                     "You might need to review your config based on the original cause of this failure:\n"
@@ -483,7 +479,7 @@ def parse_ingestion_pipeline_config_gracefully(
     """
 
     try:
-        ingestion_pipeline = IngestionPipeline.parse_obj(config_dict)
+        ingestion_pipeline = IngestionPipeline.model_validate(config_dict)
         return ingestion_pipeline
 
     except ValidationError:
@@ -518,7 +514,7 @@ def parse_automation_workflow_gracefully(
     """
 
     try:
-        automation_workflow = AutomationWorkflow.parse_obj(config_dict)
+        automation_workflow = AutomationWorkflow.model_validate(config_dict)
         return automation_workflow
 
     except ValidationError:

metadata/ingestion/api/status.py

@@ -13,9 +13,10 @@ Status output utilities
 """
 import pprint
 import time
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import AfterValidator, BaseModel, Field
+from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -26,22 +27,36 @@ from metadata.utils.logger import get_log_name, ingestion_logger
 logger = ingestion_logger()
 
 
+MAX_STACK_TRACE_LENGTH = 1_000_000
+TruncatedStr = Annotated[
+    Optional[str], AfterValidator(lambda v: v[:MAX_STACK_TRACE_LENGTH] if v else None)
+]
+
+
+class TruncatedStackTraceError(StackTraceError):
+    """
+    Update StackTraceError to limit the payload size,
+    since some connectors can make it explode
+    """
+
+    error: TruncatedStr
+    stackTrace: TruncatedStr = None
+
+
 class Status(BaseModel):
     """
     Class to handle status
     """
 
-    source_start_time: Any
-
-    records: List[Any] = Field(default_factory=list)
-    updated_records: List[Any] = Field(default_factory=list)
-    warnings: List[Any] = Field(default_factory=list)
-    filtered: List[Dict[str, str]] = Field(default_factory=list)
-    failures: List[StackTraceError] = Field(default_factory=list)
+    source_start_time: float = Field(
+        default_factory=lambda: time.time()  # pylint: disable=unnecessary-lambda
+    )
 
-    def __init__(self, **data):
-        super().__init__(**data)
-        self.source_start_time = time.time()
+    records: Annotated[List[Any], Field(default_factory=list)]
+    updated_records: Annotated[List[Any], Field(default_factory=list)]
+    warnings: Annotated[List[Any], Field(default_factory=list)]
+    filtered: Annotated[List[Dict[str, str]], Field(default_factory=list)]
+    failures: Annotated[List[TruncatedStackTraceError], Field(default_factory=list)]
 
     def scanned(self, record: Any) -> None:
         """
@@ -75,7 +90,14 @@ class Status(BaseModel):
         """
         logger.warning(error.error)
         logger.debug(error.stackTrace)
-        self.failures.append(error)
+        # Truncate StackTrace to avoid payload explosion
+        self.failures.append(
+            TruncatedStackTraceError(
+                name=error.name,
+                error=error.error,
+                stackTrace=error.stackTrace,
+            )
+        )
 
     def fail_all(self, failures: List[StackTraceError]) -> None:
         """

metadata/ingestion/api/topology_runner.py

@@ -328,7 +328,8 @@ class TopologyRunnerMixin(Generic[C]):
         """
         return PatchRequest(
             original_entity=original_entity,
-            new_entity=original_entity.copy(update=create_request.__dict__),
+            new_entity=original_entity.model_copy(update=create_request.__dict__),
+            override_metadata=self.source_config.overrideMetadata,
         )
 
     @singledispatchmethod

metadata/ingestion/bulksink/metadata_usage.py

@@ -39,6 +39,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
     StackTraceError,
 )
 from metadata.generated.schema.entity.teams.user import User
+from metadata.generated.schema.type.basic import Timestamp
 from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle
 from metadata.generated.schema.type.tableUsageCount import TableColumn, TableUsageCount
 from metadata.generated.schema.type.usageRequest import UsageRequest
@@ -99,7 +100,7 @@ class MetadataUsageBulkSink(BulkSink):
         metadata: OpenMetadata,
         pipeline_name: Optional[str] = None,
     ):
-        config = MetadataUsageSinkConfig.parse_obj(config_dict)
+        config = MetadataUsageSinkConfig.model_validate(config_dict)
         return cls(config, metadata)
 
     def __populate_table_usage_map(
@@ -109,8 +110,8 @@ class MetadataUsageBulkSink(BulkSink):
         Method Either initialise the map data or
         update existing data with information from new queries on the same table
         """
-        if not self.table_usage_map.get(table_entity.id.__root__):
-            self.table_usage_map[table_entity.id.__root__] = {
+        if not self.table_usage_map.get(table_entity.id.root):
+            self.table_usage_map[table_entity.id.root] = {
                 "table_entity": table_entity,
                 "usage_count": table_usage.count,
                 "usage_date": table_usage.date,
@@ -118,7 +119,7 @@ class MetadataUsageBulkSink(BulkSink):
                 "database_schema": table_usage.databaseSchema,
             }
         else:
-            self.table_usage_map[table_entity.id.__root__][
+            self.table_usage_map[table_entity.id.root][
                 "usage_count"
             ] += table_usage.count
 
@@ -139,10 +140,10 @@ class MetadataUsageBulkSink(BulkSink):
                     value_dict["table_entity"], table_usage_request
                 )
                 logger.info(
-                    f"Successfully table usage published for {value_dict['table_entity'].fullyQualifiedName.__root__}"
+                    f"Successfully table usage published for {value_dict['table_entity'].fullyQualifiedName.root}"
                 )
                 self.status.scanned(
-                    f"Table: {value_dict['table_entity'].fullyQualifiedName.__root__}"
+                    f"Table: {value_dict['table_entity'].fullyQualifiedName.root}"
                 )
             except ValidationError as err:
                 logger.debug(traceback.format_exc())
@@ -150,13 +151,13 @@ class MetadataUsageBulkSink(BulkSink):
                     f"Cannot construct UsageRequest from {value_dict['table_entity']}: {err}"
                 )
             except Exception as exc:
-                name = value_dict["table_entity"].fullyQualifiedName.__root__
+                name = value_dict["table_entity"].fullyQualifiedName.root
                 error = f"Failed to update usage for {name} :{exc}"
                 logger.debug(traceback.format_exc())
                 logger.warning(error)
                 self.status.failed(
                     StackTraceError(
-                        name=value_dict["table_entity"].fullyQualifiedName.__root__,
+                        name=value_dict["table_entity"].fullyQualifiedName.root,
                         error=f"Failed to update usage for {name} :{exc}",
                         stackTrace=traceback.format_exc(),
                     )
@@ -255,7 +256,7 @@ class MetadataUsageBulkSink(BulkSink):
                         )
                     )
                 except Exception as exc:
-                    name = table_entity.name.__root__
+                    name = table_entity.name.root
                     error = (
                         f"Error getting usage and join information for {name}: {exc}"
                     )
@@ -281,8 +282,10 @@ class MetadataUsageBulkSink(BulkSink):
         """
         Method to get Table Joins
         """
+        # TODO: Clean up how we are passing dates from query parsing to here to use timestamps instead of strings
+        start_date = datetime.fromtimestamp(int(table_usage.date) / 1000)
         table_joins: TableJoins = TableJoins(
-            columnJoins=[], directTableJoins=[], startDate=table_usage.date
+            columnJoins=[], directTableJoins=[], startDate=start_date
         )
         column_joins_dict = {}
         for column_join in table_usage.joins:
@@ -317,7 +320,7 @@ class MetadataUsageBulkSink(BulkSink):
             key_name = get_column_fqn(table_entity=table_entity, column=key)
             if not key_name:
                 logger.warning(
-                    f"Could not find column {key} in table {table_entity.fullyQualifiedName.__root__}"
+                    f"Could not find column {key} in table {table_entity.fullyQualifiedName.root}"
                 )
                 continue
             table_joins.columnJoins.append(
@@ -370,15 +373,15 @@ class MetadataUsageBulkSink(BulkSink):
                 query_type = get_query_type(create_query=create_query)
                 if query_type:
                     access_details = AccessDetails(
-                        timestamp=create_query.queryDate.__root__,
+                        timestamp=Timestamp(create_query.queryDate.root),
                         accessedBy=user,
                         accessedByAProcess=process_user,
                     )
                     life_cycle_attr = getattr(life_cycle, query_type)
                     if (
                         not life_cycle_attr
-                        or life_cycle_attr.timestamp.__root__
-                        < access_details.timestamp.__root__
+                        or life_cycle_attr.timestamp.root
+                        < access_details.timestamp.root
                     ):
                         setattr(life_cycle, query_type, access_details)
 

metadata/ingestion/connections/builders.py

@@ -33,6 +33,9 @@ from metadata.generated.schema.entity.services.connections.database.common.iamAu
 from metadata.ingestion.connections.headers import inject_query_header_by_conn
 from metadata.ingestion.connections.secrets import connection_with_options_secrets
 from metadata.utils.constants import BUILDER_PASSWORD_ATTR
+from metadata.utils.logger import cli_logger
+
+logger = cli_logger()
 
 
 @connection_with_options_secrets
@@ -45,8 +48,8 @@ def get_connection_args_common(connection) -> Dict[str, Any]:
     """
 
     return (
-        connection.connectionArguments.__root__
-        if connection.connectionArguments and connection.connectionArguments.__root__
+        connection.connectionArguments.root
+        if connection.connectionArguments and connection.connectionArguments.root
         else {}
     )
 
@@ -90,8 +93,8 @@ def get_connection_options_dict(connection) -> Optional[Dict[str, Any]]:
     dictionary if exists
     """
     return (
-        connection.connectionOptions.__root__
-        if connection.connectionOptions and connection.connectionOptions.__root__
+        connection.connectionOptions.root
+        if connection.connectionOptions and connection.connectionOptions.root
         else None
     )
 
@@ -101,12 +104,12 @@ def init_empty_connection_arguments() -> ConnectionArguments:
     Initialize a ConnectionArguments model with an empty dictionary.
     This helps set keys without further validations.
 
-    Running `ConnectionArguments()` returns `ConnectionArguments(__root__=None)`.
+    Running `ConnectionArguments()` returns `ConnectionArguments(root=None)`.
 
-    Instead, we want `ConnectionArguments(__root__={}})` so that
-    we can pass new keys easily as `connectionArguments.__root__["key"] = "value"`
+    Instead, we want `ConnectionArguments(root={}})` so that
+    we can pass new keys easily as `connectionArguments.root["key"] = "value"`
     """
-    return ConnectionArguments(__root__={})
+    return ConnectionArguments(root={})
 
 
 def init_empty_connection_options() -> ConnectionOptions:
@@ -114,12 +117,12 @@ def init_empty_connection_options() -> ConnectionOptions:
     Initialize a ConnectionOptions model with an empty dictionary.
     This helps set keys without further validations.
 
-    Running `ConnectionOptions()` returns `ConnectionOptions(__root__=None)`.
+    Running `ConnectionOptions()` returns `ConnectionOptions(root=None)`.
 
-    Instead, we want `ConnectionOptions(__root__={}})` so that
-    we can pass new keys easily as `ConnectionOptions.__root__["key"] = "value"`
+    Instead, we want `ConnectionOptions(root={}})` so that
+    we can pass new keys easily as `ConnectionOptions.root["key"] = "value"`
     """
-    return ConnectionOptions(__root__={})
+    return ConnectionOptions(root={})
 
 
 def _add_password(url: str, connection) -> str:
@@ -152,6 +155,9 @@ def _add_password(url: str, connection) -> str:
                         Region=connection.authType.awsConfig.awsRegion,
                     )
                 )
+    if not password:
+        logger.warning("No password has been provided in connection")
+        password = SecretStr("")
     url += f":{quote_plus(password.get_secret_value())}"
     return url
 

metadata/ingestion/connections/secrets.py

@@ -17,25 +17,27 @@ from functools import wraps
 from metadata.ingestion.models.custom_pydantic import CustomSecretStr
 
 
+# Annotated CustomSecretStr does not like the get_secret_value()
+# pylint: disable=no-member
 def update_connection_opts_args(connection):
     if (
         hasattr(connection, "connectionOptions")
         and connection.connectionOptions
-        and connection.connectionOptions.__root__
+        and connection.connectionOptions.root
     ):
-        for key, value in connection.connectionOptions.__root__.items():
+        for key, value in connection.connectionOptions.root.items():
             if isinstance(value, str):
-                connection.connectionOptions.__root__[key] = CustomSecretStr(
+                connection.connectionOptions.root[key] = CustomSecretStr(
                     value
                 ).get_secret_value()
     if (
         hasattr(connection, "connectionArguments")
         and connection.connectionArguments
-        and connection.connectionArguments.__root__
+        and connection.connectionArguments.root
     ):
-        for key, value in connection.connectionArguments.__root__.items():
+        for key, value in connection.connectionArguments.root.items():
             if isinstance(value, str):
-                connection.connectionArguments.__root__[key] = CustomSecretStr(
+                connection.connectionArguments.root[key] = CustomSecretStr(
                     value
                 ).get_secret_value()
 

metadata/ingestion/connections/test_connections.py

@@ -36,7 +36,9 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
     TestConnectionResult,
     TestConnectionStepResult,
 )
+from metadata.generated.schema.type.basic import Timestamp
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.connections import kill_active_connections
 from metadata.profiler.orm.functions.conn_test import ConnTestFn
 from metadata.utils.logger import cli_logger
 from metadata.utils.timeout import timeout
@@ -146,12 +148,16 @@ def _test_connection_steps_automation_workflow(
                     # break the workflow if the step is a short circuit step
                     break
 
-            test_connection_result.lastUpdatedAt = datetime.now().timestamp()
+            test_connection_result.lastUpdatedAt = Timestamp(
+                int(datetime.now().timestamp() * 1000)
+            )
             metadata.patch_automation_workflow_response(
                 automation_workflow, test_connection_result, WorkflowStatus.Running
             )
 
-        test_connection_result.lastUpdatedAt = datetime.now().timestamp()
+        test_connection_result.lastUpdatedAt = Timestamp(
+            int(datetime.now().timestamp() * 1000)
+        )
 
         test_connection_result.status = (
             StatusType.Failed
@@ -324,6 +330,8 @@ def test_connection_db_common(
         timeout_seconds=timeout_seconds,
     )
 
+    kill_active_connections(engine)
+
 
 def test_connection_db_schema_sources(
     metadata: OpenMetadata,
@@ -387,6 +395,8 @@ def test_connection_db_schema_sources(
         automation_workflow=automation_workflow,
     )
 
+    kill_active_connections(engine)
+
 
 def test_query(engine: Engine, statement: str):
     """

metadata/ingestion/lineage/models.py

@@ -14,7 +14,7 @@ Models related to lineage parsing
 from enum import Enum
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaType,
@@ -67,6 +67,9 @@ from metadata.generated.schema.entity.services.connections.database.sqliteConnec
 from metadata.generated.schema.entity.services.connections.database.teradataConnection import (
     TeradataType,
 )
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    DatabricksType as UnityCatalogType,
+)
 from metadata.utils.singleton import Singleton
 
 
@@ -103,6 +106,7 @@ MAP_CONNECTION_TYPE_DIALECT: Dict[str, Dialect] = {
     str(BigqueryType.BigQuery.value): Dialect.BIGQUERY,
     str(ClickhouseType.Clickhouse.value): Dialect.CLICKHOUSE,
     str(DatabricksType.Databricks.value): Dialect.DATABRICKS,
+    str(UnityCatalogType.UnityCatalog.value): Dialect.DATABRICKS,
     str(Db2Type.Db2.value): Dialect.DB2,
     str(HiveType.Hive.value): Dialect.HIVE,
     str(ImpalaType.Impala.value): Dialect.IMPALA,
@@ -144,11 +148,12 @@ class QueryParsingError(BaseModel):
         error (str): The error message of the failed query.
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     query: str = Field(..., description="query text of the failed query")
-    error: Optional[str] = Field(..., description="error message of the failed query")
+    error: Optional[str] = Field(None, description="error message of the failed query")
 
 
 class QueryParsingFailures(metaclass=Singleton):