Bugfix/ingest pipeline check (#3303)

### Description
Using a `isinstance` on the destination registry mapping breaks when
inheritance is used for the associated uploader types. This adds a
connector type field to all uploaders so that the entry can be
deterministically fetched when running check for associated stager in
pipeline.
This commit is contained in:
Roman Isecke 2024-06-27 12:35:37 -04:00 committed by GitHub
parent 087adb218f
commit 137b149be8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 25 additions and 19 deletions

View File

@ -1,4 +1,4 @@
## 0.14.9-dev7 ## 0.14.9-dev8
### Enhancements ### Enhancements

View File

@ -1 +1 @@
__version__ = "0.14.9-dev7" # pragma: no cover __version__ = "0.14.9-dev8" # pragma: no cover

View File

@ -26,6 +26,7 @@ class UploadContent:
@dataclass @dataclass
class Uploader(BaseProcess, BaseConnector, ABC): class Uploader(BaseProcess, BaseConnector, ABC):
upload_config: UploaderConfigT upload_config: UploaderConfigT
connector_type: str
def is_async(self) -> bool: def is_async(self) -> bool:
return False return False

View File

@ -87,17 +87,8 @@ class Pipeline:
# Make sure that if the set destination connector expects a stager, one is also set # Make sure that if the set destination connector expects a stager, one is also set
if not self.uploader_step: if not self.uploader_step:
return return
matching_registry_entry = [ uploader_connector_type = self.uploader_step.process.connector_type
v registry_entry = destination_registry[uploader_connector_type]
for v in destination_registry.values()
if isinstance(self.uploader_step.process, v.uploader)
]
if len(matching_registry_entry) > 1:
raise ValueError(
f"More than one entry found in destination registry "
f"for uploader type: {self.uploader_step.process}"
)
registry_entry = matching_registry_entry[0]
if registry_entry.upload_stager and self.stager_step is None: if registry_entry.upload_stager and self.stager_step is None:
raise ValueError( raise ValueError(
f"pipeline with uploader type {self.uploader_step.process.__class__.__name__} " f"pipeline with uploader type {self.uploader_step.process.__class__.__name__} "

View File

@ -92,6 +92,7 @@ class AstraUploaderConfig(UploaderConfig):
class AstraUploader(Uploader): class AstraUploader(Uploader):
connection_config: AstraConnectionConfig connection_config: AstraConnectionConfig
upload_config: AstraUploaderConfig upload_config: AstraUploaderConfig
connector_type: str = CONNECTOR_TYPE
@requires_dependencies(["astrapy"], extras="astra") @requires_dependencies(["astrapy"], extras="astra")
def get_collection(self) -> "AstraDBCollection": def get_collection(self) -> "AstraDBCollection":

View File

@ -114,6 +114,7 @@ class ChromaUploaderConfig(UploaderConfig):
@dataclass @dataclass
class ChromaUploader(Uploader): class ChromaUploader(Uploader):
connector_type: str = CONNECTOR_TYPE
upload_config: ChromaUploaderConfig upload_config: ChromaUploaderConfig
connection_config: ChromaConnectionConfig connection_config: ChromaConnectionConfig
client: Optional["Client"] = field(init=False) client: Optional["Client"] = field(init=False)

View File

@ -327,6 +327,7 @@ class ElasticsearchUploaderConfig(UploaderConfig):
@dataclass @dataclass
class ElasticsearchUploader(Uploader): class ElasticsearchUploader(Uploader):
connector_type: str = CONNECTOR_TYPE
upload_config: ElasticsearchUploaderConfig upload_config: ElasticsearchUploaderConfig
connection_config: ElasticsearchConnectionConfig connection_config: ElasticsearchConnectionConfig

View File

@ -114,6 +114,7 @@ class AzureUploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class AzureUploader(FsspecUploader): class AzureUploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: AzureConnectionConfig connection_config: AzureConnectionConfig
upload_config: AzureUploaderConfig = field(default=None) upload_config: AzureUploaderConfig = field(default=None)

View File

@ -100,7 +100,8 @@ class BoxUploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class BoxUpload(FsspecUploader): class BoxUploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: BoxConnectionConfig connection_config: BoxConnectionConfig
upload_config: BoxUploaderConfig = field(default=None) upload_config: BoxUploaderConfig = field(default=None)
@ -131,6 +132,8 @@ add_source_entry(
add_destination_entry( add_destination_entry(
destination_type=CONNECTOR_TYPE, destination_type=CONNECTOR_TYPE,
entry=DestinationRegistryEntry( entry=DestinationRegistryEntry(
uploader=BoxUpload, uploader_config=BoxUploaderConfig, connection_config=BoxConnectionConfig uploader=BoxUploader,
uploader_config=BoxUploaderConfig,
connection_config=BoxConnectionConfig,
), ),
) )

View File

@ -99,7 +99,8 @@ class DropboxUploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class DropboxUpload(FsspecUploader): class DropboxUploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: DropboxConnectionConfig connection_config: DropboxConnectionConfig
upload_config: DropboxUploaderConfig = field(default=None) upload_config: DropboxUploaderConfig = field(default=None)
@ -130,7 +131,7 @@ add_source_entry(
add_destination_entry( add_destination_entry(
destination_type=CONNECTOR_TYPE, destination_type=CONNECTOR_TYPE,
entry=DestinationRegistryEntry( entry=DestinationRegistryEntry(
uploader=DropboxUpload, uploader=DropboxUploader,
uploader_config=DropboxUploaderConfig, uploader_config=DropboxUploaderConfig,
connection_config=DropboxConnectionConfig, connection_config=DropboxConnectionConfig,
), ),

View File

@ -304,6 +304,7 @@ FsspecUploaderConfigT = TypeVar("FsspecUploaderConfigT", bound=FsspecUploaderCon
@dataclass @dataclass
class FsspecUploader(Uploader): class FsspecUploader(Uploader):
connector_type: str = CONNECTOR_TYPE
upload_config: FsspecUploaderConfigT = field(default=None) upload_config: FsspecUploaderConfigT = field(default=None)
@property @property

View File

@ -111,6 +111,7 @@ class GcsUploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class GcsUploader(FsspecUploader): class GcsUploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: GcsConnectionConfig connection_config: GcsConnectionConfig
upload_config: GcsUploaderConfig = field(default=None) upload_config: GcsUploaderConfig = field(default=None)

View File

@ -125,7 +125,8 @@ class S3UploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class S3Upload(FsspecUploader): class S3Uploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: S3ConnectionConfig connection_config: S3ConnectionConfig
upload_config: S3UploaderConfig = field(default=None) upload_config: S3UploaderConfig = field(default=None)
@ -156,7 +157,7 @@ add_source_entry(
add_destination_entry( add_destination_entry(
destination_type=CONNECTOR_TYPE, destination_type=CONNECTOR_TYPE,
entry=DestinationRegistryEntry( entry=DestinationRegistryEntry(
uploader=S3Upload, uploader=S3Uploader,
uploader_config=S3UploaderConfig, uploader_config=S3UploaderConfig,
connection_config=S3ConnectionConfig, connection_config=S3ConnectionConfig,
), ),

View File

@ -136,6 +136,7 @@ class SftpUploaderConfig(FsspecUploaderConfig):
@dataclass @dataclass
class SftpUploader(FsspecUploader): class SftpUploader(FsspecUploader):
connector_type: str = CONNECTOR_TYPE
connection_config: SftpConnectionConfig connection_config: SftpConnectionConfig
upload_config: SftpUploaderConfig = field(default=None) upload_config: SftpUploaderConfig = field(default=None)

View File

@ -160,6 +160,7 @@ class LocalUploaderConfig(UploaderConfig):
@dataclass @dataclass
class LocalUploader(Uploader): class LocalUploader(Uploader):
connector_type: str = CONNECTOR_TYPE
upload_config: LocalUploaderConfig = field(default_factory=lambda: LocalUploaderConfig()) upload_config: LocalUploaderConfig = field(default_factory=lambda: LocalUploaderConfig())
connection_config: LocalConnectionConfig = field( connection_config: LocalConnectionConfig = field(
default_factory=lambda: LocalConnectionConfig() default_factory=lambda: LocalConnectionConfig()

View File

@ -154,6 +154,7 @@ class WeaviateUploaderConfig(UploaderConfig):
@dataclass @dataclass
class WeaviateUploader(Uploader): class WeaviateUploader(Uploader):
connector_type: str = CONNECTOR_TYPE
upload_config: WeaviateUploaderConfig upload_config: WeaviateUploaderConfig
connection_config: WeaviateConnectionConfig connection_config: WeaviateConnectionConfig
client: Optional["Client"] = field(init=False) client: Optional["Client"] = field(init=False)