d54a68640f
* Improving error messages to have better messaging in datadog and the frontend * fixing tests
79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
|
|
import sys
|
|
from typing import Any, List, Mapping, Optional, Tuple
|
|
|
|
from airbyte_cdk.connector import BaseConnector
|
|
from airbyte_cdk.connector_builder.connector_builder_handler import (
|
|
TestReadLimits,
|
|
create_source,
|
|
get_limits,
|
|
list_streams,
|
|
read_stream,
|
|
resolve_manifest,
|
|
)
|
|
from airbyte_cdk.entrypoint import AirbyteEntrypoint
|
|
from airbyte_cdk.models import ConfiguredAirbyteCatalog
|
|
from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
|
|
from airbyte_cdk.utils.traced_exception import AirbyteTracedException
|
|
|
|
|
|
def get_config_and_catalog_from_args(args: List[str]) -> Tuple[str, Mapping[str, Any], Optional[ConfiguredAirbyteCatalog]]:
|
|
parsed_args = AirbyteEntrypoint.parse_args(args)
|
|
config_path, catalog_path = parsed_args.config, parsed_args.catalog
|
|
if parsed_args.command != "read":
|
|
raise ValueError("Only read commands are allowed for Connector Builder requests.")
|
|
|
|
config = BaseConnector.read_config(config_path)
|
|
|
|
if "__command" not in config:
|
|
raise ValueError(
|
|
f"Invalid config: `__command` should be provided at the root of the config but config only has keys {list(config.keys())}"
|
|
)
|
|
|
|
command = config["__command"]
|
|
if command == "test_read":
|
|
catalog = ConfiguredAirbyteCatalog.parse_obj(BaseConnector.read_config(catalog_path))
|
|
else:
|
|
catalog = None
|
|
|
|
if "__injected_declarative_manifest" not in config:
|
|
raise ValueError(
|
|
f"Invalid config: `__injected_declarative_manifest` should be provided at the root of the config but config only has keys {list(config.keys())}"
|
|
)
|
|
|
|
return command, config, catalog
|
|
|
|
|
|
def handle_connector_builder_request(
|
|
source: ManifestDeclarativeSource, command: str, config: Mapping[str, Any], catalog: Optional[ConfiguredAirbyteCatalog], limits: TestReadLimits
|
|
):
|
|
if command == "resolve_manifest":
|
|
return resolve_manifest(source)
|
|
elif command == "test_read":
|
|
assert catalog is not None, "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
|
|
return read_stream(source, config, catalog, limits)
|
|
elif command == "list_streams":
|
|
return list_streams(source, config)
|
|
else:
|
|
raise ValueError(f"Unrecognized command {command}.")
|
|
|
|
|
|
def handle_request(args: List[str]):
|
|
command, config, catalog = get_config_and_catalog_from_args(args)
|
|
limits = get_limits(config)
|
|
source = create_source(config, limits)
|
|
return handle_connector_builder_request(source, command, config, catalog, limits).json(exclude_unset=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
print(handle_request(sys.argv[1:]))
|
|
except Exception as exc:
|
|
error = AirbyteTracedException.from_exception(exc, message=f"Error handling request: {str(exc)}")
|
|
m = error.as_airbyte_message()
|
|
print(error.as_airbyte_message().json(exclude_unset=True))
|