def check(self, logger: AirbyteLogger,
config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the destination with the needed permissions
e.g: if a provided API token or password can be used to connect and write to the destination.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this destination, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
# Verify write access by attempting to write then delete
stream = str(uuid.uuid4())
# Can't override destination_path because we cannot assume we have write
# access anywhere else
with SftpClient(**config) as writer:
writer.write(stream, {"value": "_airbyte_connection_check"})
writer.delete(stream)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=
f"An exception occurred: {e}. \nStacktrace: \n{traceback.format_exc()}",
)
def check_config(self, logger: AirbyteLogger, config_path: str, config: json) -> AirbyteConnectionStatus:
token = {"refresh_token": config["refresh_token"], "token_type": "Bearer", "access_token": "wrong", "expires_in": "-30"}
extra = {"client_id": config["client_id"], "client_secret": config["client_secret"]}
sandbox = bool(config.get("sandbox", False))
user_agent = config["user_agent"]
realm_id = config["realm_id"]
session = OAuth2Session(
config["client_id"],
token=token,
auto_refresh_url=TOKEN_REFRESH_URL,
auto_refresh_kwargs=extra,
token_updater=self._write_config,
)
endpoint = f"/v3/company/{realm_id}/query"
params = {"query": "SELECT * FROM CompanyInfo"}
headers = {"Accept": "application/json", "User-Agent": user_agent}
if sandbox:
full_url = SANDBOX_ENDPOINT_BASE + endpoint
else:
full_url = PROD_ENDPOINT_BASE + endpoint
try:
session.request("GET", full_url, headers=headers, params=params)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(status=Status.FAILED, message=f"An exception occurred: {str(e)}")
def check_config(self, logger: AirbyteLogger, config_path: str,
config: json) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration
e.g: if a provided Stripe API token can be used to connect to the Stripe API.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config_path: Path to the file containing the configuration json config
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
url = "https://api.kustomerapp.com/v1/customers"
headers = {"Authorization": f"Bearer {config['api_token']}"}
try:
session = requests.get(url, headers=headers)
session.raise_for_status()
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except requests.exceptions.RequestException as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {repr(e)}")
def check(self, logger: AirbyteLogger,
config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the destination with the needed permissions
e.g: if a provided API token or password can be used to connect and write to the destination.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this destination, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
with establish_connection(config, logger) as connection:
# We can only verify correctness of connection parameters on execution
with connection.cursor() as cursor:
cursor.execute("SELECT 1")
# Test access to the bucket, if S3 strategy is used
create_firebolt_wirter(connection, config, logger)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {repr(e)}")
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the Apify integration.
This is tested by trying to access the Apify user object with the provided userId and Apify token.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
dataset_id = config["datasetId"]
dataset = ApifyClient().dataset(dataset_id).get()
if dataset is None:
raise ValueError(f"Dataset {dataset_id} does not exist")
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
else:
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def health_check(self):
try:
txid = self._client.txid()
if txid is None:
raise ValueError("Could not connect to Dawa")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(status=Status.FAILED,
message=str(e))
def check_config(self, logger: AirbyteLogger, config_path: str, config: json) -> AirbyteConnectionStatus:
try:
self.try_connect(logger, config)
except self.api_error as err:
logger.error(f"Exception while connecting to {self.tap_name}: {err}")
# this should be in UI
error_msg = f"Unable to connect to {self.tap_name} with the provided credentials. Error: {err}"
return AirbyteConnectionStatus(status=Status.FAILED, message=error_msg)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def check(self, logger: AirbyteLogger,
config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""Check connection"""
client = self._get_client(config)
alive, error = client.health_check()
if not alive:
return AirbyteConnectionStatus(status=Status.FAILED,
message=str(error))
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""Implements the Check Connection operation from the Airbyte Specification. See https://docs.airbyte.io/architecture/airbyte-specification."""
try:
check_succeeded, error = self.check_connection(logger, config)
if not check_succeeded:
return AirbyteConnectionStatus(status=Status.FAILED, message=str(error))
except Exception as e:
return AirbyteConnectionStatus(status=Status.FAILED, message=str(e))
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
client = self._get_client(config)
logger.info("Checking access to Amazon SP-API")
try:
client.check_connection()
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
try:
reader = Reader(logger, config)
client = reader.get_table_service()
tables_iterator = client.list_tables(results_per_page=1)
next(tables_iterator)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except StopIteration:
logger.log("No tables found, but credentials are correct.")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
def check(self, logger, config: Mapping) -> AirbyteConnectionStatus:
"""
Check involves verifying that the specified file is reachable with
our credentials.
"""
client = self._get_client(config)
logger.info(f"Checking access to {client.reader.full_url}...")
try:
with client.reader.open(binary=client.binary_source):
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as err:
reason = f"Failed to load {client.reader.full_url}: {repr(err)}\n{traceback.format_exc()}"
logger.error(reason)
return AirbyteConnectionStatus(status=Status.FAILED,
message=reason)
def test_run_check(entrypoint: AirbyteEntrypoint, mocker, spec_mock,
config_mock):
parsed_args = Namespace(command="check", config="config_path")
check_value = AirbyteConnectionStatus(status=Status.SUCCEEDED)
mocker.patch.object(MockSource, "check", return_value=check_value)
assert [_wrap_message(check_value)] == list(entrypoint.run(parsed_args))
assert spec_mock.called
def test_run_check(self, mocker, destination: Destination, tmp_path):
file_path = tmp_path / "config.json"
dummy_config = {"user": "******"}
write_file(file_path, dummy_config)
args = {"command": "check", "config": file_path}
parsed_args = argparse.Namespace(**args)
destination.run_cmd(parsed_args)
spec_msg = ConnectorSpecification(connectionSpecification={})
mocker.patch.object(destination, "spec", return_value=spec_msg)
validate_mock = mocker.patch(
"airbyte_cdk.destinations.destination.check_config_against_spec_or_exit"
)
expected_check_result = AirbyteConnectionStatus(
status=Status.SUCCEEDED)
mocker.patch.object(destination,
"check",
return_value=expected_check_result,
autospec=True)
returned_check_result = next(iter(destination.run_cmd(parsed_args)))
# verify method call with the correct params
# Affirm to Mypy that this is indeed a method on this mock
destination.check.assert_called_once() # type: ignore
# Affirm to Mypy that this is indeed a method on this mock
destination.check.assert_called_with(
logger=ANY, config=dummy_config) # type: ignore
# Check if config validation has been called
validate_mock.assert_called_with(dummy_config, spec_msg,
destination.logger)
# verify output was correct
assert _wrapped(expected_check_result) == returned_check_result
def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
try:
connection = create_connection(config=config)
except Exception as e:
logger.error(f"Failed to create connection. Error: {e}")
return AirbyteConnectionStatus(status=Status.FAILED, message=f"Could not create connection: {repr(e)}")
try:
channel = connection.channel()
if channel.is_open:
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
return AirbyteConnectionStatus(status=Status.FAILED, message="Could not open channel")
except Exception as e:
logger.error(f"Failed to open RabbitMQ channel. Error: {e}")
return AirbyteConnectionStatus(status=Status.FAILED, message=f"An exception occurred: {repr(e)}")
finally:
connection.close()
def check(self, logger: AirbyteLogger,
config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the destination with the needed permissions
e.g: if a provided API token or password can be used to connect and write to the destination.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this destination, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
connector_config = ConnectorConfig(**config)
try:
aws_handler = AwsHandler(connector_config, self)
except (ClientError, AttributeError) as e:
logger.error(
f"""Could not create session on {connector_config.aws_account_id} Exception: {repr(e)}"""
)
message = f"""Could not authenticate using {connector_config.credentials_type} on Account {connector_config.aws_account_id} Exception: {repr(e)}"""
return AirbyteConnectionStatus(status=Status.FAILED,
message=message)
try:
aws_handler.head_bucket()
except ClientError as e:
message = f"""Could not find bucket {connector_config.bucket_name} in aws://{connector_config.aws_account_id}:{connector_config.region} Exception: {repr(e)}"""
return AirbyteConnectionStatus(status=Status.FAILED,
message=message)
with LakeformationTransaction(aws_handler) as tx:
table_location = "s3://" + connector_config.bucket_name + "/" + connector_config.bucket_prefix + "/" + "airbyte_test/"
table = aws_handler.get_table(
txid=tx.txid,
database_name=connector_config.lakeformation_database_name,
table_name="airbyte_test",
location=table_location,
)
if table is None:
message = f"Could not create a table in database {connector_config.lakeformation_database_name}"
return AirbyteConnectionStatus(status=Status.FAILED,
message=message)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def test_run_check(entrypoint: AirbyteEntrypoint, mocker):
parsed_args = Namespace(command="check", config="config_path")
config = {"username": "******"}
check_value = AirbyteConnectionStatus(status=Status.SUCCEEDED)
mocker.patch.object(MockSource, "read_config", return_value=config)
mocker.patch.object(MockSource, "configure", return_value=config)
mocker.patch.object(MockSource, "check", return_value=check_value)
assert [_wrap_message(check_value)] == list(entrypoint.run(parsed_args))
def check_config(self, logger, config_path: str,
config: json) -> AirbyteConnectionStatus:
try:
repositories = config["repository"].split(" ")
for repository in repositories:
url = f"https://api.github.com/repos/{repository}/commits"
response = requests.get(url, auth=(config["access_token"], ""))
if response.status_code != requests.codes.ok:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"{repository} {response.text}")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as err:
logger.error(err)
error_msg = f"Unable to connect with the provided credentials. Error: {err}"
return AirbyteConnectionStatus(status=Status.FAILED,
message=error_msg)
def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the destination with the needed permissions
e.g: if a provided API token or password can be used to connect and write to the destination.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this destination, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
writer = FirestoreWriter(**config)
writer.check()
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(status=Status.FAILED, message=f"An exception occurred: {repr(e)}")
def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the destination with the needed permissions
e.g: if a provided API token or password can be used to connect and write to the destination.
"""
try:
# Verify write access by attempting to write and then delete to a random key
client = KvDbClient(**config)
random_key = str(uuid.uuid4())
client.write(random_key, {"value": "_airbyte_connection_check"})
client.delete(random_key)
except Exception as e:
traceback.print_exc()
return AirbyteConnectionStatus(
status=Status.FAILED, message=f"An exception occurred: {e}. \nStacktrace: \n{traceback.format_exc()}"
)
else:
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def test_check_connection(mocker):
source = SourceGoogleAnalyticsDataApi()
report_mock = MagicMock()
mocker.patch.object(SourceGoogleAnalyticsDataApi, "_run_report", return_value=report_mock)
logger_mock = MagicMock()
config_mock = MagicMock()
assert source.check(logger_mock, config_mock) == AirbyteConnectionStatus(status=Status.SUCCEEDED)
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
bamboo = BambooHrClient(config)
bamboo.request("employees/directory")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
try:
access_token = config["access_token"]
spreadsheet_id = config["spreadsheet_id"]
smartsheet_client = smartsheet.Smartsheet(access_token)
smartsheet_client.errors_as_exceptions(True)
smartsheet_client.Sheets.get_sheet(spreadsheet_id)
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
if isinstance(e, smartsheet.exceptions.ApiError):
err = e.error.result
code = 404 if err.code == 1006 else err.code
reason = f"{err.name}: {code} - {err.message} | Check your spreadsheet ID."
else:
reason = str(e)
logger.error(reason)
return AirbyteConnectionStatus(status=Status.FAILED)
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration
e.g: if a provided Stripe API token can be used to connect to the Stripe API.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
# Not Implemented
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
def check(self, logger: AirbyteLogger,
config: json) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration
e.g: if a provided Stripe API token can be used to connect to the Stripe API.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
try:
with establish_connection(config, logger) as connection:
# We can only verify correctness of connection parameters on execution
with connection.cursor() as cursor:
cursor.execute("SELECT 1")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED,
message=f"An exception occurred: {str(e)}")
def check(self, logger: AirbyteLogger, config: json) -> AirbyteConnectionStatus:
try:
if "max_batch_size" in config:
# Max batch size must be between 1 and 10
if config["max_batch_size"] > 10 or config["max_batch_size"] < 1:
raise Exception("max_batch_size must be between 1 and 10")
if "max_wait_time" in config:
# Max wait time must be between 1 and 20
if config["max_wait_time"] > 20 or config["max_wait_time"] < 1:
raise Exception("max_wait_time must be between 1 and 20")
# Required propeties
queue_url = config["queue_url"]
logger.debug("Amazon SQS Source Config Check - queue_url: " + queue_url)
queue_region = config["region"]
logger.debug("Amazon SQS Source Config Check - region: " + queue_region)
# Senstive Properties
access_key = config["access_key"]
logger.debug("Amazon SQS Source Config Check - access_key (ends with): " + access_key[-1])
secret_key = config["secret_key"]
logger.debug("Amazon SQS Source Config Check - secret_key (ends with): " + secret_key[-1])
logger.debug("Amazon SQS Source Config Check - Starting connection test ---")
session = boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name=queue_region)
sqs = session.resource("sqs")
queue = sqs.Queue(url=queue_url)
if hasattr(queue, "attributes"):
logger.debug("Amazon SQS Source Config Check - Connection test successful ---")
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
else:
return AirbyteConnectionStatus(status=Status.FAILED, message="Amazon SQS Source Config Check - Could not connect to queue")
except ClientError as e:
return AirbyteConnectionStatus(status=Status.FAILED, message=f"Amazon SQS Source Config Check - Error in AWS Client: {str(e)}")
except Exception as e:
return AirbyteConnectionStatus(
status=Status.FAILED, message=f"Amazon SQS Source Config Check - An exception occurred: {str(e)}"
)
def check(self, logger: AirbyteLogger, config: Dict[str, any]) -> AirbyteConnectionStatus:
"""
Tests if the input configuration can be used to successfully connect to the integration
e.g: if a provided Stripe API token can be used to connect to the Stripe API.
:param logger: Logging object to display debug/info/error to the logs
(logs will not be accessible via airbyte UI if they are not passed to this logger)
:param config: Json object containing the configuration of this source, content of this json is as specified in
the properties of the spec.json file
:return: AirbyteConnectionStatus indicating a Success or Failure
"""
# As this is an in-memory source, it always succeeds
return AirbyteConnectionStatus(status=Status.SUCCEEDED)
def test_config_validate(entrypoint: AirbyteEntrypoint, mocker, config_mock,
schema, config_valid):
parsed_args = Namespace(command="check", config="config_path")
check_value = AirbyteConnectionStatus(status=Status.SUCCEEDED)
mocker.patch.object(MockSource, "check", return_value=check_value)
mocker.patch.object(
MockSource,
"spec",
return_value=ConnectorSpecification(connectionSpecification=schema))
if config_valid:
messages = list(entrypoint.run(parsed_args))
assert [_wrap_message(check_value)] == messages
else:
with pytest.raises(Exception, match=r"(?i)Config Validation Error:.*"):
list(entrypoint.run(parsed_args))
def test_source_wrong_credentials():
source = SourceAmazonSellerPartner()
status = source.check(
logger=AirbyteLogger(),
config={
"start_date": "2021-05-27",
"refresh_token": "ABC",
"lwa_app_id": "lwa_app_id",
"lwa_client_secret": "lwa_client_secret",
"aws_access_key": "aws_access_key",
"aws_secret_key": "aws_secret_key",
"role_arn": "role_arn",
"marketplace": "USA",
},
)
assert status == AirbyteConnectionStatus(
status=Status.FAILED, message="An exception occurred: ('invalid_client', 'Client authentication failed', 401)"
)
def test_check_v2(test_config_v2):
responses.add(
responses.GET,
"https://airbyte-test.chargebee.com/api/v2/subscriptions",
json={
"list": [{
"subscription": {
"id": "cbdemo_cancelled_sub"
},
"customer": {},
"card": {}
}]
},
)
source = SourceChargebee()
logger_mock = MagicMock()
assert source.check(
logger_mock,
test_config_v2) == AirbyteConnectionStatus(status=Status.SUCCEEDED)
assert len(responses.calls) == 1