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

metadata/ingestion/source/metadata/atlas/metadata.py

@@ -75,7 +75,7 @@ class AtlasSource(Source):
         super().__init__()
         self.config = config
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
 
         self.atlas_client = get_connection(self.service_connection)
         self.connection_obj = self.atlas_client
@@ -96,8 +96,8 @@ class AtlasSource(Source):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: AtlasConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: AtlasConnection = config.serviceConnection.root.config
         if not isinstance(connection, AtlasConnection):
             raise InvalidSourceException(
                 f"Expected AtlasConnection, but got {connection}"
@@ -209,7 +209,7 @@ class AtlasSource(Source):
                     database_fqn = fqn.build(
                         self.metadata,
                         entity_type=Database,
-                        service_name=self.service.name.__root__,
+                        service_name=self.service.name.root,
                         database_name=database_name,
                     )
                     database_object = self.metadata.get_by_name(
@@ -226,7 +226,7 @@ class AtlasSource(Source):
                     database_schema_fqn = fqn.build(
                         self.metadata,
                         entity_type=DatabaseSchema,
-                        service_name=self.service.name.__root__,
+                        service_name=self.service.name.root,
                         database_name=database_name,
                         schema_name=db_entity["displayText"],
                     )
@@ -252,7 +252,7 @@ class AtlasSource(Source):
                     table_fqn = fqn.build(
                         metadata=self.metadata,
                         entity_type=Table,
-                        service_name=self.service.name.__root__,
+                        service_name=self.service.name.root,
                         database_name=database_name,
                         schema_name=db_entity["displayText"],
                         table_name=tbl_attrs["name"],
@@ -335,7 +335,6 @@ class AtlasSource(Source):
                 col_guid = col["guid"]
                 col_ref_entity = referred_entities[col_guid]
                 column = col_ref_entity["attributes"]
-                col_data_length = "1"
                 om_column = Column(
                     name=column["name"],
                     description=column.get("comment", None),
@@ -343,7 +342,7 @@ class AtlasSource(Source):
                         column["dataType"].upper()
                     ),
                     dataTypeDisplay=column["dataType"],
-                    dataLength=col_data_length,
+                    dataLength=1,
                     ordinalPosition=ordinal_pos,
                 )
                 om_cols.append(om_column)
@@ -377,7 +376,7 @@ class AtlasSource(Source):
                 from_fqn = fqn.build(
                     self.metadata,
                     entity_type=Table,
-                    service_name=self.service.name.__root__,
+                    service_name=self.service.name.root,
                     database_name=get_database_name_for_lineage(
                         self.service, db_entity["displayText"]
                     ),
@@ -406,7 +405,7 @@ class AtlasSource(Source):
                         to_fqn = fqn.build(
                             self.metadata,
                             entity_type=Table,
-                            service_name=self.service.name.__root__,
+                            service_name=self.service.name.root,
                             database_name=get_database_name_for_lineage(
                                 self.service, db_entity["displayText"]
                             ),
@@ -469,11 +468,11 @@ class AtlasSource(Source):
         if entity_type == "table":
             table: Table = self.metadata.get_by_name(entity=Table, fqn=to_fqn)
             if table:
-                return EntityReference(id=table.id.__root__, type="table")
+                return EntityReference(id=table.id.root, type="table")
         if entity_type == "pipeline":
             pipeline: Pipeline = self.metadata.get_by_name(entity=Pipeline, fqn=to_fqn)
             if pipeline:
-                return EntityReference(id=pipeline.id.__root__, type="pipeline")
+                return EntityReference(id=pipeline.id.root, type="pipeline")
         return None
 
     def test_connection(self) -> None:

metadata/ingestion/source/mlmodel/mlflow/metadata.py

@@ -35,6 +35,12 @@ 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,
+    FullyQualifiedEntityName,
+    Markdown,
+    SourceUrl,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -58,8 +64,8 @@ class MlflowSource(MlModelServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: MlflowConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: MlflowConnection = config.serviceConnection.root.config
         if not isinstance(connection, MlflowConnection):
             raise InvalidSourceException(
                 f"Expected MlFlowConnection, but got {connection}"
@@ -120,16 +126,16 @@ class MlflowSource(MlModelServiceSource):
         )
 
         mlmodel_request = CreateMlModelRequest(
-            name=model.name,
-            description=model.description,
+            name=EntityName(model.name),
+            description=Markdown(model.description) if model.description else None,
             algorithm=self._get_algorithm(),  # Setting this to a constant
             mlHyperParameters=self._get_hyper_params(run.data),
             mlFeatures=self._get_ml_features(
                 run.data, latest_version.run_id, model.name
             ),
             mlStore=self._get_ml_store(latest_version),
-            service=self.context.get().mlmodel_service,
-            sourceUrl=source_url,
+            service=FullyQualifiedEntityName(self.context.get().mlmodel_service),
+            sourceUrl=SourceUrl(source_url),
         )
         yield Either(right=mlmodel_request)
         self.register_record(mlmodel_request=mlmodel_request)

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -14,6 +14,9 @@ Base class for ingesting mlmodel services
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Set
 
+from pydantic import Field
+from typing_extensions import Annotated
+
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
 from metadata.generated.schema.entity.data.mlmodel import (
     MlFeature,
@@ -59,7 +62,9 @@ class MlModelServiceTopology(ServiceTopology):
     data that has been produced by any parent node.
     """
 
-    root = TopologyNode(
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
         producer="get_services",
         stages=[
             NodeStage(
@@ -74,7 +79,9 @@ class MlModelServiceTopology(ServiceTopology):
         children=["mlmodel"],
         post_process=["mark_mlmodels_as_deleted"],
     )
-    mlmodel = TopologyNode(
+    mlmodel: Annotated[
+        TopologyNode, Field(description="ML Model Processing Node")
+    ] = TopologyNode(
         producer="get_mlmodels",
         stages=[
             NodeStage(
@@ -97,7 +104,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
     source_config: MlModelServiceMetadataPipeline
     config: WorkflowSource
     # Big union of types we want to fetch dynamically
-    service_connection: MlModelConnection.__fields__["config"].type_
+    service_connection: MlModelConnection.model_fields["config"].annotation
 
     topology = MlModelServiceTopology()
     context = TopologyContextManager(topology)
@@ -111,7 +118,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.source_config: MlModelServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
@@ -190,8 +197,8 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
         mlmodel_fqn = fqn.build(
             self.metadata,
             entity_type=MlModel,
-            service_name=mlmodel_request.service.__root__,
-            mlmodel_name=mlmodel_request.name.__root__,
+            service_name=mlmodel_request.service.root,
+            mlmodel_name=mlmodel_request.name.root,
         )
 
         self.mlmodel_source_state.add(mlmodel_fqn)

metadata/ingestion/source/mlmodel/sagemaker/metadata.py

@@ -13,7 +13,7 @@
 import traceback
 from typing import Iterable, List, Optional
 
-from pydantic import BaseModel, Extra, Field, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
 from metadata.generated.schema.entity.data.mlmodel import (
@@ -30,7 +30,18 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.tagLabel import TagLabel
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
+from metadata.generated.schema.type.tagLabel import (
+    LabelType,
+    State,
+    TagFQN,
+    TagLabel,
+    TagSource,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -42,8 +53,9 @@ logger = ingestion_logger()
 
 
 class SageMakerModel(BaseModel):
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     name: str = Field(..., description="Model name", title="Model Name")
     arn: str = Field(..., description="Model ARN in AWS account", title="Model ARN")
@@ -70,8 +82,8 @@ class SagemakerSource(MlModelServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SageMakerConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SageMakerConnection = config.serviceConnection.root.config
         if not isinstance(connection, SageMakerConnection):
             raise InvalidSourceException(
                 f"Expected SageMakerConnection, but got {connection}"
@@ -86,14 +98,14 @@ class SagemakerSource(MlModelServiceSource):
         try:
             while has_more_models:
                 response = self.sagemaker.list_models(**args)
-                models.append(response["Models"])
+                models.extend(response["Models"])
                 has_more_models = response.get("NextToken")
                 args["NextToken"] = response.get("NextToken")
         except Exception as err:
             logger.debug(traceback.format_exc())
             logger.error(f"Failed to fetch models list - {err}")
 
-        for model in response["Models"]:
+        for model in models:
             try:
                 if filter_by_mlmodel(
                     self.source_config.mlModelFilterPattern,
@@ -135,10 +147,10 @@ class SagemakerSource(MlModelServiceSource):
         """
         try:
             mlmodel_request = CreateMlModelRequest(
-                name=model.name,
+                name=EntityName(model.name),
                 algorithm=self._get_algorithm(),  # Setting this to a constant
                 mlStore=self._get_ml_store(model.name),
-                service=self.context.get().mlmodel_service,
+                service=FullyQualifiedEntityName(self.context.get().mlmodel_service),
             )
             yield Either(right=mlmodel_request)
             self.register_record(mlmodel_request=mlmodel_request)
@@ -178,17 +190,18 @@ class SagemakerSource(MlModelServiceSource):
 
     def _get_tags(self, model_arn: str) -> Optional[List[TagLabel]]:
         try:
-            tags = self.sagemaker.list_tags(ResourceArn=model_arn)["Tags"]
-            return [
-                TagLabel(
-                    tagFQN=tag["Key"],
-                    description=tag["Value"],
-                    source="Classification",
-                    labelType="Propagated",
-                    state="Confirmed",
-                )
-                for tag in tags
-            ]
+            tags = self.sagemaker.list_tags(ResourceArn=model_arn).get("Tags")
+            if tags:
+                return [
+                    TagLabel(
+                        tagFQN=TagFQN(tag["Key"]),
+                        description=Markdown(tag["Value"]),
+                        source=TagSource.Classification,
+                        labelType=LabelType.Automated,
+                        state=State.Confirmed,
+                    )
+                    for tag in tags
+                ]
         except ValidationError as err:
             logger.debug(traceback.format_exc())
             logger.warning(

metadata/ingestion/source/pipeline/airbyte/client.py

@@ -30,7 +30,7 @@ class AirbyteClient:
     def __init__(self, config: AirbyteConnection):
         self.config = config
         client_config: ClientConfig = ClientConfig(
-            base_url=self.config.hostPort,
+            base_url=str(self.config.hostPort),
             api_version="api/v1",
             auth_header=AUTHORIZATION_HEADER,
             auth_token=lambda: (NO_ACCESS_TOKEN, 0),

metadata/ingestion/source/pipeline/airbyte/metadata.py

@@ -33,6 +33,12 @@ from metadata.generated.schema.entity.services.databaseService import DatabaseSe
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    SourceUrl,
+    Timestamp,
+)
 from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
@@ -78,8 +84,8 @@ class AirbyteSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: AirbyteConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: AirbyteConnection = config.serviceConnection.root.config
         if not isinstance(connection, AirbyteConnection):
             raise InvalidSourceException(
                 f"Expected AirbyteConnection, but got {connection}"
@@ -94,7 +100,7 @@ class AirbyteSource(PipelineServiceSource):
             Task(
                 name=connection["connectionId"],
                 displayName=connection["name"],
-                sourceUrl=f"{connection_url}/status",
+                sourceUrl=SourceUrl(f"{connection_url}/status"),
             )
         ]
 
@@ -112,13 +118,13 @@ class AirbyteSource(PipelineServiceSource):
             f"/connections/{pipeline_details.connection.get('connectionId')}"
         )
         pipeline_request = CreatePipelineRequest(
-            name=pipeline_details.connection.get("connectionId"),
+            name=EntityName(pipeline_details.connection.get("connectionId")),
             displayName=pipeline_details.connection.get("name"),
-            sourceUrl=connection_url,
+            sourceUrl=SourceUrl(connection_url),
             tasks=self.get_connections_jobs(
                 pipeline_details.connection, connection_url
             ),
-            service=self.context.get().pipeline_service,
+            service=FullyQualifiedEntityName(self.context.get().pipeline_service),
         )
         yield Either(right=pipeline_request)
         self.register_record(pipeline_request=pipeline_request)
@@ -132,7 +138,7 @@ class AirbyteSource(PipelineServiceSource):
 
         # Airbyte does not offer specific attempt link, just at pipeline level
         log_link = (
-            f"{self.service_connection.hostPort}/workspaces/{pipeline_details.workspace.get('workspaceId')}"
+            f"{self.service_connection.hostPort}workspaces/{pipeline_details.workspace.get('workspaceId')}"
             f"/connections/{pipeline_details.connection.get('connectionId')}/status"
         )
 
@@ -168,7 +174,7 @@ class AirbyteSource(PipelineServiceSource):
                         attempt["status"].lower(), StatusType.Pending
                     ).value,
                     taskStatus=task_status,
-                    timestamp=created_at,
+                    timestamp=Timestamp(created_at),
                 )
                 pipeline_fqn = fqn.build(
                     metadata=self.metadata,
@@ -245,9 +251,7 @@ class AirbyteSource(PipelineServiceSource):
             )
 
             lineage_details = LineageDetails(
-                pipeline=EntityReference(
-                    id=pipeline_entity.id.__root__, type="pipeline"
-                ),
+                pipeline=EntityReference(id=pipeline_entity.id.root, type="pipeline"),
                 source=LineageSource.PipelineLineage,
             )
 

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

@@ -73,7 +73,7 @@ from functools import singledispatch
 from typing import Any, DefaultDict, Dict, List, Optional, Type
 
 import attr
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.models import T
@@ -129,12 +129,11 @@ class XLets(BaseModel):
     Group inlets and outlets from all tasks in a DAG
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     inlets: List[OMEntity]
     outlets: List[OMEntity]
 
-    class Config:
-        arbitrary_types_allowed = True
-
 
 def concat_dict_values(
     dict_1: DefaultDict[str, List[Any]], dict_2: Optional[Dict[str, List[Any]]]

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

@@ -43,9 +43,17 @@ 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,
+    FullyQualifiedEntityName,
+    Markdown,
+    SourceUrl,
+    Timestamp,
+)
 from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.connections.session import create_and_bind_session
@@ -116,8 +124,8 @@ class AirflowSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ) -> "AirflowSource":
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: AirflowConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: AirflowConnection = config.serviceConnection.root.config
         if not isinstance(connection, AirflowConnection):
             raise InvalidSourceException(
                 f"Expected AirflowConnection, but got {connection}"
@@ -149,7 +157,7 @@ class AirflowSource(PipelineServiceSource):
             )
             .filter(DagRun.dag_id == dag_id)
             .order_by(DagRun.execution_date.desc())
-            .limit(self.config.serviceConnection.__root__.config.numberOfStatus)
+            .limit(self.config.serviceConnection.root.config.numberOfStatus)
             .all()
         )
 
@@ -252,12 +260,13 @@ class AirflowSource(PipelineServiceSource):
                         if task.task_id in self.context.get().task_names
                     ]
 
+                    timestamp = datetime_to_ts(dag_run.execution_date)
                     pipeline_status = PipelineStatus(
                         taskStatus=task_statuses,
                         executionStatus=STATUS_MAP.get(
                             dag_run.state, StatusType.Pending.value
                         ),
-                        timestamp=datetime_to_ts(dag_run.execution_date),
+                        timestamp=Timestamp(timestamp) if timestamp else None,
                     )
                     pipeline_fqn = fqn.build(
                         metadata=self.metadata,
@@ -315,7 +324,7 @@ class AirflowSource(PipelineServiceSource):
                 dag = AirflowDagDetails(
                     dag_id=serialized_dag[0],
                     fileloc=serialized_dag[2],
-                    data=AirflowDag.parse_obj(serialized_dag[1]),
+                    data=AirflowDag.model_validate(serialized_dag[1]),
                     max_active_runs=data.get("max_active_runs", None),
                     description=data.get("_description", None),
                     start_date=data.get("start_date", None),
@@ -376,7 +385,7 @@ class AirflowSource(PipelineServiceSource):
             Task(
                 name=task.task_id,
                 description=task.doc_md,
-                sourceUrl=(
+                sourceUrl=SourceUrl(
                     f"{clean_uri(host_port)}/taskinstance/list/"
                     f"?flt1_dag_id_equals={dag.dag_id}&_flt_3_task_id={task.task_id}"
                 ),
@@ -386,12 +395,12 @@ class AirflowSource(PipelineServiceSource):
                 startDate=task.start_date.isoformat() if task.start_date else None,
                 endDate=task.end_date.isoformat() if task.end_date else None,
                 taskType=task.task_type,
-                owner=self.get_owner(task.owner),
+                owners=self.get_owner(task.owner),
             )
             for task in cast(Iterable[BaseOperator], dag.tasks)
         ]
 
-    def get_owner(self, owner) -> Optional[EntityReference]:
+    def get_owner(self, owner) -> Optional[EntityReferenceList]:
         """
         Fetching users by name via ES to keep things as fast as possible.
 
@@ -403,7 +412,7 @@ class AirflowSource(PipelineServiceSource):
         until the next run.
         """
         try:
-            return self.metadata.get_reference_by_name(name=owner)
+            return self.metadata.get_reference_by_name(name=owner, is_owner=True)
         except Exception as exc:
             logger.warning(f"Error while getting details of user {owner} - {exc}")
         return None
@@ -422,9 +431,11 @@ class AirflowSource(PipelineServiceSource):
             source_url = f"{clean_uri(self.service_connection.hostPort)}/dags/{pipeline_details.dag_id}/grid"
 
             pipeline_request = CreatePipelineRequest(
-                name=pipeline_details.dag_id,
-                description=pipeline_details.description,
-                sourceUrl=source_url,
+                name=EntityName(pipeline_details.dag_id),
+                description=Markdown(pipeline_details.description)
+                if pipeline_details.description
+                else None,
+                sourceUrl=SourceUrl(source_url),
                 concurrency=pipeline_details.max_active_runs,
                 pipelineLocation=pipeline_details.fileloc,
                 startDate=pipeline_details.start_date.isoformat()
@@ -433,8 +444,8 @@ class AirflowSource(PipelineServiceSource):
                 tasks=self.get_tasks_from_dag(
                     pipeline_details, self.service_connection.hostPort
                 ),
-                service=self.context.get().pipeline_service,
-                owner=self.get_owner(pipeline_details.owner),
+                service=FullyQualifiedEntityName(self.context.get().pipeline_service),
+                owners=self.get_owner(pipeline_details.owner),
                 scheduleInterval=pipeline_details.schedule_interval,
             )
             yield Either(right=pipeline_request)
@@ -497,7 +508,7 @@ class AirflowSource(PipelineServiceSource):
 
         lineage_details = LineageDetails(
             pipeline=EntityReference(
-                id=pipeline_entity.id.__root__,
+                id=pipeline_entity.id.root,
                 type=ENTITY_REFERENCE_TYPE_MAP[Pipeline.__name__],
             ),
             source=LineageSource.PipelineLineage,
@@ -538,12 +549,12 @@ class AirflowSource(PipelineServiceSource):
                         else:
                             logger.warning(
                                 f"Could not find [{to_xlet.entity.__name__}] [{to_xlet.fqn}] from "
-                                f"[{pipeline_entity.fullyQualifiedName.__root__}] outlets"
+                                f"[{pipeline_entity.fullyQualifiedName.root}] outlets"
                             )
                 else:
                     logger.warning(
                         f"Could not find [{from_xlet.entity.__name__}] [{from_xlet.fqn}] from "
-                        f"[{pipeline_entity.fullyQualifiedName.__root__}] inlets"
+                        f"[{pipeline_entity.fullyQualifiedName.root}] inlets"
                     )
 
     def close(self):

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

@@ -16,7 +16,7 @@ Tableau Source Model module
 from datetime import datetime
 from typing import Any, List, Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class AirflowBaseModel(BaseModel):
@@ -24,51 +24,48 @@ class AirflowBaseModel(BaseModel):
     Tableau basic configurations
     """
 
-    class Config:
-        extra = Extra.allow
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
 
     dag_id: str
 
 
 class AirflowTask(BaseModel):
-    pool: Optional[str]
-    doc_md: Optional[str]
-    inlets: Optional[List[Any]] = Field(alias="_inlets")
+    pool: Optional[str] = None
+    doc_md: Optional[str] = None
+    inlets: Optional[List[Any]] = Field(None, alias="_inlets")
     task_id: str
-    outlets: Optional[List[Any]] = Field(alias="_outlets")
-    task_type: Optional[Any] = Field(alias="_task_type")
-    downstream_task_ids: Optional[List[str]]
-    start_date: Optional[datetime]
-    end_date: Optional[datetime]
-    owner: Optional[str]
+    outlets: Optional[List[Any]] = Field(None, alias="_outlets")
+    task_type: Optional[Any] = Field(None, alias="_task_type")
+    downstream_task_ids: Optional[List[str]] = None
+    start_date: Optional[datetime] = None
+    end_date: Optional[datetime] = None
+    owner: Optional[str] = None
 
     # Allow picking up data from key `inlets` and `_inlets`
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class TaskList(BaseModel):
-    __root__: List[AirflowTask]
+    root: List[AirflowTask]
 
 
 class Dag(BaseModel):
     fileloc: str
-    tags: Optional[List[str]]
-    start_date: Optional[float]
+    tags: Optional[List[str]] = None
+    start_date: Optional[float] = None
     _processor_dags_folder: str
 
 
 class AirflowDag(BaseModel):
-    dag: Optional[Dag]
+    dag: Optional[Dag] = None
 
 
 class AirflowDagDetails(AirflowBaseModel):
     fileloc: str
     data: AirflowDag
-    max_active_runs: Optional[int]
-    description: Optional[str]
-    start_date: Optional[datetime]
+    max_active_runs: Optional[int] = None
+    description: Optional[str] = None
+    start_date: Optional[datetime] = None
     tasks: List[AirflowTask]
-    owner: Optional[str]
-    schedule_interval: Optional[str]
+    owner: Optional[str] = None
+    schedule_interval: Optional[str] = None