def run(self, parsed_args: argparse.Namespace) -> Iterable[str]: cmd = parsed_args.command if not cmd: raise Exception("No command passed") if hasattr(parsed_args, "debug") and parsed_args.debug: self.logger.setLevel(logging.DEBUG) self.logger.debug("Debug logs enabled") else: self.logger.setLevel(logging.INFO) # todo: add try catch for exceptions with different exit codes source_spec: ConnectorSpecification = self.source.spec(self.logger) with tempfile.TemporaryDirectory() as temp_dir: if cmd == "spec": message = AirbyteMessage(type=Type.SPEC, spec=source_spec) yield message.json(exclude_unset=True) else: raw_config = self.source.read_config(parsed_args.config) config = self.source.configure(raw_config, temp_dir) # Now that we have the config, we can use it to get a list of ai airbyte_secrets # that we should filter in logging to avoid leaking secrets config_secrets = get_secrets( source_spec.connectionSpecification, config) update_secrets(config_secrets) # Remove internal flags from config before validating so # jsonschema's additionalProperties flag wont fail the validation connector_config, _ = split_config(config) if self.source.check_config_against_spec or cmd == "check": check_config_against_spec_or_exit(connector_config, source_spec) if cmd == "check": check_result = self.source.check(self.logger, config) if check_result.status == Status.SUCCEEDED: self.logger.info("Check succeeded") else: self.logger.error("Check failed") output_message = AirbyteMessage( type=Type.CONNECTION_STATUS, connectionStatus=check_result).json(exclude_unset=True) yield output_message elif cmd == "discover": catalog = self.source.discover(self.logger, config) yield AirbyteMessage( type=Type.CATALOG, catalog=catalog).json(exclude_unset=True) elif cmd == "read": config_catalog = self.source.read_catalog( parsed_args.catalog) state = self.source.read_state(parsed_args.state) generator = self.source.read(self.logger, config, config_catalog, state) for message in generator: yield message.json(exclude_unset=True) else: raise Exception("Unexpected command " + cmd)
def run(self, parsed_args: argparse.Namespace) -> Iterable[str]: cmd = parsed_args.command if not cmd: raise Exception("No command passed") # todo: add try catch for exceptions with different exit codes source_spec = self.source.spec(self.logger) with tempfile.TemporaryDirectory() as temp_dir: if cmd == "spec": message = AirbyteMessage(type=Type.SPEC, spec=source_spec) yield message.json(exclude_unset=True) else: raw_config = self.source.read_config(parsed_args.config) config = self.source.configure(raw_config, temp_dir) # Remove internal flags from config before validating so # jsonschema's additionalProperties flag wont fail the validation config, internal_config = split_config(config) if self.source.check_config_against_spec or cmd == "check": check_config_against_spec_or_exit(config, source_spec, self.logger) # Put internal flags back to config dict config.update(internal_config.dict()) if cmd == "check": check_result = self.source.check(self.logger, config) if check_result.status == Status.SUCCEEDED: self.logger.info("Check succeeded") else: self.logger.error("Check failed") output_message = AirbyteMessage( type=Type.CONNECTION_STATUS, connectionStatus=check_result).json(exclude_unset=True) yield output_message elif cmd == "discover": catalog = self.source.discover(self.logger, config) yield AirbyteMessage( type=Type.CATALOG, catalog=catalog).json(exclude_unset=True) elif cmd == "read": config_catalog = self.source.read_catalog( parsed_args.catalog) state = self.source.read_state(parsed_args.state) generator = self.source.read(self.logger, config, config_catalog, state) for message in generator: yield message.json(exclude_unset=True) else: raise Exception("Unexpected command " + cmd)
def test_configure_catalog(): stream = AirbyteStream(name="stream", supported_sync_modes=[SyncMode.full_refresh], json_schema={}) catalog = AirbyteCatalog(streams=[stream]) catalog_message = AirbyteMessage(type=Type.CATALOG, catalog=catalog) sys.stdin = io.StringIO(catalog_message.json()) expected_configured_catalog = ConfiguredAirbyteCatalog(streams=[ ConfiguredAirbyteStream( stream=stream, sync_mode=SyncMode.full_refresh, destination_sync_mode=DestinationSyncMode.append) ]) expected_configured_catalog_json = json.loads( expected_configured_catalog.json()) with tempfile.TemporaryDirectory() as temp_dir: os.chdir(temp_dir) configure_catalog() assert os.path.exists("integration_tests/configured_catalog.json") with open("integration_tests/configured_catalog.json") as f: configured_catalog_json = json.loads(f.read()) assert configured_catalog_json == expected_configured_catalog_json
def format(self, record: logging.LogRecord) -> str: """Return a JSON representation of the log message""" message = super().format(record) airbyte_level = self.level_mapping.get(record.levelno, "INFO") log_message = AirbyteMessage(type="LOG", log=AirbyteLogMessage(level=airbyte_level, message=message)) return log_message.json(exclude_unset=True)
def run(self, parsed_args: argparse.Namespace) -> Iterable[str]: cmd = parsed_args.command if not cmd: raise Exception("No command passed") # todo: add try catch for exceptions with different exit codes with tempfile.TemporaryDirectory() as temp_dir: if cmd == "spec": message = AirbyteMessage(type=Type.SPEC, spec=self.source.spec(logger)) yield message.json(exclude_unset=True) else: raw_config = self.source.read_config(parsed_args.config) config = self.source.configure(raw_config, temp_dir) if cmd == "check": check_result = self.source.check(logger, config) if check_result.status == Status.SUCCEEDED: logger.info("Check succeeded") else: logger.error("Check failed") output_message = AirbyteMessage( type=Type.CONNECTION_STATUS, connectionStatus=check_result).json(exclude_unset=True) yield output_message elif cmd == "discover": catalog = self.source.discover(logger, config) yield AirbyteMessage( type=Type.CATALOG, catalog=catalog).json(exclude_unset=True) elif cmd == "read": config_catalog = self.source.read_catalog( parsed_args.catalog) state = self.source.read_state(parsed_args.state) generator = self.source.read(logger, config, config_catalog, state) for message in generator: yield message.json(exclude_unset=True) else: raise Exception("Unexpected command " + cmd)
def test_unhandled_logger(): cmd = "from airbyte_cdk.logger import init_logger; init_logger('airbyte'); raise 1" expected_message = ( "exceptions must derive from BaseException\n" "Traceback (most recent call last):\n" ' File "<string>", line 1, in <module>\n' "TypeError: exceptions must derive from BaseException" ) log_message = AirbyteMessage(type="LOG", log=AirbyteLogMessage(level="FATAL", message=expected_message)) expected_output = log_message.json(exclude_unset=True) with pytest.raises(subprocess.CalledProcessError) as err: subprocess.check_output([sys.executable, "-c", cmd], stderr=subprocess.STDOUT) assert not err.value.stderr, "nothing on the stderr" assert err.value.output.decode("utf-8").strip() == expected_output, "Error should be printed in expected form"
def _wrap_message( submessage: Union[AirbyteConnectionStatus, ConnectorSpecification, AirbyteRecordMessage, AirbyteCatalog] ) -> str: if isinstance(submessage, AirbyteConnectionStatus): message = AirbyteMessage(type=Type.CONNECTION_STATUS, connectionStatus=submessage) elif isinstance(submessage, ConnectorSpecification): message = AirbyteMessage(type=Type.SPEC, spec=submessage) elif isinstance(submessage, AirbyteCatalog): message = AirbyteMessage(type=Type.CATALOG, catalog=submessage) elif isinstance(submessage, AirbyteRecordMessage): message = AirbyteMessage(type=Type.RECORD, record=submessage) else: raise Exception(f"Unknown message type: {submessage}") return message.json(exclude_unset=True)
def start(self, args): # set up parent parsers parent_parser = argparse.ArgumentParser(add_help=False) main_parser = argparse.ArgumentParser() subparsers = main_parser.add_subparsers(title="commands", dest="command") # spec subparsers.add_parser( "spec", help="outputs the json configuration specification", parents=[parent_parser]) # check check_parser = subparsers.add_parser( "check", help="checks the config can be used to connect", parents=[parent_parser]) required_check_parser = check_parser.add_argument_group( "required named arguments") required_check_parser.add_argument( "--config", type=str, required=True, help="path to the json configuration file") # discover discover_parser = subparsers.add_parser( "discover", help="outputs a catalog describing the source's schema", parents=[parent_parser]) required_discover_parser = discover_parser.add_argument_group( "required named arguments") required_discover_parser.add_argument( "--config", type=str, required=True, help="path to the json configuration file") # read read_parser = subparsers.add_parser( "read", help="reads the source and outputs messages to STDOUT", parents=[parent_parser]) read_parser.add_argument("--state", type=str, required=False, help="path to the json-encoded state file") required_read_parser = read_parser.add_argument_group( "required named arguments") required_read_parser.add_argument( "--config", type=str, required=True, help="path to the json configuration file") required_read_parser.add_argument( "--catalog", type=str, required=True, help="path to the catalog used to determine which data to read") # parse the args parsed_args = main_parser.parse_args(args) # execute cmd = parsed_args.command if not cmd: raise Exception("No command passed") # todo: add try catch for exceptions with different exit codes with tempfile.TemporaryDirectory() as temp_dir: if cmd == "spec": message = AirbyteMessage(type=Type.SPEC, spec=self.source.spec(logger)) print(message.json(exclude_unset=True)) sys.exit(0) raw_config = self.source.read_config(parsed_args.config) config = self.source.configure(raw_config, temp_dir) if cmd == "check": check_result = self.source.check(logger, config) if check_result.status == Status.SUCCEEDED: logger.info("Check succeeded") else: logger.error("Check failed") output_message = AirbyteMessage( type=Type.CONNECTION_STATUS, connectionStatus=check_result).json(exclude_unset=True) print(output_message) sys.exit(0) elif cmd == "discover": catalog = self.source.discover(logger, config) print( AirbyteMessage(type=Type.CATALOG, catalog=catalog).json(exclude_unset=True)) sys.exit(0) elif cmd == "read": catalog = self.source.read_catalog(parsed_args.catalog) state = self.source.read_state(parsed_args.state) generator = self.source.read(logger, config, catalog, state) for message in generator: print(message.json(exclude_unset=True)) sys.exit(0) else: raise Exception("Unexpected command " + cmd)
def log(self, level, message): log_record = AirbyteLogMessage(level=level, message=message) log_message = AirbyteMessage(type="LOG", log=log_record) print(log_message.json(exclude_unset=True))