def get_connection(self, context):
"""
Establish a connection to the Postgres server
:param context: execution context
:return: server object or None if unable to connect
:rtype: PostgresDb
"""
if isinstance(self._postgres_cfg, str):
client = PostgresDb(cfg_filename=self._postgres_cfg)
elif isinstance(self._postgres_cfg, dict):
client = PostgresDb(cfg_dict=self._postgres_cfg)
else:
raise Failure(f'No configuration provided for PostgresWarehouse')
server = client['host']
if client.get_connection() is not None:
context.log.info(f'Connected to Postgres: {server}')
else:
context.log.info(f'Unable to connect to Postgres: {server}')
client.close_connection()
if self._fatal:
raise Failure(f'Unable to connect to Postgres: {server}')
client = None
return client
def get_connection(self, context):
"""
Establish a connection to the mongoDB server
:param context: execution context
:return: server object or None if unable to connect
:rtype: MongoDb
"""
if isinstance(self._mongo_cfg, str):
client = MongoDb(cfg_filename=self._mongo_cfg)
elif isinstance(self._mongo_cfg, dict):
client = MongoDb(cfg_dict=self._mongo_cfg)
else:
raise Failure(f'No configuration provided for MongoWarehouse')
server = client['server']
if client.is_authenticated():
context.log.info(f'Connected to mongoDB: {server}')
else:
context.log.info(f'Unable to connect to mongoDB: {server}')
client.close_connection()
if self._fatal:
raise Failure(f'Unable to connect to mongoDB: {server}')
client = None
return client
def less_simple_data_frame_type_check(value):
if not isinstance(value, list):
raise Failure(
'LessSimpleDataFrame should be a list of dicts, got '
'{type_}'.format(type_=type(value))
)
fields = [field for field in value[0].keys()]
for i in range(len(value)):
row = value[i]
if not isinstance(row, dict):
raise Failure(
'LessSimpleDataFrame should be a list of dicts, '
'got {type_} for row {idx}'.format(
type_=type(row), idx=(i + 1)
)
)
row_fields = [field for field in row.keys()]
if fields != row_fields:
raise Failure(
'Rows in LessSimpleDataFrame should have the same fields, '
'got {actual} for row {idx}, expected {expected}'.format(
actual=row_fields, idx=(i + 1), expected=fields
)
)
def poll_sync(
self,
connector_id: str,
initial_last_sync_completion: datetime.datetime,
poll_interval: float = DEFAULT_POLL_INTERVAL,
poll_timeout: float = None,
) -> Dict[str, Any]:
"""
Given a Fivetran connector and the timestamp at which the previous sync completed, poll
until the next sync completes.
The previous sync completion time is necessary because the only way to tell when a sync
completes is when this value changes.
Args:
connector_id (str): The Fivetran Connector ID. You can retrieve this value from the
"Setup" tab of a given connector in the Fivetran UI.
initial_last_sync_completion (datetime.datetime): The timestamp of the last completed sync
(successful or otherwise) for this connector, prior to running this method.
poll_interval (float): The time (in seconds) that will be waited between successive polls.
poll_timeout (float): The maximum time that will waited before this operation is timed
out. By default, this will never time out.
Returns:
Dict[str, Any]: Parsed json data representing the API response.
"""
poll_start = datetime.datetime.now()
while True:
(
curr_last_sync_completion,
curr_last_sync_succeeded,
curr_sync_state,
) = self.get_connector_sync_status(connector_id)
self._log.info(f"Polled '{connector_id}'. Status: [{curr_sync_state}]")
if curr_last_sync_completion > initial_last_sync_completion:
break
if poll_timeout and datetime.datetime.now() > poll_start + datetime.timedelta(
seconds=poll_timeout
):
raise Failure(
f"Sync for connector '{connector_id}' timed out after {datetime.datetime.now() - poll_start}."
)
# Sleep for the configured time interval before polling again.
time.sleep(poll_interval)
connector_details = self.get_connector_details(connector_id)
if not curr_last_sync_succeeded:
raise Failure(
f"Sync for connector '{connector_id}' failed!",
metadata={
"connector_details": MetadataValue.json(connector_details),
"log_url": MetadataValue.url(get_fivetran_logs_url(connector_details)),
},
)
return connector_details
def file_exists_at_path_type_check(value):
if not isinstance(value, six.string_types):
raise Failure(
'FileExistsAtPath must be a string in memory. Got {value}'.format(
value=repr(value)))
if not safe_isfile(value):
raise Failure(
('FileExistsAtPath must be a path that points to a file that '
'exists. "{value}" does not exist on disk').format(value=value))
def raise_for_rpc_error(context: SolidExecutionContext,
resp: Response) -> None:
error = resp.json().get("error")
if error is not None:
if error["code"] in [
DBTErrors.project_currently_compiling_error,
DBTErrors.runtime_error,
DBTErrors.server_error,
]:
context.log.warning(error["message"])
raise RetryRequested(max_retries=5,
seconds_to_wait=30) # TODO backoff logic
elif error["code"] == DBTErrors.project_compile_failure_error:
raise Failure(
description=error["message"],
metadata_entries=[
EventMetadataEntry.text(text=str(error["code"]),
label="RPC Error Code"),
EventMetadataEntry.text(
text=error["data"]["cause"]["message"],
label="RPC Error Cause"),
],
)
elif error["code"] == DBTErrors.rpc_process_killed_error:
raise Failure(
description=error["message"],
metadata_entries=[
EventMetadataEntry.text(text=str(error["code"]),
label="RPC Error Code"),
EventMetadataEntry.text(text=str(error["data"]["signum"]),
label="RPC Signum"),
EventMetadataEntry.text(text=error["data"]["message"],
label="RPC Error Message"),
],
)
elif error["code"] == DBTErrors.rpc_timeout_error:
raise Failure(
description=error["message"],
metadata_entries=[
EventMetadataEntry.text(text=str(error["code"]),
label="RPC Error Code"),
EventMetadataEntry.text(text=str(error["data"]["timeout"]),
label="RPC Timeout"),
EventMetadataEntry.text(text=error["data"]["message"],
label="RPC Error Message"),
],
)
else:
raise Failure(
description=error["message"],
metadata_entries=[
EventMetadataEntry.text(text=str(error["code"]),
label="RPC Error Code"),
],
)
def drop_tables():
"""
Drop postgres tables
"""
drop_sales_table = drop_table.alias('drop_sales_table')
drop_tracking_table = drop_table.alias('drop_tracking_table')
if not drop_sales_table():
raise Failure('Sales table not dropped')
if not drop_tracking_table():
raise Failure('Tracking table not dropped')
def sync_and_poll(
self,
connection_id: str,
poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
poll_timeout: float = None,
):
"""
Initializes a sync operation for the given connector, and polls until it completes.
Args:
connection_id (str): The Airbyte Connector ID. You can retrieve this value from the
"Connection" tab of a given connection in the Arbyte UI.
poll_interval (float): The time (in seconds) that will be waited between successive polls.
poll_timeout (float): The maximum time that will waited before this operation is timed
out. By default, this will never time out.
Returns:
:py:class:`~AirbyteOutput`:
Details of the sync job.
"""
job_details = self.start_sync(connection_id)
job_id = job_details.get("job", {}).get("id")
self._log.info(
f"Job {job_id} initialized for connection_id={connection_id}.")
start = time.monotonic()
while True:
if poll_timeout and start + poll_timeout < time.monotonic():
raise Failure(
f"Timeout: Airbyte job {job_id} is not ready after the timeout {poll_timeout} seconds"
)
time.sleep(poll_interval)
job_details = self.get_job_status(job_id)
state = job_details.get("job", {}).get("status")
if state in (AirbyteState.RUNNING, AirbyteState.PENDING,
AirbyteState.INCOMPLETE):
continue
elif state == AirbyteState.SUCCEEDED:
break
elif state == AirbyteState.ERROR:
raise Failure(f"Job failed: {job_id}")
elif state == AirbyteState.CANCELLED:
raise Failure(f"Job was cancelled: {job_id}")
else:
raise Failure(
f"Encountered unexpected state `{state}` for job_id {job_id}"
)
return AirbyteOutput(job_details=job_details.get("job", {}))
def _assert_syncable_connector(self, connector_id: str):
"""
Confirms that a given connector is eligible to sync. Will raise a Failure in the event that
the connector is either paused or not fully setup.
Args:
connector_id (str): The Fivetran Connector ID. You can retrieve this value from the
"Setup" tab of a given connector in the Fivetran UI.
"""
connector_details = self.get_connector_details(connector_id)
if connector_details["paused"]:
raise Failure("Connector '{connector_id}' cannot be synced as it is currently paused.")
if connector_details["status"]["setup_state"] != "connected":
raise Failure("Connector '{connector_id}' cannot be synced as it has not been setup")
def positive_num_check(_, value):
if value > 0:
return True
else:
raise Failure(
"Numbers cannot be 0 or negative, got {num} for PositiveNumber type"
.format(num=value))
def dbt_test_solid(context):
executable_path = context.solid_config['dbt_executable']
if not distutils.spawn.find_executable(executable_path):
raise Failure(
'Could not find dbt executable at "{executable_path}". Please ensure that '
'dbt is installed and on the PATH.'.format(
executable_path=executable_path))
cmd = '{executable_path} test --project-dir {project_dir}'.format(
executable_path=executable_path, project_dir=project_dir)
if profiles_dir:
cmd += ' --profiles-dir {}'.format(profiles_dir)
args = shlex.split(cmd)
proc = subprocess.Popen(args, stdout=subprocess.PIPE)
for line in io.TextIOWrapper(proc.stdout, encoding='utf-8'):
text = line.rstrip()
if not text:
continue
# print to stdout
print(text)
# remove colors
text = ANSI_ESCAPE.sub('', text)
expt = try_parse_test(text)
if expt:
yield expt
yield Output(value=None, output_name='test_complete')
def throw():
raise Failure(
description="it Failure",
metadata_entries=[
EventMetadataEntry.text(label="label", text="text", description="description")
],
)
def generate_currency_table_fields_str(context, table_data_columns: Dict):
"""
Upload a DataFrame to the Postgres server, creating the table if it doesn't exist
:param context: execution context
:param table_data_columns: dict containing list of column names and definitions for the database table
"""
# TODO basically same as generate_tracking_table_fields_str, so make a generic solid
names = table_data_columns['value']['names']
defs = table_data_columns['value']['defs']
auto = table_data_columns['value']['auto']
if len(names) != len(defs) or len(names) != len(auto):
raise Failure(
f'Configuration error: column definition counts do not match: names ({len(names)}), '
f'definitions ({len(defs)}), auto ({len(auto)})')
# add fields from the table description
create_columns = ''
insert_columns = ''
for idx in range(len(names)):
if len(create_columns) > 0:
create_columns += ', '
create_columns += f"{names[idx]} {defs[idx]}"
if not auto[idx]:
if len(insert_columns) > 0:
insert_columns += ', '
insert_columns += f"{names[idx]}"
yield Output(create_columns, 'create_currency_columns')
yield Output(insert_columns, 'insert_currency_columns')
def _poll_rpc(context: SolidExecutionContext,
request_token: str,
should_yield_materializations: bool = True) -> DbtRpcOutput:
"""Polls the dbt RPC server for the status of a request until the state is ``success``."""
from ..utils import generate_materializations
logs_start = 0
interval = context.solid_config.get("interval")
elapsed_time = -1
current_state = None
while True:
# Poll for the dbt RPC request.
context.log.debug(f"RequestToken: {request_token}")
resp = context.resources.dbt_rpc.poll(
request_token=request_token,
logs=context.solid_config["logs"],
logs_start=logs_start)
raise_for_rpc_error(context, resp)
resp_json = resp.json()
resp_result = resp_json.get("result", {})
# Pass dbt RPC logs into the Dagster/Dagit logger.
if context.solid_config["logs"]:
logs = resp_result.get("logs", [])
if len(logs) > 0:
log_rpc(context, logs)
logs_start += len(logs)
current_state = resp_result.get("state")
# Stop polling if request's state is no longer "running".
if current_state != "running":
break
elapsed_time = resp_result.get("elapsed", 0)
# Sleep for the configured time interval before polling again.
context.log.debug(
f"Request {request_token} currently in state '{current_state}' (elapsed time "
f"{elapsed_time} seconds). Sleeping for {interval}s...")
time.sleep(interval)
if current_state != "success":
raise Failure(description=(
f"Request {request_token} finished with state '{current_state}' in "
f"{elapsed_time} seconds"), )
context.log.info(
f"Request {request_token} finished with state '{current_state}' in {elapsed_time} seconds"
)
context.log.debug(json.dumps(resp_result, indent=2))
polled_run_results = DbtRpcOutput.from_dict(resp_result)
if should_yield_materializations:
for materialization in generate_materializations(polled_run_results):
yield materialization
yield Output(polled_run_results)
def dbt_solid(_):
cmd = 'dbt run --project-dir {}'.format(project_dir)
if profiles_dir:
cmd += ' --profiles-dir {}'.format(profiles_dir)
args = shlex.split(cmd)
proc = subprocess.Popen(args, stdout=subprocess.PIPE)
# if https://github.com/fishtown-analytics/dbt/issues/1237 gets done
# we should definitely switch to parsing the json output, as that
# would be much more reliable/resilient
for line in io.TextIOWrapper(proc.stdout, encoding='utf-8'):
text = line.rstrip()
if not text:
continue
# print to stdout
print(text)
# remove colors
text = ANSI_ESCAPE.sub('', text)
mat = try_parse_run(text)
if mat:
yield mat
proc.wait()
if proc.returncode != 0:
raise Failure('Dbt invocation errored')
yield Output(value=None, output_name='run_complete')
def throws_failure():
raise Failure(
description='Always fails.',
metadata_entries=[
EventMetadataEntry.text('why', label='always_fails')
],
)
def collect_releases(context):
ocds_loader: OCDSLoader = context.resources.ocds_loader
last_loading = ocds_loader.get_last_data_loading_records(
n_records=1).iloc[0]
app_dir = os.path.dirname(os.path.abspath(__file__))
script_path = os.path.join(app_dir, 'shell', 'kingfisher-collect.sh')
from_date = context.solid_config.get('from_date')
last_item_date = last_loading.last_item_date
start_date = from_date or last_item_date.strftime('%Y-%m-%d')
shell_command = 'bash {} {}'.format(script_path, start_date)
output, return_code = execute(shell_command=shell_command,
log=context.log,
**without_keys(context.solid_config,
['from_date']))
if return_code:
raise Failure(
description='Shell command execution failed with output: {output}'.
format(output=output))
return output
def make_request(self, endpoint: str, data: Optional[Dict[str, Any]]):
"""
Creates and sends a request to the desired Airbyte REST API endpoint.
Args:
endpoint (str): The Airbyte API endpoint to send this request to.
data (Optional[str]): JSON-formatted data string to be included in the request.
Returns:
Dict[str, Any]: Parsed json data from the response to this request
"""
headers = {"accept": "application/json"}
num_retries = 0
while True:
try:
response = requests.request(
method="POST",
url=self.api_base_url + endpoint,
headers=headers,
json=data,
)
response.raise_for_status()
return response.json()
except RequestException as e:
self._log.error("Request to Airbyte API failed: %s", e)
if num_retries == self._request_max_retries:
break
num_retries += 1
time.sleep(self._request_retry_delay)
raise Failure("Exceeded max number of retries.")
def should_fail(_, _obj):
raise Failure(
description="Foolure",
metadata_entries=[
EventMetadataEntry.text(label="label", text="text", description="description")
],
)
def drop_currency_tables():
"""
Drop postgres tables
"""
drop_currency_table = drop_table.alias('drop_currency_table')
if not drop_currency_table():
raise Failure('Currency table not dropped')
def type_check(self, value):
if not isinstance(value, dict):
raise Failure(
'Value {value} should be of type {type_name}.'.format(
value=value, type_name=self.name))
for key in value:
if not key in permitted_key_names:
raise Failure(
'Key {name} is not a permitted value, values can only be of: {name_list}'
.format(name=value.name,
name_list=permitted_key_names))
return TypeCheck(metadata_entries=[
EventMetadataEntry.text(label='row_count',
text=str(len(value))),
EventMetadataEntry.text(label='series_names',
text=', '.join(value.keys())),
])
def retry_op(context):
if context.retry_number < 3:
if is_explicit:
raise Failure(description="some failure description",
metadata={"foo": 1.23})
else:
_ = "x" + 1
return context.retry_number
def _execute_item(self, context, item_name, forward_resource_stacks,
backward_resources):
"""Executes the given item using the given forward resource stacks and backward resources.
Returns list of output resource stacks.
Called by ``_make_forward_solid_def.compute_fn``.
For each element yielded by ``_filtered_resources_iterator``, spawns a thread that runs ``_execute_item_filtered``.
Args:
context
item_name (str)
forward_resource_stacks (list(tuple(ProjectItemResource)))
backward_resources (list(ProjectItemResource))
Returns:
list(tuple(ProjectItemResource))
"""
success = [ItemExecutionFinishState.NEVER_FINISHED]
output_resources_list = []
threads = []
resources_iterator = self._filtered_resources_iterator(
item_name, forward_resource_stacks, backward_resources,
create_timestamp())
for flt_fwd_resources, flt_bwd_resources, filter_id in resources_iterator:
item = self._make_item(item_name, ED.FORWARD)
if not item.ready_to_execute(self._settings):
if not self._execution_permits[self._solid_names[
item_name]]: # Exclude if not selected
success[0] = ItemExecutionFinishState.EXCLUDED
else: # Fail if selected
context.log.error(
f"compute_fn() FAILURE in: '{item_name}', not ready for forward execution"
)
success[0] = ItemExecutionFinishState.FAILURE
else:
item.filter_id = filter_id
thread = threading.Thread(
target=self._execute_item_filtered,
args=(item, flt_fwd_resources, flt_bwd_resources,
output_resources_list, success),
)
threads.append(thread)
thread.start()
for thread in threads:
thread.join()
if success[0] == ItemExecutionFinishState.FAILURE:
context.log.error(
f"compute_fn() FAILURE with item: {item_name} failed to execute"
)
raise Failure()
for resources in output_resources_list:
for connection in self._connections_by_source.get(item_name, []):
connection.receive_resources_from_source(resources)
if connection.has_filters():
connection.fetch_database_items()
return list(zip(*output_resources_list)), success[0]
def dbt_rpc_poll(
context,
request_token: str,
should_yield_materializations: bool = True) -> DbtRpcPollResult:
resp = context.resources.dbt_rpc.poll(request_token=request_token,
logs=context.solid_config["logs"])
raise_for_rpc_error(context, resp)
interval = context.solid_config["interval"]
logs_start = 0
while resp.json().get("result").get("state") == "running":
context.log.debug(
f"Request {request_token} currently in state '{resp.json().get('result').get('state')}' (elapsed time {resp.json().get('result').get('elapsed', 0)} seconds). Sleeping for {interval}s.."
)
if context.solid_config["logs"]:
logs = resp.json().get("result").get("logs")
if len(logs) > 0:
log_rpc(context, logs)
logs_start += len(logs)
time.sleep(interval)
resp = context.resources.dbt_rpc.poll(
request_token=request_token,
logs=context.solid_config["logs"],
logs_start=logs_start)
raise_for_rpc_error(context, resp)
if resp.json().get("result").get("state") != "success":
if context.solid_config["logs"]:
logs = resp.json().get("result").get("logs")
if len(logs) > 0:
log_rpc(context, logs)
raise Failure(
description=
f"Request {request_token} finished with state '{resp.json().get('result').get('state')}' in {resp.json().get('result').get('elapsed')} seconds",
)
else:
context.log.info(
f"Request {request_token} finished with state '{resp.json().get('result').get('state')}' in {resp.json().get('result').get('elapsed')} seconds"
)
if context.solid_config["logs"]:
logs = resp.json().get("result").get("logs")
if len(logs) > 0:
log_rpc(context, logs)
context.log.debug(json.dumps(resp.json().get("result"), indent=2))
rpc_poll_result = DbtRpcPollResult.from_results(resp.json().get("result"))
if should_yield_materializations:
for materialization in generate_materializations(
rpc_poll_result=rpc_poll_result):
yield materialization
yield Output(value=rpc_poll_result, output_name="result")
def throw():
raise Failure(
description='it Failure',
metadata_entries=[
EventMetadataEntry.text(label='label',
text='text',
description='description')
],
)
def compute_fn(context, inputs):
if self.state() == SpineEngineState.USER_STOPPED:
context.log.error(
f"compute_fn() FAILURE with item: {item_name} stopped by the user"
)
raise Failure()
context.log.info(f"Item Name: {item_name}")
item = self._make_item(item_name, ED.BACKWARD)
resources = item.output_resources(ED.BACKWARD)
yield Output(value=resources, output_name=f"{ED.BACKWARD}_output")
def dbt_solid(context):
executable_path = context.solid_config['dbt_executable']
if not distutils.spawn.find_executable(executable_path):
raise Failure(
'Could not find dbt executable at "{executable_path}". Please ensure that '
'dbt is installed and on the PATH.'.format(
executable_path=executable_path))
cmd = '{executable_path} run --project-dir {project_dir}'.format(
executable_path=executable_path, project_dir=project_dir)
if profiles_dir:
cmd += ' --profiles-dir {profiles_dir}'.format(
profiles_dir=profiles_dir)
args = shlex.split(cmd)
proc = subprocess.Popen(args, stdout=subprocess.PIPE)
# if https://github.com/fishtown-analytics/dbt/issues/1237 gets done
# we should definitely switch to parsing the json output, as that
# would be much more reliable/resilient
for line in io.TextIOWrapper(proc.stdout, encoding='utf-8'):
text = line.rstrip()
if not text:
continue
# print to stdout
print(text)
# remove colors
text = ANSI_ESCAPE.sub('', text)
mat = try_parse_run(text)
if mat:
yield mat
proc.wait()
if proc.returncode != 0:
raise Failure('Dbt invocation errored.')
yield Output(value=None, output_name='run_complete')
def raise_for_rpc_error(context: SolidExecutionContext,
resp: Response) -> None:
error = resp.json().get("error")
if error is not None:
if error["code"] in [
DBTErrors.project_currently_compiling_error.value,
DBTErrors.runtime_error.value,
DBTErrors.server_error.value,
]:
context.log.warning(error["message"])
raise RetryRequested(max_retries=5, seconds_to_wait=30)
elif error["code"] == DBTErrors.project_compile_failure_error.value:
raise Failure(
description=error["message"],
metadata={
"RPC Error Code": str(error["code"]),
"RPC Error Cause": error["data"]["cause"]["message"],
},
)
elif error["code"] == DBTErrors.rpc_process_killed_error.value:
raise Failure(
description=error["message"],
metadata={
"RPC Error Code": str(error["code"]),
"RPC Signum": str(error["data"]["signum"]),
"RPC Error Message": error["data"]["message"],
},
)
elif error["code"] == DBTErrors.rpc_timeout_error.value:
raise Failure(
description=error["message"],
metadata={
"RPC Error Code": str(error["code"]),
"RPC Timeout": str(error["data"]["timeout"]),
"RPC Error Message": error["data"]["message"],
},
)
else:
raise Failure(
description=error["message"],
metadata={"RPC Error Code": str(error["code"])},
)
def _shell_solid(context):
output, return_code = execute(shell_command=shell_command,
log=context.log,
**context.solid_config)
if return_code:
raise Failure(
description=
"Shell command execution failed with output: {output}".format(
output=output))
return output
def _bash_solid(context):
output, return_code = execute(bash_command=bash_command,
log=context.log,
**context.solid_config)
if return_code:
raise Failure(
description=
'Bash command execution failed with output: {output}'.format(
output=output))
return output
