airbyte_cdk.destinations

 31class Destination(Connector, ABC):
 32    VALID_CMDS = {"spec", "check", "write"}
 33
 34    @abstractmethod
 35    def write(
 36        self,
 37        config: Mapping[str, Any],
 38        configured_catalog: ConfiguredAirbyteCatalog,
 39        input_messages: Iterable[AirbyteMessage],
 40    ) -> Iterable[AirbyteMessage]:
 41        """Implement to define how the connector writes data to the destination"""
 42
 43    def _run_check(self, config: Mapping[str, Any]) -> AirbyteMessage:
 44        check_result = self.check(logger, config)
 45        return AirbyteMessage(type=Type.CONNECTION_STATUS, connectionStatus=check_result)
 46
 47    def _parse_input_stream(self, input_stream: io.TextIOWrapper) -> Iterable[AirbyteMessage]:
 48        """Reads from stdin, converting to Airbyte messages"""
 49        for line in input_stream:
 50            try:
 51                yield AirbyteMessageSerializer.load(orjson.loads(line))
 52            except orjson.JSONDecodeError:
 53                logger.info(
 54                    f"ignoring input which can't be deserialized as Airbyte Message: {line}"
 55                )
 56
 57    def _run_write(
 58        self,
 59        config: Mapping[str, Any],
 60        configured_catalog_path: str,
 61        input_stream: io.TextIOWrapper,
 62    ) -> Iterable[AirbyteMessage]:
 63        catalog = ConfiguredAirbyteCatalogSerializer.load(
 64            orjson.loads(open(configured_catalog_path).read())
 65        )
 66        input_messages = self._parse_input_stream(input_stream)
 67        logger.info("Begin writing to the destination...")
 68        yield from self.write(
 69            config=config, configured_catalog=catalog, input_messages=input_messages
 70        )
 71        logger.info("Writing complete.")
 72
 73    def parse_args(self, args: List[str]) -> argparse.Namespace:
 74        """
 75        :param args: commandline arguments
 76        :return:
 77        """
 78
 79        parent_parser = argparse.ArgumentParser(add_help=False)
 80        main_parser = argparse.ArgumentParser()
 81        subparsers = main_parser.add_subparsers(title="commands", dest="command")
 82
 83        # spec
 84        subparsers.add_parser(
 85            "spec", help="outputs the json configuration specification", parents=[parent_parser]
 86        )
 87
 88        # check
 89        check_parser = subparsers.add_parser(
 90            "check", help="checks the config can be used to connect", parents=[parent_parser]
 91        )
 92        required_check_parser = check_parser.add_argument_group("required named arguments")
 93        required_check_parser.add_argument(
 94            "--config", type=str, required=True, help="path to the json configuration file"
 95        )
 96
 97        # write
 98        write_parser = subparsers.add_parser(
 99            "write", help="Writes data to the destination", parents=[parent_parser]
100        )
101        write_required = write_parser.add_argument_group("required named arguments")
102        write_required.add_argument(
103            "--config", type=str, required=True, help="path to the JSON configuration file"
104        )
105        write_required.add_argument(
106            "--catalog", type=str, required=True, help="path to the configured catalog JSON file"
107        )
108
109        parsed_args = main_parser.parse_args(args)
110        cmd = parsed_args.command
111        if not cmd:
112            raise Exception("No command entered. ")
113        elif cmd not in ["spec", "check", "write"]:
114            # This is technically dead code since parse_args() would fail if this was the case
115            # But it's non-obvious enough to warrant placing it here anyways
116            raise Exception(f"Unknown command entered: {cmd}")
117
118        return parsed_args
119
120    def run_cmd(self, parsed_args: argparse.Namespace) -> Iterable[AirbyteMessage]:
121        cmd = parsed_args.command
122        if cmd not in self.VALID_CMDS:
123            raise Exception(f"Unrecognized command: {cmd}")
124
125        spec = self.spec(logger)
126        if cmd == "spec":
127            yield AirbyteMessage(type=Type.SPEC, spec=spec)
128            return
129        config = self.read_config(config_path=parsed_args.config)
130        # Register the config's secrets with the log filter before anything can print them,
131        # as the source entrypoint does. A config validation error quotes config values.
132        update_secrets(get_secrets(spec.connectionSpecification, config))
133        if self.check_config_against_spec or cmd == "check":
134            try:
135                check_config_against_spec_or_exit(config, spec)
136            except AirbyteTracedException as traced_exc:
137                connection_status = traced_exc.as_connection_status_message()
138                if connection_status and cmd == "check":
139                    yield connection_status
140                    return
141                raise traced_exc
142
143        if cmd == "check":
144            yield self._run_check(config=config)
145        elif cmd == "write":
146            # Wrap in UTF-8 to override any other input encodings
147            wrapped_stdin = io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")
148            yield from self._run_write(
149                config=config,
150                configured_catalog_path=parsed_args.catalog,
151                input_stream=wrapped_stdin,
152            )
153
154    def run(self, args: List[str]) -> None:
155        init_uncaught_exception_handler(logger)
156        parsed_args = self.parse_args(args)
157        output_messages = self.run_cmd(parsed_args)
158        for message in output_messages:
159            print(orjson.dumps(AirbyteMessageSerializer.dump(message)).decode())