openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/sas/client.py

@@ -35,7 +35,7 @@ class SASClient:
             config.serverHost, config.username, config.password.get_secret_value()
         )
         client_config: ClientConfig = ClientConfig(
-            base_url=config.serverHost,
+            base_url=str(config.serverHost),
             auth_header="Authorization",
             auth_token=self.get_auth_token,
             api_version="",
@@ -77,7 +77,7 @@ class SASClient:
         return response
 
     def get_information_catalog_link(self, instance_id):
-        return f"{self.config.serverHost}/SASInformationCatalog/details/~fs~catalog~fs~instances~fs~{instance_id}"
+        return f"{self.config.serverHost}SASInformationCatalog/details/~fs~catalog~fs~instances~fs~{instance_id}"
 
     def list_assets(self, assets):
         """

metadata/ingestion/source/database/sas/metadata.py

@@ -17,8 +17,8 @@ SAS source to extract metadata
 import copy
 import json
 import re
-import time
 import traceback
+from datetime import datetime, timezone
 from typing import Any, Iterable, Optional, Tuple, Union
 
 from requests.exceptions import HTTPError
@@ -65,6 +65,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import EntityName, Timestamp
 from metadata.generated.schema.type.entityLineage import EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.common import Entity
@@ -101,7 +102,7 @@ class SasSource(
         self.source_config: DatabaseServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
 
         self.sas_client = get_connection(self.service_connection)
         self.connection_obj = self.sas_client
@@ -121,6 +122,8 @@ class SasSource(
         self.databases = None
         self.database_schemas = None
 
+        self.timestamp = Timestamp(int(datetime.now().timestamp() * 1000))
+
     @classmethod
     def create(
         cls,
@@ -129,8 +132,8 @@ class SasSource(
         pipeline_name: Optional[str] = None,
     ):
         logger.info(f"running create {config_dict}")
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SASConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SASConnection = config.serviceConnection.root.config
         if not isinstance(connection, SASConnection):
             raise InvalidSourceException(
                 f"Expected SASConnection, but got {connection}"
@@ -412,8 +415,7 @@ class SasSource(
                         col_profile_dict["valuesCount"]
                         - col_profile_dict["missingCount"]
                     )
-            timestamp = time.time() - 100000
-            col_profile_dict["timestamp"] = timestamp
+            col_profile_dict["timestamp"] = self.timestamp
             col_profile_dict["name"] = parsed_string["name"]
             column_profile = ColumnProfile(**col_profile_dict)
             col_profile_list.append(column_profile)
@@ -466,8 +468,8 @@ class SasSource(
             # if the table entity already exists, we don't need to create it again
             # only update it when either the sourceUrl or analysisTimeStamp changed
             if not table_entity or (
-                table_url != table_entity.sourceUrl.__root__
-                or table_entity.extension.__root__.get("analysisTimeStamp")
+                table_url != table_entity.sourceUrl.root
+                or table_entity.extension.root.get("analysisTimeStamp")
                 != table_extension.get("analysisTimeStamp")
             ):
                 # create the columns of the table
@@ -530,10 +532,10 @@ class SasSource(
                 )
                 # update the description
                 logger.debug(
-                    f"Updating description for {table_entity.id.__root__} with {table_description}"
+                    f"Updating description for {table_entity.id.root} with {table_description}"
                 )
                 self.metadata.client.patch(
-                    path=f"/tables/{table_entity.id.__root__}",
+                    path=f"/tables/{table_entity.id.root}",
                     data=json.dumps(
                         [
                             {
@@ -547,10 +549,10 @@ class SasSource(
 
                 # update the custom properties
                 logger.debug(
-                    f"Updating custom properties for {table_entity.id.__root__} with {extension_attributes}"
+                    f"Updating custom properties for {table_entity.id.root} with {extension_attributes}"
                 )
                 self.metadata.client.patch(
-                    path=f"/tables/{table_entity.id.__root__}",
+                    path=f"/tables/{table_entity.id.root}",
                     data=json.dumps(
                         [
                             {
@@ -570,35 +572,26 @@ class SasSource(
                 ):
                     return
 
-                # update table profile
-                table_profile_dict = {
-                    "timestamp": time.time() - 100000,
-                    "createDateTime": table_entity_instance["creationTimeStamp"],
-                    "rowCount": (
-                        0
-                        if "rowCount" not in table_extension
-                        else table_extension["rowCount"]
-                    ),
-                    "columnCount": (
-                        0
-                        if "columnCount" not in table_extension
-                        else table_extension["columnCount"]
-                    ),
-                    "sizeInByte": (
-                        0
-                        if "dataSize" not in extension_attributes
-                        else table_extension["dataSize"]
-                    ),
-                }
+                raw_create_date: Optional[datetime] = table_entity_instance.get(
+                    "creationTimeStamp"
+                )
+                if raw_create_date:
+                    raw_create_date = raw_create_date.replace(tzinfo=timezone.utc)
 
                 # create Profiles & Data Quality Column
                 table_profile_request = CreateTableProfileRequest(
-                    tableProfile=TableProfile(**table_profile_dict),
+                    tableProfile=TableProfile(
+                        timestamp=self.timestamp,
+                        createDateTime=raw_create_date,
+                        rowCount=int(table_extension.get("rowCount", 0)),
+                        columnCount=int(table_extension.get("columnCount", 0)),
+                        sizeInByte=int(table_extension.get("dataSize", 0)),
+                    ),
                     columnProfile=col_profile_list,
                 )
                 self.metadata.client.put(
-                    path=f"{self.metadata.get_suffix(Table)}/{table_entity.id.__root__}/tableProfile",
-                    data=table_profile_request.json(),
+                    path=f"{self.metadata.get_suffix(Table)}/{table_entity.id.root}/tableProfile",
+                    data=table_profile_request.model_dump_json(),
                 )
 
         except Exception as exc:
@@ -607,7 +600,7 @@ class SasSource(
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception to create table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
         finally:
@@ -693,7 +686,7 @@ class SasSource(
                 left=StackTraceError(
                     name=dashboard_service_name,
                     error=f"Unexpected exception to create dashboard service for [{dashboard_service_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -733,10 +726,8 @@ class SasSource(
         return Either(
             right=AddLineageRequest(
                 edge=EntitiesEdge(
-                    fromEntity=EntityReference(
-                        id=from_entity.id.__root__, type=from_type
-                    ),
-                    toEntity=EntityReference(id=to_entity.id.__root__, type=in_type),
+                    fromEntity=EntityReference(id=from_entity.id.root, type=from_type),
+                    toEntity=EntityReference(id=to_entity.id.root, type=in_type),
                 )
             )
         )
@@ -787,7 +778,7 @@ class SasSource(
                 left=StackTraceError(
                     name=report_name,
                     error=f"Unexpected exception to create report [{report['id']}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -844,7 +835,7 @@ class SasSource(
                 left=StackTraceError(
                     name=data_flow_id,
                     error=f"Unexpected exception to create data flow [{data_flow_id}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -857,7 +848,7 @@ class SasSource(
     ) -> Iterable[Either[CreateDatabaseRequest]]:
         yield Either(
             right=CreateDatabaseRequest(
-                name=database_name,
+                name=EntityName(database_name),
                 service=self.context.get().database_service,
             )
         )

metadata/ingestion/source/database/singlestore/metadata.py

@@ -45,8 +45,8 @@ class SinglestoreSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SingleStoreConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SingleStoreConnection = config.serviceConnection.root.config
         if not isinstance(connection, SingleStoreConnection):
             raise InvalidSourceException(
                 f"Expected SingleStoreConnection, but got {connection}"

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

@@ -54,7 +54,7 @@ logger = ingestion_logger()
 class SnowflakeEngineWrapper(BaseModel):
     service_connection: SnowflakeConnection
     engine: Any
-    database_name: Optional[str]
+    database_name: Optional[str] = None
 
 
 def get_connection_url(connection: SnowflakeConnection) -> str:
@@ -125,10 +125,10 @@ def get_connection(connection: SnowflakeConnection) -> Engine:
             encryption_algorithm=serialization.NoEncryption(),
         )
 
-        connection.connectionArguments.__root__["private_key"] = pkb
+        connection.connectionArguments.root["private_key"] = pkb
 
     if connection.clientSessionKeepAlive:
-        connection.connectionArguments.__root__[
+        connection.connectionArguments.root[
             "client_session_keep_alive"
         ] = connection.clientSessionKeepAlive
 

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

@@ -20,7 +20,7 @@ import sqlparse
 from snowflake.sqlalchemy.custom_types import VARIANT
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect, ischema_names
 from sqlalchemy.engine.reflection import Inspector
-from sqlparse.sql import Function, Identifier
+from sqlparse.sql import Function, Identifier, Token
 
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
@@ -44,7 +44,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, SourceUrl
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    SourceUrl,
+)
 from metadata.ingestion.api.delete import delete_entity_by_name
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
@@ -61,6 +65,9 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 from metadata.ingestion.source.database.incremental_metadata_extraction import (
     IncrementalConfig,
 )
+from metadata.ingestion.source.database.life_cycle_query_mixin import (
+    LifeCycleQueryMixin,
+)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.snowflake.models import (
     STORED_PROC_LANGUAGE_MAP,
@@ -78,6 +85,7 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_SCHEMA_COMMENTS,
     SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES,
     SNOWFLAKE_GET_STORED_PROCEDURES,
+    SNOWFLAKE_LIFE_CYCLE_QUERY,
     SNOWFLAKE_SESSION_TAG_QUERY,
 )
 from metadata.ingestion.source.database.snowflake.utils import (
@@ -140,6 +148,7 @@ SnowflakeDialect._get_schema_foreign_keys = get_schema_foreign_keys
 
 
 class SnowflakeSource(
+    LifeCycleQueryMixin,
     StoredProcedureMixin,
     ExternalTableLineageMixin,
     CommonDbSourceService,
@@ -165,6 +174,7 @@ class SnowflakeSource(
 
         self._account: Optional[str] = None
         self._org_name: Optional[str] = None
+        self.life_cycle_query = SNOWFLAKE_LIFE_CYCLE_QUERY
         self.context.get_global().deleted_tables = []
         self.pipeline_name = pipeline_name
         self.incremental = incremental_configuration
@@ -180,8 +190,8 @@ class SnowflakeSource(
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SnowflakeConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SnowflakeConnection = config.serviceConnection.root.config
         if not isinstance(connection, SnowflakeConnection):
             raise InvalidSourceException(
                 f"Expected SnowflakeConnection, but got {connection}"
@@ -279,7 +289,7 @@ class SnowflakeSource(
             yield row[1]
 
     def get_database_names(self) -> Iterable[str]:
-        configured_db = self.config.serviceConnection.__root__.config.database
+        configured_db = self.config.serviceConnection.root.config.database
         if configured_db:
             self.set_inspector(configured_db)
             self.set_session_query_tag()
@@ -320,6 +330,21 @@ class SnowflakeSource(
                         f"Error trying to connect to database {new_database}: {exc}"
                     )
 
+    def __clean_append(self, token: Token, result_list: List) -> None:
+        """
+        Appends the real name of the given token to the result list if it exists.
+
+        Args:
+            token (Token): The token whose real name is to be appended.
+            result_list (List): The list to which the real name will be appended.
+
+        Returns:
+            None
+        """
+        name = token.get_real_name()
+        if name is not None:
+            result_list.append(name)
+
     def __get_identifier_from_function(self, function_token: Function) -> List:
         identifiers = []
         for token in function_token.get_parameters():
@@ -327,7 +352,7 @@ class SnowflakeSource(
                 # get column names from nested functions
                 identifiers.extend(self.__get_identifier_from_function(token))
             elif isinstance(token, Identifier):
-                identifiers.append(token.get_real_name())
+                self.__clean_append(token, identifiers)
         return identifiers
 
     def parse_column_name_from_expr(self, cluster_key_expr: str) -> Optional[List[str]]:
@@ -341,7 +366,7 @@ class SnowflakeSource(
                 if isinstance(token, Function):
                     result.extend(self.__get_identifier_from_function(token))
                 elif isinstance(token, Identifier):
-                    result.append(token.get_real_name())
+                    self.__clean_append(token, result)
             return result
         except Exception as err:
             logger.debug(traceback.format_exc())
@@ -427,8 +452,10 @@ class SnowflakeSource(
                 row = list(res)
                 fqn_elements = [name for name in row[2:] if name]
                 yield from get_ometa_tag_and_classification(
-                    tag_fqn=fqn._build(  # pylint: disable=protected-access
-                        self.context.get().database_service, *fqn_elements
+                    tag_fqn=FullyQualifiedEntityName(
+                        fqn._build(  # pylint: disable=protected-access
+                            self.context.get().database_service, *fqn_elements
+                        )
                     ),
                     tags=[row[1]],
                     classification_name=row[0],
@@ -617,7 +644,7 @@ class SnowflakeSource(
                 )
             ).all()
             for row in results:
-                stored_procedure = SnowflakeStoredProcedure.parse_obj(dict(row))
+                stored_procedure = SnowflakeStoredProcedure.model_validate(dict(row))
                 if stored_procedure.definition is None:
                     logger.debug(
                         f"Missing ownership permissions on procedure {stored_procedure.name}."
@@ -655,7 +682,7 @@ class SnowflakeSource(
 
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
-                name=EntityName(__root__=stored_procedure.name),
+                name=EntityName(stored_procedure.name),
                 description=stored_procedure.comment,
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
@@ -669,7 +696,7 @@ class SnowflakeSource(
                     schema_name=self.context.get().database_schema,
                 ),
                 sourceUrl=SourceUrl(
-                    __root__=self._get_source_url_root(
+                    self._get_source_url_root(
                         database_name=self.context.get().database,
                         schema_name=self.context.get().database_schema,
                     )

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

@@ -15,7 +15,7 @@ import urllib
 from datetime import datetime
 from typing import List, Optional
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 from requests.utils import quote
 
 from metadata.generated.schema.entity.data.storedProcedure import Language
@@ -35,16 +35,16 @@ class SnowflakeStoredProcedure(BaseModel):
     """Snowflake stored procedure list query results"""
 
     name: str = Field(..., alias="NAME")
-    owner: Optional[str] = Field(..., alias="OWNER")
+    owner: Optional[str] = Field(None, alias="OWNER")
     language: str = Field(..., alias="LANGUAGE")
-    definition: str = Field(None, alias="DEFINITION")
+    definition: Optional[str] = Field(None, alias="DEFINITION")
     signature: Optional[str] = Field(
-        ..., alias="SIGNATURE", description="Used to build the source URL"
+        None, alias="SIGNATURE", description="Used to build the source URL"
     )
-    comment: Optional[str] = Field(..., alias="COMMENT")
+    comment: Optional[str] = Field(None, alias="COMMENT")
 
     # Update the signature to clean it up on read
-    @validator("signature")
+    @field_validator("signature")
     def clean_signature(  # pylint: disable=no-self-argument
         cls, signature
     ) -> Optional[str]:
@@ -82,7 +82,7 @@ class SnowflakeTable(BaseModel):
     """
 
     name: str
-    deleted: Optional[datetime]
+    deleted: Optional[datetime] = None
 
 
 class SnowflakeTableList(BaseModel):

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

@@ -43,8 +43,8 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SnowflakeConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SnowflakeConnection = config.serviceConnection.root.config
         if not isinstance(connection, SnowflakeConnection):
             raise InvalidSourceException(
                 f"Expected SnowflakeConnection, but got {connection}"

metadata/ingestion/source/database/sql_column_handler.py

@@ -273,7 +273,7 @@ class SqlColumnHandlerMixin:
                     col_data_length = 1 if col_data_length is None else col_data_length
                     om_column = Column(
                         name=ColumnName(
-                            __root__=column["name"]
+                            root=column["name"]
                             # Passing whitespace if column name is an empty string
                             # since pydantic doesn't accept empty string
                             if column["name"]
@@ -291,8 +291,9 @@ class SqlColumnHandlerMixin:
                         ordinalPosition=column.get("ordinalPosition"),
                     )
                     if precision:
-                        om_column.precision = precision[0]
-                        om_column.scale = precision[1]
+                        # Precision and scale must be integer values
+                        om_column.precision = int(precision[0])
+                        om_column.scale = int(precision[1])
 
                 else:
                     col_obj = self._process_complex_col_type(

metadata/ingestion/source/database/sqlite/metadata.py

@@ -37,8 +37,8 @@ class SqliteSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection = config.serviceConnection.root.config
         if not isinstance(connection, SQLiteConnection):
             raise InvalidSourceException(
                 f"Expected SQLiteConnection, but got {connection}"

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -18,7 +18,7 @@ from collections import defaultdict
 from datetime import datetime
 from typing import Dict, Iterable, List, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -53,18 +53,17 @@ class QueryByProcedure(BaseModel):
 
     procedure_name: str = Field(None, alias="PROCEDURE_NAME")
     query_type: str = Field(..., alias="QUERY_TYPE")
-    query_database_name: str = Field(None, alias="QUERY_DATABASE_NAME")
-    query_schema_name: str = Field(None, alias="QUERY_SCHEMA_NAME")
+    query_database_name: Optional[str] = Field(None, alias="QUERY_DATABASE_NAME")
+    query_schema_name: Optional[str] = Field(None, alias="QUERY_SCHEMA_NAME")
     procedure_text: str = Field(..., alias="PROCEDURE_TEXT")
     procedure_start_time: datetime = Field(..., alias="PROCEDURE_START_TIME")
     procedure_end_time: datetime = Field(..., alias="PROCEDURE_END_TIME")
-    query_start_time: Optional[datetime] = Field(..., alias="QUERY_START_TIME")
+    query_start_time: Optional[datetime] = Field(None, alias="QUERY_START_TIME")
     query_duration: Optional[float] = Field(None, alias="QUERY_DURATION")
     query_text: str = Field(..., alias="QUERY_TEXT")
     query_user_name: Optional[str] = Field(None, alias="QUERY_USER_NAME")
 
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class StoredProcedureMixin(ABC):
@@ -107,7 +106,7 @@ class StoredProcedureMixin(ABC):
 
         for row in results:
             try:
-                query_by_procedure = QueryByProcedure.parse_obj(dict(row))
+                query_by_procedure = QueryByProcedure.model_validate(dict(row))
                 procedure_name = (
                     query_by_procedure.procedure_name
                     or get_procedure_name_from_call(
@@ -181,11 +180,11 @@ class StoredProcedureMixin(ABC):
 
         yield Either(
             right=CreateQueryRequest(
-                query=SqlQuery(__root__=query_by_procedure.query_text),
+                query=SqlQuery(query_by_procedure.query_text),
                 query_type=query_by_procedure.query_type,
                 duration=query_by_procedure.query_duration,
                 queryDate=Timestamp(
-                    __root__=convert_timestamp_to_milliseconds(
+                    root=convert_timestamp_to_milliseconds(
                         int(query_by_procedure.query_start_time.timestamp())
                     )
                 ),
@@ -210,13 +209,13 @@ class StoredProcedureMixin(ABC):
             queries_dict = self.get_stored_procedure_queries_dict()
             # Then for each procedure, iterate over all its queries
             for procedure_fqn in self.context.get().stored_procedures:
-                procedure = self.metadata.get_by_name(
+                procedure: StoredProcedure = self.metadata.get_by_name(
                     entity=StoredProcedure, fqn=procedure_fqn
                 )
                 if procedure:
                     logger.debug(f"Processing Lineage for [{procedure.name}]")
                     for query_by_procedure in (
-                        queries_dict.get(procedure.name.__root__.lower()) or []
+                        queries_dict.get(procedure.name.root.lower()) or []
                     ):
                         yield from self.yield_procedure_lineage(
                             query_by_procedure=query_by_procedure, procedure=procedure

metadata/ingestion/source/database/teradata/metadata.py

@@ -65,8 +65,8 @@ class TeradataSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection = config.serviceConnection.root.config
         if not isinstance(connection, TeradataConnection):
             raise InvalidSourceException(
                 f"Expected TeradataConnection, but got {connection}"
@@ -83,7 +83,7 @@ class TeradataSource(CommonDbSourceService):
             ).all()
             for row in results:
                 try:
-                    stored_procedure = TeradataStoredProcedure.parse_obj(dict(row))
+                    stored_procedure = TeradataStoredProcedure.model_validate(dict(row))
                     stored_procedure.definition = self.describe_procedure_definition(
                         stored_procedure
                     )
@@ -119,7 +119,7 @@ class TeradataSource(CommonDbSourceService):
 
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
-                name=EntityName(__root__=stored_procedure.procedure_name),
+                name=EntityName(stored_procedure.procedure_name),
                 description=None,
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(

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

@@ -65,13 +65,13 @@ def get_connection_url(connection: TrinoConnection) -> str:
     if isinstance(connection.authType, jwtAuth.JwtAuth):
         if not connection.connectionOptions:
             connection.connectionOptions = init_empty_connection_options()
-        connection.connectionOptions.__root__[
+        connection.connectionOptions.root[
             "access_token"
         ] = connection.authType.jwt.get_secret_value()
     if connection.connectionOptions is not None:
         params = "&".join(
             f"{key}={quote_plus(value)}"
-            for (key, value) in connection.connectionOptions.__root__.items()
+            for (key, value) in connection.connectionOptions.root.items()
             if value
         )
         url = f"{url}?{params}"
@@ -86,7 +86,7 @@ def get_connection_args(connection: TrinoConnection):
         if not connection.connectionArguments:
             connection.connectionArguments = init_empty_connection_arguments()
 
-        connection.connectionArguments.__root__["http_session"] = session
+        connection.connectionArguments.root["http_session"] = session
 
     return get_connection_args_common(connection)
 
@@ -99,9 +99,7 @@ def get_connection(connection: TrinoConnection) -> Engine:
         connection.connectionArguments = (
             connection.connectionArguments or init_empty_connection_arguments()
         )
-        connection.connectionArguments.__root__["verify"] = {
-            "verify": connection.verify
-        }
+        connection.connectionArguments.root["verify"] = {"verify": connection.verify}
     if hasattr(connection.authType, "azureConfig"):
         azure_client = AzureClient(connection.authType.azureConfig).create_client()
         if not connection.authType.azureConfig.scopes:
@@ -113,7 +111,7 @@ def get_connection(connection: TrinoConnection) -> Engine:
         )
         if not connection.connectionOptions:
             connection.connectionOptions = init_empty_connection_options()
-        connection.connectionOptions.__root__["access_token"] = access_token_obj.token
+        connection.connectionOptions.root["access_token"] = access_token_obj.token
     return create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,

metadata/ingestion/source/database/trino/metadata.py

@@ -186,8 +186,8 @@ class TrinoSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config = WorkflowSource.parse_obj(config_dict)
-        connection: TrinoConnection = config.serviceConnection.__root__.config
+        config = WorkflowSource.model_validate(config_dict)
+        connection: TrinoConnection = config.serviceConnection.root.config
         if not isinstance(connection, TrinoConnection):
             raise InvalidSourceException(
                 f"Expected TrinoConnection, but got {connection}"

metadata/ingestion/source/database/trino/query_parser.py

@@ -37,8 +37,8 @@ class TrinoQueryParserSource(QueryParserSource, ABC):
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
         """Create class instance"""
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: TrinoConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: TrinoConnection = config.serviceConnection.root.config
         if not isinstance(connection, TrinoConnection):
             raise InvalidSourceException(
                 f"Expected TrinoConnection, but got {connection}"