openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

metadata/ingestion/source/pipeline/gluepipeline/metadata.py

@@ -34,6 +34,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,
+    SourceUrl,
+    Timestamp,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
@@ -75,8 +81,8 @@ class GluepipelineSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: GluePipelineConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: GluePipelineConnection = config.serviceConnection.root.config
         if not isinstance(connection, GluePipelineConnection):
             raise InvalidSourceException(
                 f"Expected GlueConnection, but got {connection}"
@@ -95,17 +101,17 @@ class GluepipelineSource(PipelineServiceSource):
         self, pipeline_details: Any
     ) -> Iterable[Either[CreatePipelineRequest]]:
         """Method to Get Pipeline Entity"""
-        source_url = (
+        source_url = SourceUrl(
             f"https://{self.service_connection.awsConfig.awsRegion}.console.aws.amazon.com/glue/home?"
             f"region={self.service_connection.awsConfig.awsRegion}#/v2/etl-configuration/"
             f"workflows/view/{pipeline_details[NAME]}"
         )
         self.job_name_list = set()
         pipeline_request = CreatePipelineRequest(
-            name=pipeline_details[NAME],
+            name=EntityName(pipeline_details[NAME]),
             displayName=pipeline_details[NAME],
             tasks=self.get_tasks(pipeline_details),
-            service=self.context.get().pipeline_service,
+            service=FullyQualifiedEntityName(self.context.get().pipeline_service),
             sourceUrl=source_url,
         )
         yield Either(right=pipeline_request)
@@ -154,18 +160,24 @@ class GluepipelineSource(PipelineServiceSource):
                             executionStatus=STATUS_MAP.get(
                                 attempt["JobRunState"].lower(), StatusType.Pending
                             ).value,
-                            startTime=convert_timestamp_to_milliseconds(
-                                attempt["StartedOn"].timestamp()
+                            startTime=Timestamp(
+                                convert_timestamp_to_milliseconds(
+                                    attempt["StartedOn"].timestamp()
+                                )
                             ),
-                            endTime=convert_timestamp_to_milliseconds(
-                                attempt["CompletedOn"].timestamp()
+                            endTime=Timestamp(
+                                convert_timestamp_to_milliseconds(
+                                    attempt["CompletedOn"].timestamp()
+                                )
                             ),
                         )
                     )
                     pipeline_status = PipelineStatus(
                         taskStatus=task_status,
-                        timestamp=convert_timestamp_to_milliseconds(
-                            attempt["StartedOn"].timestamp()
+                        timestamp=Timestamp(
+                            convert_timestamp_to_milliseconds(
+                                attempt["StartedOn"].timestamp()
+                            )
                         ),
                         executionStatus=STATUS_MAP.get(
                             attempt["JobRunState"].lower(), StatusType.Pending

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

@@ -36,6 +36,7 @@ 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, 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
@@ -72,8 +73,8 @@ class KafkaconnectSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: KafkaConnectConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: KafkaConnectConnection = config.serviceConnection.root.config
         if not isinstance(connection, KafkaConnectConnection):
             raise InvalidSourceException(
                 f"Expected KafkaConnectConnection, but got {connection}"
@@ -87,10 +88,10 @@ class KafkaconnectSource(PipelineServiceSource):
         Method to Get Pipeline Entity
         """
         try:
-            connection_url = f"{clean_uri(self.service_connection.hostPort)}"
+            connection_url = SourceUrl(f"{clean_uri(self.service_connection.hostPort)}")
 
             pipeline_request = CreatePipelineRequest(
-                name=pipeline_details.name,
+                name=EntityName(pipeline_details.name),
                 sourceUrl=connection_url,
                 tasks=[
                     Task(
@@ -193,16 +194,13 @@ class KafkaconnectSource(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,
             )
 
             dataset_entity = self.get_dataset_entity(pipeline_details=pipeline_details)
 
             for topic in pipeline_details.topics or []:
-
                 topic_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Topic,
@@ -292,7 +290,7 @@ class KafkaconnectSource(PipelineServiceSource):
                     pipeline_details.status, StatusType.Pending
                 ),
                 taskStatus=task_status,
-                timestamp=datetime_to_ts(datetime.now())
+                timestamp=Timestamp(datetime_to_ts(datetime.now()))
                 # Kafka connect doesn't provide any details with exec time
             )
 

metadata/ingestion/source/pipeline/nifi/metadata.py

@@ -25,6 +25,11 @@ from metadata.generated.schema.entity.services.connections.pipeline.nifiConnecti
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    SourceUrl,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
@@ -46,7 +51,7 @@ class NifiProcessor(BaseModel):
     """
 
     id_: str
-    name: Optional[str]
+    name: Optional[str] = None
     type_: str
     uri: str
 
@@ -68,7 +73,7 @@ class NifiPipelineDetails(BaseModel):
     """
 
     id_: str
-    name: Optional[str]
+    name: Optional[str] = None
     uri: str
     processors: List[NifiProcessor]
     connections: List[NifiProcessorConnections]
@@ -84,8 +89,8 @@ class NifiSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: NifiConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: NifiConnection = config.serviceConnection.root.config
         if not isinstance(connection, NifiConnection):
             raise InvalidSourceException(
                 f"Expected NifiConnection, but got {connection}"
@@ -115,7 +120,9 @@ class NifiSource(PipelineServiceSource):
                 Task(
                     name=processor.id_,
                     displayName=processor.name,
-                    sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{processor.uri}",
+                    sourceUrl=SourceUrl(
+                        f"{clean_uri(self.service_connection.hostPort)}{processor.uri}"
+                    ),
                     taskType=processor.type_,
                     downstreamTasks=self._get_downstream_tasks_from(
                         source_id=processor.id_,
@@ -140,11 +147,13 @@ class NifiSource(PipelineServiceSource):
         :return: Create Pipeline request with tasks
         """
         pipeline_request = CreatePipelineRequest(
-            name=pipeline_details.id_,
+            name=EntityName(pipeline_details.id_),
             displayName=pipeline_details.name,
-            sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{pipeline_details.uri}",
+            sourceUrl=SourceUrl(
+                f"{clean_uri(self.service_connection.hostPort)}{pipeline_details.uri}"
+            ),
             tasks=self._get_tasks_from_details(pipeline_details),
-            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)

metadata/ingestion/source/pipeline/openlineage/connection.py

@@ -47,9 +47,18 @@ def get_connection(connection: OpenLineageConnection) -> KafkaConsumer:
             config.update(
                 {
                     "security.protocol": connection.securityProtocol.value,
-                    "ssl.ca.location": connection.sslConfig.__root__.caCertificate,
-                    "ssl.certificate.location": connection.sslConfig.__root__.sslCertificate,
-                    "ssl.key.location": connection.sslConfig.__root__.sslKey,
+                    "ssl.ca.location": connection.sslConfig.root.caCertificate,
+                    "ssl.certificate.location": connection.sslConfig.root.sslCertificate,
+                    "ssl.key.location": connection.sslConfig.root.sslKey,
+                }
+            )
+        elif connection.securityProtocol.value == KafkaSecProtocol.SASL_SSL.value:
+            config.update(
+                {
+                    "security.protocol": connection.securityProtocol.value,
+                    "sasl.mechanism": connection.saslConfig.saslMechanism.value,
+                    "sasl.username": connection.saslConfig.saslUsername,
+                    "sasl.password": connection.saslConfig.saslPassword,
                 }
             )
 

metadata/ingestion/source/pipeline/openlineage/metadata.py

@@ -81,8 +81,8 @@ class OpenlineageSource(PipelineServiceSource):
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
         """Create class instance"""
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: OpenLineageConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: OpenLineageConnection = config.serviceConnection.root.config
         if not isinstance(connection, OpenLineageConnection):
             raise InvalidSourceException(
                 f"Expected OpenLineageConnection, but got {connection}"
@@ -146,30 +146,6 @@ class OpenlineageSource(PipelineServiceSource):
             except FQNNotFoundException:
                 return None
 
-    def _get_table_fqn_from_om(self, table_details: TableDetails) -> Optional[str]:
-        """
-        Based on partial schema and table names look for matching table object in open metadata.
-        :param schema: schema name
-        :param table: table name
-        :return: fully qualified name of a Table in Open Metadata
-        """
-        result = None
-        services = self.get_db_service_names()
-        for db_service in services:
-            result = fqn.build(
-                metadata=self.metadata,
-                entity_type=Table,
-                service_name=db_service,
-                database_name=None,
-                schema_name=table_details.schema,
-                table_name=table_details.name,
-            )
-        if not result:
-            raise FQNNotFoundException(
-                f"Table FQN not found for table: {table_details} within services: {services}"
-            )
-        return result
-
     def _get_schema_fqn_from_om(self, schema: str) -> Optional[str]:
         """
         Based on partial schema name look for any matching DatabaseSchema object in open metadata.
@@ -452,7 +428,7 @@ class OpenlineageSource(PipelineServiceSource):
                         ),
                         lineageDetails=LineageDetails(
                             pipeline=EntityReference(
-                                id=pipeline_entity.id.__root__,
+                                id=pipeline_entity.id.root,
                                 type="pipeline",
                             ),
                             description=f"Lineage extracted from OpenLineage job: {pipeline_details.job['name']}",

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -14,9 +14,13 @@ Base class for ingesting database 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.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.pipeline import Pipeline
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.pipelineService import (
     PipelineConnection,
     PipelineService,
@@ -33,6 +37,7 @@ from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -42,6 +47,8 @@ from metadata.ingestion.models.topology import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
+from metadata.ingestion.source.pipeline.openlineage.utils import FQNNotFoundException
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_pipeline
 from metadata.utils.logger import ingestion_logger
@@ -57,7 +64,9 @@ class PipelineServiceTopology(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(
@@ -72,7 +81,9 @@ class PipelineServiceTopology(ServiceTopology):
         children=["pipeline"],
         post_process=["mark_pipelines_as_deleted"],
     )
-    pipeline = TopologyNode(
+    pipeline: Annotated[
+        TopologyNode, Field(description="Processing Pipelines Node")
+    ] = TopologyNode(
         producer="get_pipeline",
         stages=[
             NodeStage(
@@ -113,7 +124,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     source_config: PipelineServiceMetadataPipeline
     config: WorkflowSource
     # Big union of types we want to fetch dynamically
-    service_connection: PipelineConnection.__fields__["config"].type_
+    service_connection: PipelineConnection.model_fields["config"].annotation
 
     topology = PipelineServiceTopology()
     context = TopologyContextManager(topology)
@@ -127,7 +138,7 @@ class PipelineServiceSource(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: PipelineServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
@@ -170,10 +181,43 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def yield_pipeline_lineage(
         self, pipeline_details: Any
-    ) -> Iterable[Either[AddLineageRequest]]:
+    ) -> Iterable[Either[OMetaLineageRequest]]:
         """Yields lineage if config is enabled"""
         if self.source_config.includeLineage:
-            yield from self.yield_pipeline_lineage_details(pipeline_details) or []
+            for lineage in self.yield_pipeline_lineage_details(pipeline_details) or []:
+                if lineage.right is not None:
+                    yield Either(
+                        right=OMetaLineageRequest(
+                            lineage_request=lineage.right,
+                            override_lineage=self.source_config.overrideLineage,
+                        )
+                    )
+                else:
+                    yield lineage
+
+    def _get_table_fqn_from_om(self, table_details: TableDetails) -> Optional[str]:
+        """
+        Based on partial schema and table names look for matching table object in open metadata.
+        :param schema: schema name
+        :param table: table name
+        :return: fully qualified name of a Table in Open Metadata
+        """
+        result = None
+        services = self.get_db_service_names()
+        for db_service in services:
+            result = fqn.build(
+                metadata=self.metadata,
+                entity_type=Table,
+                service_name=db_service,
+                database_name=None,
+                schema_name=table_details.schema,
+                table_name=table_details.name,
+            )
+            if result:
+                return result
+        raise FQNNotFoundException(
+            f"Table FQN not found for table: {table_details} within services: {services}"
+        )
 
     def yield_tag(self, *args, **kwargs) -> Iterable[Either[OMetaTagAndClassification]]:
         """Method to fetch pipeline tags"""
@@ -214,8 +258,8 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
         pipeline_fqn = fqn.build(
             self.metadata,
             entity_type=Pipeline,
-            service_name=pipeline_request.service.__root__,
-            pipeline_name=pipeline_request.name.__root__,
+            service_name=pipeline_request.service.root,
+            pipeline_name=pipeline_request.name.root,
         )
 
         self.pipeline_source_state.add(pipeline_fqn)

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

@@ -38,7 +38,7 @@ class SplineClient:
     def __init__(self, config: SplineConnection):
         self.config = config
         client_config: ClientConfig = ClientConfig(
-            base_url=clean_uri(self.config.hostPort),
+            base_url=clean_uri(str(self.config.hostPort)),
             api_version="consumer",
             auth_header=AUTHORIZATION_HEADER,
             auth_token=lambda: (NO_ACCESS_TOKEN, 0),

metadata/ingestion/source/pipeline/spline/metadata.py

@@ -59,8 +59,8 @@ class SplineSource(PipelineServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: SplineConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SplineConnection = config.serviceConnection.root.config
         if not isinstance(connection, SplineConnection):
             raise InvalidSourceException(
                 f"Expected SplineConnection, but got {connection}"
@@ -235,7 +235,7 @@ class SplineSource(PipelineServiceSource):
                             edge=EntitiesEdge(
                                 lineageDetails=LineageDetails(
                                     pipeline=EntityReference(
-                                        id=pipeline_entity.id.__root__,
+                                        id=pipeline_entity.id.root,
                                         type="pipeline",
                                     ),
                                     columnsLineage=[

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

@@ -17,9 +17,9 @@ from pydantic import BaseModel, Field
 
 
 class ExecutionEvent(BaseModel):
-    executionEventId: Optional[str]
-    executionPlanId: Optional[str]
-    applicationName: Optional[str]
+    executionEventId: Optional[str] = None
+    executionPlanId: Optional[str] = None
+    applicationName: Optional[str] = None
 
 
 class ExecutionEvents(BaseModel):
@@ -30,15 +30,15 @@ class ExecutionEvents(BaseModel):
 
 
 class Inputs(BaseModel):
-    source: Optional[str]
+    source: Optional[str] = None
 
 
 class Output(BaseModel):
-    source: Optional[str]
+    source: Optional[str] = None
 
 
 class AttributesNames(BaseModel):
-    id: Optional[str]
+    id: Optional[str] = None
 
 
 class Extra(BaseModel):
@@ -46,8 +46,8 @@ class Extra(BaseModel):
 
 
 class ExecutionPlan(BaseModel):
-    id: Optional[str] = Field(..., alias="_id")
-    name: Optional[str]
+    id: Optional[str] = Field(None, alias="_id")
+    name: Optional[str] = None
     inputs: Optional[List[Inputs]] = []
     output: Optional[Output] = None
     extra: Optional[Extra] = None
@@ -58,13 +58,13 @@ class ExecutionDetail(BaseModel):
 
 
 class ColNodes(BaseModel):
-    id: Optional[str] = Field(..., alias="_id")
-    name: Optional[str]
+    id: Optional[str] = Field(None, alias="_id")
+    name: Optional[str] = None
 
 
 class ColLineage(BaseModel):
-    source: Optional[str]
-    target: Optional[str]
+    source: Optional[str] = None
+    target: Optional[str] = None
 
 
 class Lineage(BaseModel):

metadata/ingestion/source/search/elasticsearch/connection.py

@@ -172,11 +172,11 @@ def get_connection(connection: ElasticsearchConnection) -> Elasticsearch:
         ssl_context = get_ssl_context(connection.sslConfig)
 
     return Elasticsearch(
-        connection.hostPort,
+        str(connection.hostPort),
         http_auth=basic_auth,
         api_key=api_key,
         ssl_context=ssl_context,
-        **connection.connectionArguments.__root__,
+        **connection.connectionArguments.root,
     )
 
 

metadata/ingestion/source/search/elasticsearch/metadata.py

@@ -30,6 +30,7 @@ from metadata.generated.schema.entity.services.connections.search.elasticSearchC
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
 from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
@@ -59,8 +60,8 @@ class ElasticsearchSource(SearchServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: ElasticsearchConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: ElasticsearchConnection = config.serviceConnection.root.config
         if not isinstance(connection, ElasticsearchConnection):
             raise InvalidSourceException(
                 f"Expected ElasticsearchConnection, but got {connection}"
@@ -93,12 +94,12 @@ class ElasticsearchSource(SearchServiceSource):
         index_name = self.get_search_index_name(search_index_details)
         if index_name:
             search_index_request = CreateSearchIndexRequest(
-                name=index_name,
+                name=EntityName(index_name),
                 displayName=index_name,
                 searchIndexSettings=search_index_details.get(index_name, {}).get(
                     "settings", {}
                 ),
-                service=self.context.get().search_service,
+                service=FullyQualifiedEntityName(self.context.get().search_service),
                 fields=parse_es_index_mapping(
                     search_index_details.get(index_name, {}).get("mappings")
                 ),

metadata/ingestion/source/search/search_service.py

@@ -14,6 +14,9 @@ Base class for ingesting search index 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.createSearchIndex import (
     CreateSearchIndexRequest,
 )
@@ -63,7 +66,9 @@ class SearchServiceTopology(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(
@@ -78,7 +83,9 @@ class SearchServiceTopology(ServiceTopology):
         children=["search_index"],
         post_process=["mark_search_indexes_as_deleted"],
     )
-    search_index = TopologyNode(
+    search_index: Annotated[
+        TopologyNode, Field(description="Search Index Processing Node")
+    ] = TopologyNode(
         producer="get_search_index",
         stages=[
             NodeStage(
@@ -107,7 +114,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     source_config: SearchServiceMetadataPipeline
     config: WorkflowSource
     # Big union of types we want to fetch dynamically
-    service_connection: SearchConnection.__fields__["config"].type_
+    service_connection: SearchConnection.model_fields["config"].annotation
 
     topology = SearchServiceTopology()
     context = TopologyContextManager(topology)
@@ -124,7 +131,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
         self.source_config: SearchServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.connection = get_connection(self.service_connection)
 
         # Flag the connection for the test connection
@@ -205,8 +212,8 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
         index_fqn = fqn.build(
             self.metadata,
             entity_type=SearchIndex,
-            service_name=search_index_request.service.__root__,
-            search_index_name=search_index_request.name.__root__,
+            service_name=search_index_request.service.root,
+            search_index_name=search_index_request.name.root,
         )
 
         self.index_source_state.add(index_fqn)

metadata/ingestion/source/storage/gcs/client.py

@@ -0,0 +1,63 @@
+#  Copyright 2024 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""A client for Google Cloud Storage that supports multiple projects."""
+from functools import partial
+from typing import List, Optional, Type, Union
+
+from google import auth
+from google.cloud.monitoring_v3 import MetricServiceClient
+from google.cloud.storage import Client
+
+NoProject = object()
+
+
+class MultiProjectClient:
+    """Google Cloud Client does not support ad-hoc project switching. This class wraps the client and allows
+    switching between projects. If no project is specified, the client will not have a project set and will try
+    to resolve it from ADC.
+    Example usage:
+    ```
+    from google.cloud.storage import Client
+    client = MultiProjectClient(Client, project_ids=["project1", "project2"])
+    buckets_project1 = client.list_buckets("project1")
+    buckets_project2 = client.list_buckets("project2")
+    """
+
+    def __init__(
+        self,
+        client_class: Union[Type[Client], Type[MetricServiceClient]],
+        project_ids: Optional[List[str]] = None,
+        **client_kwargs,
+    ):
+        if project_ids:
+            self.clients = {
+                project_id: client_class(project=project_id, **client_kwargs)
+                for project_id in project_ids
+            }
+        else:
+            self.clients = {NoProject: client_class(**client_kwargs)}
+
+    def project_ids(self):
+        if NoProject in self.clients:
+            _, project_id = auth.default()
+            return [project_id]
+        return list(self.clients.keys())
+
+    def __getattr__(self, client_method):
+        """Return the underlying client method as a partial function so we can inject the project_id."""
+        return partial(self._call, client_method)
+
+    def _call(self, method, project_id, *args, **kwargs):
+        """Call the method on the client for the given project_id. The args and kwargs are passed through."""
+        client = self.clients.get(project_id, self.clients.get(NoProject))
+        if not client:
+            raise ValueError(f"Project {project_id} not found")
+        return getattr(client, method)(*args, **kwargs)