diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1fc7838..3011d45 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,7 +7,7 @@ repos: - id: isort stages: [commit,push] name: isort - entry: poetry run isort -rc + entry: poetry run isort language: system types: [python] - id: black diff --git a/dbtc/cli.py b/dbtc/cli.py index f629f3f..4d10aaf 100644 --- a/dbtc/cli.py +++ b/dbtc/cli.py @@ -1,4 +1,5 @@ # stdlib +from enum import auto import json from typing import List, Optional @@ -255,11 +256,10 @@ def create_environment_variables( def create_job( ctx: typer.Context, account_id: int = ACCOUNT_ID, - project_id: int = PROJECT_ID, payload: str = PAYLOAD, ): """Create a job in a project.""" - _dbt_cloud_request(ctx, 'create_job', account_id, project_id, json.loads(payload)) + _dbt_cloud_request(ctx, 'create_job', account_id, json.loads(payload)) @app.command() @@ -969,6 +969,74 @@ def test_connection( ) +@app.command() +def trigger_job_with_autoscaling( + ctx: typer.Context, + account_id: int = ACCOUNT_ID, + job_id: int = JOB_ID, + payload: str = PAYLOAD, + should_poll: bool = typer.Option( + True, + help='Poll until job completion (status is one of success, failure, or ' + 'cancelled)', + ), + poll_interval: int = typer.Option( + 10, '--poll-interval', help='Number of seconds to wait in between polling.' + ), + autoscale_delete_post_run: bool = typer.Option( + False, help='Delete the cloned job immediately after it runs' + ), + autoscale_job_identifier: str = typer.Option( + False, + help=( + 'A string to append to the name of the job with id specified in job_id ' + 'E.g. if the job is named "my_dbt_cloud_job" and the autoscale_job_identifier ' + 'is "my_new_job", the cloned job will be named "my_dbt_cloud_job-my_new_job" ' + ), + ), +): + """Trigger job to run.""" + _dbt_cloud_request( + ctx, + 'trigger_job_with_autoscaling', + account_id, + job_id, + json.loads(payload), + should_poll=should_poll, + poll_interval=poll_interval, + autoscale_delete_post_run=autoscale_delete_post_run, + autoscale_job_identifier=autoscale_job_identifier, + ) + + +@app.command() +def trigger_job_restart_from_failure( + ctx: typer.Context, + account_id: int = ACCOUNT_ID, + job_id: int = JOB_ID, + payload: str = PAYLOAD, + should_poll: bool = typer.Option( + True, + help='Poll until job completion (status is one of success, failure, or ' + 'cancelled)', + ), + poll_interval: int = typer.Option( + 10, '--poll-interval', help='Number of seconds to wait in between polling.' + ) +): + """Trigger job to rerun from the point of failure on its most recent run + Does nothing if no failures are identified on the most recent run.""" + _dbt_cloud_request( + ctx, + 'trigger_job_restart_from_failure', + account_id, + job_id, + json.loads(payload), + should_poll=should_poll, + poll_interval=poll_interval, + ) + + @app.command() def trigger_job( ctx: typer.Context, diff --git a/dbtc/client/cloud/base.py b/dbtc/client/cloud/base.py index 0226dc9..47d3524 100644 --- a/dbtc/client/cloud/base.py +++ b/dbtc/client/cloud/base.py @@ -1,42 +1,24 @@ # stdlib import argparse -import enum +import json import shlex import time +from datetime import datetime from functools import partial, wraps -from typing import Dict, Iterable, List +from typing import Dict, Iterable, List, Union # third party import requests # first party from dbtc.client.base import _Client - - -class JobRunStatus(enum.IntEnum): - QUEUED = 1 - STARTING = 2 - RUNNING = 3 - SUCCESS = 10 - ERROR = 20 - CANCELLED = 30 - - -RUN_COMMANDS = ['build', 'run', 'test', 'seed', 'snapshot'] -GLOBAL_CLI_ARGS = { - 'warn_error': {'flags': ('--warn-error',), 'action': 'store_true'}, - 'use_experimental_parser': { - 'flags': ('--use-experimental-parser',), - 'action': 'store_true', - }, -} -SUB_COMMAND_CLI_ARGS = { - 'vars': {'flags': ('--vars',)}, - 'args': {'flags': ('--args',)}, - 'fail_fast': {'flags': ('-x', '--fail-fast'), 'action': 'store_true'}, - 'full_refresh': {'flags': ('--full-refresh',), 'action': 'store_true'}, - 'store_failures': {'flags': ('--store-failures',), 'action': 'store_true'}, -} +from dbtc.client.cloud import models +from dbtc.client.cloud.configs.dbt_core_cli import ( + global_cli_args, + run_commands, + sub_command_cli_args, +) +from dbtc.client.cloud.configs.enums import JobRunStatus def _version_decorator(func, version): @@ -49,6 +31,7 @@ def wrapper(self, *args, **kwargs): return wrapper +# Version Decorators v2 = partial(_version_decorator, version='v2') v3 = partial(_version_decorator, version='v3') v4 = partial(_version_decorator, version='v4') @@ -60,7 +43,7 @@ def __init__(self, **kwargs): self.session = requests.Session() self.session.headers = self.headers self.parser = argparse.ArgumentParser() - all_cli_args = {**GLOBAL_CLI_ARGS, **SUB_COMMAND_CLI_ARGS} + all_cli_args = {**global_cli_args, **sub_command_cli_args} for arg_specs in all_cli_args.values(): flags = arg_specs['flags'] self.parser.add_argument( @@ -77,10 +60,34 @@ def _header_property(self): return 'api_key' + def _clone_resource(self, resource: str, account_id: int, **kwargs): + payload = getattr(self, f'get_{resource}')(account_id=account_id, **kwargs)['data'] + resource_fields = getattr(models, resource.capitalize()).__fields__ + + for k in list(payload): + # only map the fields we're aware of + if k not in resource_fields: + payload.pop(k, None) + + # for optional fields, don't copy them if none in the source payload + elif not resource_fields[k].required and payload[k] is None: + payload.pop(k, None) + + # Can't recreate a resource with an ID + payload['id'] = None + return payload + def _make_request( self, path: str, *, method: str = 'get', **kwargs ) -> requests.Response: """Make request to API.""" + + # Model is not an argument that the request method accepts, needs to be removed + model = kwargs.pop('model', None) + if model is not None: + + # This will validate the payload as well as add any optional fields + kwargs['json'] = model(**kwargs['json']).dict() full_url = self.full_url(path) response = self.session.request(method=method, url=full_url, **kwargs) return response @@ -194,6 +201,27 @@ def cancel_run(self, account_id: int, run_id: int) -> Dict: method='post', ) + @v2 + def clone_job( + self, + account_id: int, + job_id: int, + ): + + """Create a job using the configuration of another + + !!! tip + If a job is currently running, replicate the job definition to a new job, + and trigger + + Args: + account_id (int): Numeric ID of the account to retrieve + job_id (int): Numeric ID of the job to trigger + """ + return self._clone_resource( + 'job', account_id=account_id, job_id=job_id, create_args=['account_id'] + ) + @v3 def create_adapter(self, account_id: int, project_id: int, payload: Dict) -> Dict: """Create an adapter @@ -291,6 +319,7 @@ def create_job(self, account_id: int, payload: Dict) -> Dict: f'accounts/{account_id}/jobs/', method='post', json=payload, + #model=models.Job, ) @v3 @@ -302,7 +331,10 @@ def create_project(self, account_id: int, payload: Dict) -> Dict: payload (dict): Dictionary representing the project to create """ return self._simple_request( - f'accounts/{account_id}/projects/', method='post', json=payload + f'accounts/{account_id}/projects/', + method='post', + json=payload, + model=models.Project, ) @v3 @@ -796,7 +828,11 @@ def list_invited_users(self, account_id: int) -> Dict: @v2 def list_jobs( - self, account_id: int, *, order_by: str = None, project_id: int = None + self, + account_id: int, + *, + order_by: str = None, + project_id: int = None, ) -> Dict: """List jobs in an account or specific project. @@ -865,7 +901,7 @@ def list_runs( order_by: str = None, offset: int = None, limit: int = None, - status: str = None, + status: Union[str, List[str]] = None, ) -> Dict: """List runs in an account. @@ -883,15 +919,21 @@ def list_runs( Use with limit to paginate results. limit (int, optional): The limit to apply when listing runs. Use with offset to paginate results. - status (str, optional): The status to apply when listing runs. + status (str or list, optional): The status to apply when listing runs. Options include queued, starting, running, success, error, and cancelled """ if status is not None: try: - status = getattr(JobRunStatus, status.upper()).value + if isinstance(status, list): + status = [getattr(JobRunStatus, s.upper()).value for s in status] + else: + status = [getattr(JobRunStatus, status.upper()).value] except AttributeError: - pass + raise + else: + status = json.dumps(status) + return self._simple_request( f'accounts/{account_id}/runs', params={ @@ -900,7 +942,7 @@ def list_runs( 'order_by': order_by, 'offset': offset, 'limit': limit, - 'status': status, + 'status__in': status, }, ) @@ -1002,6 +1044,277 @@ def test_connection(self, account_id: int, payload: Dict) -> Dict: f'accounts/{account_id}/connections/test/', method='post', json=payload ) + @v2 + def _get_restart_job_definition( + self, + account_id: int, + job_id: int, + payload: Dict, + ): + + """Identifies whether there was a failure on the previous run of the job. + When failures are identified, returns an updated job definition to + restart from the point of failure. + + Args: + account_id (int): Numeric ID of the account to retrieve + job_id (int): Numeric ID of the job to trigger + payload (dict): Payload required for post request + """ + + def parse_args(cli_args: Iterable[str], namespace: argparse.Namespace): + string = '' + for arg in cli_args: + value = getattr(namespace, arg, None) + if value: + arg = arg.replace('_', '-') + if isinstance(value, bool): + string += f' --{arg}' + else: + string += f" --{arg} '{value}'" + return string + + has_failures = False + + last_run_data = self.list_runs( + account_id=account_id, + include_related=['run_steps'], + job_definition_id=job_id, + order_by='-id', + limit=1, + )['data'][0] + + last_run_status = last_run_data['status_humanized'].lower() + last_run_id = last_run_data['id'] + + if last_run_status == 'error': + rerun_steps = [] + + for run_step in last_run_data['run_steps']: + status = run_step['status_humanized'].lower() + # Skipping cloning, profile setup, and dbt deps - always + # the first three steps in any run + if run_step['index'] <= 3 or status == 'success': + self.console.log( + f'Skipping rerun for command "{run_step["name"]}" ' + 'as it does not need to be repeated.' + ) + + else: + + # get the dbt command used within this step + command = run_step['name'].partition('`')[2].partition('`')[0] + namespace, remaining = self.parser.parse_known_args( + shlex.split(command) + ) + sub_command = remaining[1] + + if ( + sub_command not in run_commands + and status in ['error', 'cancelled', 'skipped'] + ) or (sub_command in run_commands and status == 'skipped'): + rerun_steps.append(command) + + # errors and failures are when we need to inspect to figure + # out the point of failure + else: + + # get the run results scoped to the step which had an error + # an error here indicates that either: + # 1) the fail-fast flag was set, in which case + # the run_results.json file was never created; or + # 2) there was a problem on dbt Cloud's side saving + # this artifact + try: + step_results = self.get_run_artifact( + account_id=account_id, + run_id=last_run_id, + path='run_results.json', + step=run_step['index'], + )['results'] + + # If the artifact isn't found, the API returns a 404 with + # no json. The ValueError will catch the JSONDecodeError + except ValueError: + rerun_steps.append(command) + else: + rerun_nodes = ' '.join( + [ + record['unique_id'].split('.')[2] + for record in step_results + if record['status'] in ['error', 'skipped', 'fail'] + ] + ) + global_args = parse_args(global_cli_args.keys(), namespace) + sub_command_args = parse_args( + sub_command_cli_args.keys(), namespace + ) + modified_command = f'dbt{global_args} {sub_command} -s {rerun_nodes}{sub_command_args}' # noqa: E501 + rerun_steps.append(modified_command) + self.console.log( + f'Modifying command "{command}" as an error ' + 'or failure was encountered.' + ) + if len(rerun_steps) > 0: + has_failures = True + payload.update({"steps_override": rerun_steps}) + + return payload, has_failures + + @v2 + def trigger_job_with_autoscaling( + self, + account_id: int, + job_id: int, + payload: Dict, + *, + should_poll: bool = True, + poll_interval: int = 10, + autoscale_delete_post_run: bool = True, + autoscale_job_identifier: str = None, + ): + """Check if job with id = job_id is actively running. If it is, create a + clone of the target job and then trigger the clone to run using self.trigger_job + + Args: + account_id (int): Numeric ID of the account to retrieve + job_id (int): Numeric ID of the job to trigger + payload (dict): Payload required for post request + should_poll (bool, optional): Poll until completion if `True`, completion + is one of success, failure, or cancelled + poll_interval (int, optional): Number of seconds to wait in between + polling + autoscale_delete_post_run (bool, optional): Only relevant when job_run_strategy = 'autoscale' + Remove a job replicated via autoscaling after it finishes running. + autoscale_job_identifier (str, optional): Only relevant when job_run_strategy = 'autoscale' + append value to the existing job name when replicating the job definition. + If None defaults to the current timestamp on job creation + """ + autoscale_job_created = False + self.console.log( + 'Triggered with autoscaling set to True. ' + 'Detecting any running instances' + ) + + # dbt Cloud API will remove jobs that are queued mid-run if a DELETE is issued. + # we don't want this behavior so do some config validation + if autoscale_delete_post_run and not should_poll: + self.console.log( + 'autoscale_delete_post_run set to True and should_poll set to False. ' + 'This has the effect that, after a new dbt Cloud job replica is created, ' + 'it will be removed before the run completes, so this configuration is disallowed. ' + ) + raise Exception('Invalid configuration') + + try: + most_recent_job_run = self.list_runs( + account_id=account_id, job_definition_id=job_id, limit=1, order_by='-id' + )['data'][0] + most_recent_job_run_status = most_recent_job_run['status_humanized'] + self.console.log( + f'Status for most recent run of job {job_id} ' + f'is {most_recent_job_run_status}.' + ) + except IndexError: + self.console.log( + f'Failed to get status for most recent run of job: {job_id} ' + f'This happens for jobs that have not previously run. ' + f'Triggering a new run for this job' + ) + + if most_recent_job_run_status not in ['Queued', 'Starting', 'Running']: + self.console.log( + f'autoscale set to true but base job with id {job_id} is free. ' + 'triggering base job and ignoring autoscale configuration.' + ) + + else: + self.console.log(f'job_id {job_id} has an active run. Cloning job.') + job_definition = self._clone_resource( + 'job', + account_id=account_id, + job_id=job_id + ) + + if not autoscale_job_identifier: + creation_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S') + new_job_name = '-'.join([job_definition['name'], creation_time]) + else: + new_job_name = '-'.join([job_definition['name'], autoscale_job_identifier]) + + job_definition['name'] = new_job_name + job = self.create_job(account_id=account_id, payload=job_definition) + job_id = job['data']['id'] + self.console.log(f'Created new job: {job_id}') + autoscale_job_created = True + + self.console.log(f'Triggering job: {job_id}') + run = self.trigger_job( + account_id=account_id, + job_id=job_id, + payload=payload, + should_poll=should_poll, + poll_interval=poll_interval + ) + + if run['status']['code'] in [200, 201] \ + and autoscale_delete_post_run \ + and autoscale_job_created: + + self.console.log(f'Deleting autoscaled job with job_id: {job_id} post run') + self.delete_job( + account_id=account_id, + job_id=job_id + ) + + return run + + @v2 + def trigger_job_restart_from_failure( + self, + account_id: int, + job_id: int, + payload: Dict, + *, + should_poll: bool = True, + poll_interval: int = 10, + ): + """Check if job with job_id had failed steps on its prior run. + If failed steps are found, modify the job run commands to restart + from the point of failure. + If no failed steps are found, return None + + Args: + account_id (int): Numeric ID of the account to retrieve + job_id (int): Numeric ID of the job to trigger + payload (dict): Payload required for post request + should_poll (bool, optional): Poll until completion if `True`, completion + is one of success, failure, or cancelled + poll_interval (int, optional): Number of seconds to wait in between + polling + """ + + self.console.log(f'Detecting failures on prior run of job {job_id}') + payload, has_failures = self._get_restart_job_definition( + account_id=account_id, + job_id=job_id, + payload=payload + ) + + if has_failures: + run = self.trigger_job( + account_id=account_id, + job_id=job_id, + payload=payload, + should_poll=should_poll, + poll_interval=poll_interval + ) + return run + + self.console.log(f'No failed steps identified for job {job_id}. Exiting') + + return None + @v2 def trigger_job( self, @@ -1012,7 +1325,7 @@ def trigger_job( should_poll: bool = True, poll_interval: int = 10, restart_from_failure: bool = False, - trigger_on_failure_only: bool = False, + trigger_on_failure_only: bool = True, ): """Trigger a job by its ID @@ -1030,7 +1343,6 @@ def trigger_job( restart_from_failure to True. This has the effect of only triggering the job when the prior invocation was not successful. Otherwise, the function will exit prior to triggering the job. - """ def run_status_formatted(run: Dict, time: float) -> str: @@ -1046,125 +1358,26 @@ def run_status_formatted(run: Dict, time: float) -> str: f', View here: {url}' ) - def parse_args(cli_args: Iterable[str], namespace: argparse.Namespace): - string = '' - for arg in cli_args: - value = getattr(namespace, arg, None) - if value: - arg = arg.replace('_', '-') - if isinstance(value, bool): - string += f' --{arg}' - else: - string += f" --{arg} '{value}'" - return string - + # this is here to not break existing stuff 09.26.2022 if restart_from_failure: - self.console.log(f'Restarting job {job_id} from last failed state.') - last_run_data = self.list_runs( + run = self.trigger_job_restart_from_failure( account_id=account_id, - include_related=['run_steps'], - job_definition_id=job_id, - order_by='-id', - limit=1, - )['data'][0] - - last_run_status = last_run_data['status_humanized'].lower() - last_run_id = last_run_data['id'] - - if last_run_status == 'error': - rerun_steps = [] - - for run_step in last_run_data['run_steps']: - - status = run_step['status_humanized'].lower() - # Skipping cloning, profile setup, and dbt deps - always - # the first three steps in any run - if run_step['index'] <= 3 or status == 'success': - self.console.log( - f'Skipping rerun for command "{run_step["name"]}" ' - 'as it does not need to be repeated.' - ) - - else: - - # get the dbt command used within this step - command = run_step['name'].partition('`')[2].partition('`')[0] - namespace, remaining = self.parser.parse_known_args( - shlex.split(command) - ) - sub_command = remaining[1] - - if ( - sub_command not in RUN_COMMANDS - and status in ['error', 'cancelled', 'skipped'] - ) or (sub_command in RUN_COMMANDS and status == 'skipped'): - rerun_steps.append(command) - - # errors and failures are when we need to inspect to figure - # out the point of failure - else: - - # get the run results scoped to the step which had an error - # an error here indicates that either: - # 1) the fail-fast flag was set, in which case - # the run_results.json file was never created; or - # 2) there was a problem on dbt Cloud's side saving - # this artifact - try: - step_results = self.get_run_artifact( - account_id=account_id, - run_id=last_run_id, - path='run_results.json', - step=run_step['index'], - )['results'] - - # If the artifact isn't found, the API returns a 404 with - # no json. The ValueError will catch the JSONDecodeError - except ValueError: - rerun_steps.append(command) - else: - rerun_nodes = ' '.join( - [ - record['unique_id'].split('.')[2] - for record in step_results - if record['status'] - in ['error', 'skipped', 'fail'] - ] - ) - global_args = parse_args( - GLOBAL_CLI_ARGS.keys(), namespace - ) - sub_command_args = parse_args( - SUB_COMMAND_CLI_ARGS.keys(), namespace - ) - modified_command = f'dbt{global_args} {sub_command} -s {rerun_nodes}{sub_command_args}' # noqa: E501 - rerun_steps.append(modified_command) - self.console.log( - f'Modifying command "{command}" as an error ' - 'or failure was encountered.' - ) - - payload.update({"steps_override": rerun_steps}) - self.console.log( - f'Triggering modified job to re-run failed steps: {rerun_steps}' - ) - - else: + job_id=job_id, + payload=payload + ) + + if not run and trigger_on_failure_only: self.console.log( - 'Process triggered with restart_from_failure set to True but no ' - 'failed run steps found.' - ) - if trigger_on_failure_only: - self.console.log( 'Not triggering job because prior run was successful.' ) - return - + return + run = self._simple_request( f'accounts/{account_id}/jobs/{job_id}/run/', method='post', json=payload, ) + if not run['status']['is_success']: self.console.log(f'Run NOT triggered for job {job_id}. See run response.') return run @@ -1200,7 +1413,7 @@ def update_connection( payload (dict): Dictionary representing the connection to update """ return self._simple_request( - f'accounts/{account_id}/projects/{project_id}/connections/{connection_id}/', + f'accounts/{account_id}/projects/{project_id}/connections/{connection_id}/', # noqa: E501 method='post', json=payload, ) @@ -1218,7 +1431,7 @@ def update_credentials( payload (dict): Dictionary representing the credentials to update """ return self._simple_request( - f'accounts/{account_id}/projects/{project_id}/credentials/{credentials_id}/', # noqa: E50 + f'accounts/{account_id}/projects/{project_id}/credentials/{credentials_id}/', # noqa: E501 method='post', json=payload, ) diff --git a/dbtc/client/cloud/configs/__init__.py b/dbtc/client/cloud/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dbtc/client/cloud/configs/dbt_core_cli.py b/dbtc/client/cloud/configs/dbt_core_cli.py new file mode 100644 index 0000000..bc9df6a --- /dev/null +++ b/dbtc/client/cloud/configs/dbt_core_cli.py @@ -0,0 +1,17 @@ +run_commands = ['build', 'run', 'test', 'seed', 'snapshot'] + +global_cli_args = { + 'warn_error': {'flags': ('--warn-error',), 'action': 'store_true'}, + 'use_experimental_parser': { + 'flags': ('--use-experimental-parser',), + 'action': 'store_true', + }, +} + +sub_command_cli_args = { + 'vars': {'flags': ('--vars',)}, + 'args': {'flags': ('--args',)}, + 'fail_fast': {'flags': ('-x', '--fail-fast'), 'action': 'store_true'}, + 'full_refresh': {'flags': ('--full-refresh',), 'action': 'store_true'}, + 'store_failures': {'flags': ('--store-failures',), 'action': 'store_true'}, +} diff --git a/dbtc/client/cloud/configs/enums.py b/dbtc/client/cloud/configs/enums.py new file mode 100644 index 0000000..7691ff0 --- /dev/null +++ b/dbtc/client/cloud/configs/enums.py @@ -0,0 +1,11 @@ +# stdlib +import enum + + +class JobRunStatus(enum.IntEnum): + QUEUED = 1 + STARTING = 2 + RUNNING = 3 + SUCCESS = 10 + ERROR = 20 + CANCELLED = 30 diff --git a/dbtc/client/cloud/models/__init__.py b/dbtc/client/cloud/models/__init__.py new file mode 100644 index 0000000..93852cb --- /dev/null +++ b/dbtc/client/cloud/models/__init__.py @@ -0,0 +1,2 @@ +from .job import Job # noqa: F401 +from .project import Project # noqa: F401 diff --git a/dbtc/client/cloud/models/constants.py b/dbtc/client/cloud/models/constants.py new file mode 100644 index 0000000..e861841 --- /dev/null +++ b/dbtc/client/cloud/models/constants.py @@ -0,0 +1,7 @@ +# stdlib +import enum + + +class State(enum.IntEnum): + active = 1 + deleted = 2 diff --git a/dbtc/client/cloud/models/job.py b/dbtc/client/cloud/models/job.py new file mode 100644 index 0000000..86291ac --- /dev/null +++ b/dbtc/client/cloud/models/job.py @@ -0,0 +1,53 @@ +# stdlib +from typing import Any, Dict, List, Literal, Optional + +# third party +from pydantic import BaseModel + + +class _JobExecution(BaseModel): + timeout_seconds: int + + +class _JobSchedule(BaseModel): + cron: str + date: Dict[str, Any] + time: Dict[str, Any] + + +class _JobSettings(BaseModel): + threads: int + target_name: str + + +class _JobTrigger(BaseModel): + custom_branch_only: Optional[bool] + git_provider_webhook: Optional[bool] + github_webhook: bool + schedule: bool + + +class Job(BaseModel): + + # Required + account_id: int + dbt_version: str + environment_id: int + execution: _JobExecution + generate_docs: bool + name: str + project_id: int + run_generate_sources: bool + schedule: _JobSchedule + settings: _JobSettings + triggers: _JobTrigger + state: Literal[1, 2] + + # Optional + deactivated: Optional[bool] = False + deferring_job_definition_id: Optional[int] + execute_steps: Optional[List[str]] + id: Optional[int] + lifecycle_webhooks: Optional[bool] + lifecycle_webhooks_url: Optional[str] + run_failure_count: Optional[int] = 0 diff --git a/dbtc/client/cloud/models/project.py b/dbtc/client/cloud/models/project.py new file mode 100644 index 0000000..0136288 --- /dev/null +++ b/dbtc/client/cloud/models/project.py @@ -0,0 +1,23 @@ +# stdlib +from typing import Optional + +# third party +from pydantic import BaseModel + +from .constants import State + + +class Project(BaseModel): + + # Required + account_id: int + name: str + + # Optional + id: Optional[int] = None + connection_id: Optional[int] = None + dbt_project_subdirectory: Optional[str] = None + docs_job_id: Optional[int] = None + freshness_job_id: Optional[int] = None + repository_id: Optional[int] = None + state: int = State.active diff --git a/poetry.lock b/poetry.lock index 261a20a..47a2e4f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -600,6 +600,21 @@ category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" +[[package]] +name = "pydantic" +version = "1.10.2" +description = "Data validation and settings management using python type hints" +category = "main" +optional = false +python-versions = ">=3.7" + +[package.dependencies] +typing-extensions = ">=4.1.0" + +[package.extras] +dotenv = ["python-dotenv (>=0.10.4)"] +email = ["email-validator (>=1.0.3)"] + [[package]] name = "pyflakes" version = "2.4.0" @@ -938,7 +953,7 @@ testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest- [metadata] lock-version = "1.1" python-versions = "^3.8" -content-hash = "02325e1d8e94719b09b921bed4f10cec541135cb29e02514367fa34594346b7b" +content-hash = "ecbd26f1cd1c50c8dd6a3199923a7c4bc888401255ea46a61097f8c3578ddb90" [metadata.files] appnope = [ @@ -1178,6 +1193,7 @@ pycodestyle = [ {file = "pycodestyle-2.8.0-py2.py3-none-any.whl", hash = "sha256:720f8b39dde8b293825e7ff02c475f3077124006db4f440dcbc9a20b76548a20"}, {file = "pycodestyle-2.8.0.tar.gz", hash = "sha256:eddd5847ef438ea1c7870ca7eb78a9d47ce0cdb4851a5523949f2601d0cbbe7f"}, ] +pydantic = [] pyflakes = [ {file = "pyflakes-2.4.0-py2.py3-none-any.whl", hash = "sha256:3bb3a3f256f4b7968c9c788781e4ff07dce46bdf12339dcda61053375426ee2e"}, {file = "pyflakes-2.4.0.tar.gz", hash = "sha256:05a85c2872edf37a4ed30b0cce2f6093e1d0581f8c19d7393122da7e25b2b24c"}, diff --git a/pyproject.toml b/pyproject.toml index c23c73d..c0fc77e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ python = "^3.8" sgqlc = "^15.0" requests = "^2.27.1" typer = {extras = ["all"], version = "^0.6.1"} +pydantic = "^1.10.2" [tool.poetry.dev-dependencies] black = "^22.1.0" diff --git a/tests/test_cloud.py b/tests/test_cloud.py index 15085cf..35f3983 100644 --- a/tests/test_cloud.py +++ b/tests/test_cloud.py @@ -144,6 +144,37 @@ def test_list_runs(dbtc_client): account_id=pytest.account_id, job_definition_id=pytest.job_id, ) + + +@pytest.mark.dependency(depends=['test_list_jobs']) +def test_list_runs_list_status(dbtc_client): + _test_cloud_method( + dbtc_client, + 'list_runs', + job_definition_id=pytest.job_id, + status=['error', 'success'] + ) + + +@pytest.mark.dependency(depends=['test_list_jobs']) +def test_list_runs_str_status(dbtc_client): + _test_cloud_method( + dbtc_client, + 'list_runs', + job_definition_id=pytest.job_id, + status='success' + ) + + +@pytest.mark.dependency(depends=['test_list_jobs']) +def test_list_runs_bad_status(dbtc_client): + with pytest.raises(AttributeError): + _test_cloud_method( + dbtc_client, + 'list_runs', + job_definition_id=pytest.job_id, + status='successs' + ) @pytest.mark.dependency(depends=['test_list_jobs']) diff --git a/tests/test_trigger_job_with_autoscaling.py b/tests/test_trigger_job_with_autoscaling.py new file mode 100644 index 0000000..d4abb09 --- /dev/null +++ b/tests/test_trigger_job_with_autoscaling.py @@ -0,0 +1,53 @@ +import time + +# This dictionary contains job_ids and what the associated +# override steps should be when restarting from failure. +JOB_ASSERTIONS = { + 133168: {'execute_steps': ['dbt build -s state:modified+']}, +} + +ACCOUNT_ID = 28885 + +def _test_job(dbtc_client, job_id: int): + + first_run = dbtc_client.cloud.trigger_job( + ACCOUNT_ID, + job_id, + payload={'cause': 'Testing dbtc'}, + should_poll=False, + )['data'] + + # wait a few seconds to make sure the first job's status has been updated + time.sleep(3) + second_run = dbtc_client.cloud.trigger_job_with_autoscaling( + ACCOUNT_ID, + job_id, + payload={'cause': 'Testing dbtc'}, + autoscale_delete_post_run=True, + )['data'] + + # check that we triggered distinct jobs + assert first_run['job_definition_id'] != second_run['job_definition_id'] + + # get the first and second runs with run steps included + first_run_data = dbtc_client.cloud.get_run( + account_id=ACCOUNT_ID, + run_id=first_run['id'], + include_related=['run_steps'] + )['data'] + + second_run_data = dbtc_client.cloud.get_run( + account_id=ACCOUNT_ID, + run_id=second_run['id'], + include_related=['run_steps'] + )['data'] + # check that the run steps are the same for the original and replicated jobs + first_run_step_names = [step['name'] for step in first_run_data['run_steps']] + second_run_step_names = [step['name'] for step in second_run_data['run_steps']] + + assert first_run_step_names == second_run_step_names + + +def test_trigger_job_with_autoscaling(dbtc_client): + for job_id in JOB_ASSERTIONS.keys(): + _test_job(dbtc_client, job_id)