AirbyteLib: Progress Printer

airbyte-lib/airbyte_lib/_factories/connector_factories.py

@@ -1,11 +1,13 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 from __future__ import annotations
 
+import shutil
+from pathlib import Path
 from typing import Any
 
-from airbyte_lib._executor import Executor, PathExecutor, VenvExecutor
-from airbyte_lib.exceptions import AirbyteLibInputError
-from airbyte_lib.registry import get_connector_metadata
+from airbyte_lib import exceptions as exc
+from airbyte_lib._executor import PathExecutor, VenvExecutor
+from airbyte_lib.registry import ConnectorMetadata, get_connector_metadata
 from airbyte_lib.source import Source
 
 
@@ -15,7 +17,7 @@ def get_connector(
     pip_url: str | None = None,
     config: dict[str, Any] | None = None,
     *,
-    use_local_install: bool = False,
+    local_executable: Path | str | None = None,
     install_if_missing: bool = True,
 ) -> Source:
     """Get a connector by name and version.
@@ -29,34 +31,58 @@ def get_connector(
             connector name.
         config: connector config - if not provided, you need to set it later via the set_config
             method.
-        use_local_install: whether to use a virtual environment to run the connector. If True, the
-            connector is expected to be available on the path (e.g. installed via pip). If False,
-            the connector will be installed automatically in a virtual environment.
-        install_if_missing: whether to install the connector if it is not available locally. This
-            parameter is ignored if use_local_install is True.
+        local_executable: If set, the connector will be assumed to already be installed and will be
+            executed using this path or executable name. Otherwise, the connector will be installed
+            automatically in a virtual environment.
+        install_if_missing: Whether to install the connector if it is not available locally. This
+            parameter is ignored when local_executable is set.
     """
-    metadata = get_connector_metadata(name)
-    if use_local_install:
+    if local_executable:
         if pip_url:
-            raise AirbyteLibInputError(
-                message="Param 'pip_url' is not supported when 'use_local_install' is True."
+            raise exc.AirbyteLibInputError(
+                message="Param 'pip_url' is not supported when 'local_executable' is set."
             )
         if version:
-            raise AirbyteLibInputError(
-                message="Param 'version' is not supported when 'use_local_install' is True."
+            raise exc.AirbyteLibInputError(
+                message="Param 'version' is not supported when 'local_executable' is set."
             )
-        executor: Executor = PathExecutor(
-            metadata=metadata,
-            target_version=version,
-        )
 
-    else:
-        executor = VenvExecutor(
-            metadata=metadata,
-            target_version=version,
-            install_if_missing=install_if_missing,
-            pip_url=pip_url,
+        if isinstance(local_executable, str):
+            if "/" in local_executable or "\\" in local_executable:
+                # Assume this is a path
+                local_executable = Path(local_executable).absolute()
+            else:
+                which_executable = shutil.which(local_executable)
+                if which_executable is None:
+                    raise FileNotFoundError(local_executable)
+                local_executable = Path(which_executable).absolute()
+
+        return Source(
+            name=name,
+            config=config,
+            executor=PathExecutor(
+                name=name,
+                path=local_executable,
+            ),
         )
+
+    metadata: ConnectorMetadata | None = None
+    try:
+        metadata = get_connector_metadata(name)
+    except exc.AirbyteConnectorNotRegisteredError:
+        if not pip_url:
+            # We don't have a pip url or registry entry, so we can't install the connector
+            raise
+
+    executor = VenvExecutor(
+        name=name,
+        metadata=metadata,
+        target_version=version,
+        pip_url=pip_url,
+    )
+    if install_if_missing:
+        executor.ensure_installation()
+
     return Source(
         executor=executor,
         name=name,

airbyte-lib/airbyte_lib/_processors.py

@@ -32,6 +32,7 @@
 
 from airbyte_lib import exceptions as exc
 from airbyte_lib._util import protocol_util  # Internal utility functions
+from airbyte_lib.progress import progress
 
 
 if TYPE_CHECKING:
@@ -40,7 +41,7 @@
     from airbyte_lib.config import CacheConfigBase
 
 
-DEFAULT_BATCH_SIZE = 10000
+DEFAULT_BATCH_SIZE = 10_000
 
 
 class BatchHandle:
@@ -95,7 +96,7 @@ def register_source(
         For now, only one source at a time is supported.
         If this method is called multiple times, the last call will overwrite the previous one.
 
-        TODO: Expand this to handle mutliple sources.
+        TODO: Expand this to handle multiple sources.
         """
         _ = source_name
         self.source_catalog = incoming_source_catalog
@@ -157,6 +158,7 @@ def process_airbyte_messages(
                 if len(stream_batch) >= max_batch_size:
                     record_batch = pa.Table.from_pylist(stream_batch)
                     self._process_batch(stream_name, record_batch)
+                    progress.log_batch_written(stream_name, len(stream_batch))
                     stream_batch.clear()
 
             elif message.type is Type.STATE:
@@ -180,14 +182,16 @@ def process_airbyte_messages(
                 )
 
         # We are at the end of the stream. Process whatever else is queued.
-        for stream_name, batch in stream_batches.items():
-            if batch:
-                record_batch = pa.Table.from_pylist(batch)
+        for stream_name, stream_batch in stream_batches.items():
+            if stream_batch:
+                record_batch = pa.Table.from_pylist(stream_batch)
                 self._process_batch(stream_name, record_batch)
+                progress.log_batch_written(stream_name, len(stream_batch))
 
         # Finalize any pending batches
         for stream_name in list(self._pending_batches.keys()):
             self._finalize_batches(stream_name)
+            progress.log_stream_finalized(stream_name)
 
     @final
     def _process_batch(
@@ -287,7 +291,10 @@ def _finalizing_batches(
         state_messages_to_finalize = self._pending_state_messages[stream_name].copy()
         self._pending_batches[stream_name].clear()
         self._pending_state_messages[stream_name].clear()
+
+        progress.log_batches_finalizing(stream_name, len(batches_to_finalize))
         yield batches_to_finalize
+        progress.log_batches_finalized(stream_name, len(batches_to_finalize))
 
         self._finalized_batches[stream_name].update(batches_to_finalize)
         self._finalized_state_messages[stream_name] += state_messages_to_finalize

airbyte-lib/airbyte_lib/exceptions.py

@@ -174,6 +174,14 @@ class AirbyteConnectorRegistryError(AirbyteError):
     """Error when accessing the connector registry."""
 
 
+@dataclass
+class AirbyteConnectorNotRegisteredError(AirbyteConnectorRegistryError):
+    """Connector not found in registry."""
+
+    connector_name: str | None = None
+    guidance = "Please double check the connector name."
+
+
 # Connector Errors
 
 
@@ -184,8 +192,8 @@ class AirbyteConnectorError(AirbyteError):
     connector_name: str | None = None
 
 
-class AirbyteConnectorNotFoundError(AirbyteConnectorError):
-    """Connector not found."""
+class AirbyteConnectorExecutableNotFoundError(AirbyteConnectorError):
+    """Connector executable not found."""
 
 
 class AirbyteConnectorInstallationError(AirbyteConnectorError):