From 25fce9b1468b97a00b630f14de61d03cbcc4930d Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 15 Jan 2026 18:21:02 -0600 Subject: [PATCH 1/6] 58941 (feature): moved test connection execution to workers --- airflow-core/docs/img/airflow_erd.sha256 | 2 +- airflow-core/docs/img/airflow_erd.svg | 4542 +++++++++-------- airflow-core/docs/migrations-ref.rst | 5 +- .../core_api/datamodels/connection_test.py | 44 + .../openapi/v2-rest-api-generated.yaml | 219 +- .../core_api/routes/public/connections.py | 164 +- .../datamodels/connection_test.py | 54 + .../execution_api/routes/__init__.py | 4 + .../execution_api/routes/connection_tests.py | 126 + .../src/airflow/executors/base_executor.py | 29 +- .../executors/connection_test_runner.py | 152 + .../src/airflow/executors/local_executor.py | 76 +- .../src/airflow/executors/workloads.py | 45 +- .../src/airflow/jobs/scheduler_job_runner.py | 39 + ...3_2_0_add_connection_test_request_table.py | 69 + airflow-core/src/airflow/models/__init__.py | 2 + .../src/airflow/models/connection_test.py | 189 + .../airflow/ui/openapi-gen/queries/common.ts | 8 +- .../ui/openapi-gen/queries/ensureQueryData.ts | 13 + .../ui/openapi-gen/queries/prefetch.ts | 13 + .../airflow/ui/openapi-gen/queries/queries.ts | 43 +- .../ui/openapi-gen/queries/suspense.ts | 13 + .../ui/openapi-gen/requests/schemas.gen.ts | 86 +- .../ui/openapi-gen/requests/services.gen.ts | 82 +- .../ui/openapi-gen/requests/types.gen.ts | 111 +- .../ui/public/i18n/locales/en/admin.json | 10 + .../Connections/TestConnectionButton.tsx | 3 +- .../ui/src/queries/useTestConnection.ts | 113 +- airflow-core/src/airflow/utils/db.py | 2 +- .../routes/public/test_connections.py | 196 +- .../versions/head/test_connection_tests.py | 301 ++ .../tests/unit/models/test_connection_test.py | 392 ++ .../airflowctl/api/datamodels/generated.py | 21 +- .../src/tests_common/test_utils/db.py | 9 + .../airflow/providers/dbt/cloud/hooks/dbt.py | 5 +- .../airflow/sdk/api/datamodels/_generated.py | 49 + .../sdk/execution_time/execute_workload.py | 51 +- 37 files changed, 4782 insertions(+), 2500 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/core_api/datamodels/connection_test.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py create mode 100644 airflow-core/src/airflow/executors/connection_test_runner.py create mode 100644 airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py create mode 100644 airflow-core/src/airflow/models/connection_test.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py create mode 100644 airflow-core/tests/unit/models/test_connection_test.py diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 index e4b17237c3c7b..55e734468a288 100644 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -1 +1 @@ -174cc04e341bf41946034e5ba04dcc1da7b5b2c7e09cd74b378e07bb87de2ee6 \ No newline at end of file +a2c06c083c41afde4bce7352ed0e510df75a3b2584ca49c4e1f5e602448f0e03 \ No newline at end of file diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg index 49f9a5c66a75b..3068f7d3f1600 100644 --- a/airflow-core/docs/img/airflow_erd.svg +++ b/airflow-core/docs/img/airflow_erd.svg @@ -4,2662 +4,2718 @@ - - + + %3 - + job - -job - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -end_date - - [TIMESTAMP] - -executor_class - - [VARCHAR(500)] - -hostname - - [VARCHAR(500)] - -job_type - - [VARCHAR(30)] - -latest_heartbeat - - [TIMESTAMP] - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -unixname - - [VARCHAR(1000)] + +job + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +end_date + + [TIMESTAMP] + +executor_class + + [VARCHAR(500)] + +hostname + + [VARCHAR(500)] + +job_type + + [VARCHAR(30)] + +latest_heartbeat + + [TIMESTAMP] + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +unixname + + [VARCHAR(1000)] + + + +connection_test_request + +connection_test_request + +id + + [VARCHAR(36)] + NOT NULL + +completed_at + + [TIMESTAMP] + +conn_type + + [VARCHAR(500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +encrypted_connection_uri + + [TEXT] + NOT NULL + +result_message + + [TEXT] + +result_status + + [BOOLEAN] + +started_at + + [TIMESTAMP] + +state + + [VARCHAR(10)] + NOT NULL + +timeout + + [INTEGER] + NOT NULL + +worker_hostname + + [VARCHAR(500)] - + partitioned_asset_key_log - -partitioned_asset_key_log - -id - - [INTEGER] - NOT NULL - -asset_event_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -asset_partition_dag_run_id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -source_partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -target_partition_key - - [VARCHAR(250)] - NOT NULL + +partitioned_asset_key_log + +id + + [INTEGER] + NOT NULL + +asset_event_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +asset_partition_dag_run_id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +source_partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +target_partition_key + + [VARCHAR(250)] + NOT NULL - + log - -log - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -dttm - - [TIMESTAMP] - NOT NULL - -event - - [VARCHAR(60)] - NOT NULL - -extra - - [TEXT] - -logical_date - - [TIMESTAMP] - -map_index - - [INTEGER] - -owner - - [VARCHAR(500)] - -owner_display_name - - [VARCHAR(500)] - -run_id - - [VARCHAR(250)] - -task_id - - [VARCHAR(250)] - -try_number - - [INTEGER] + +log + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +dttm + + [TIMESTAMP] + NOT NULL + +event + + [VARCHAR(60)] + NOT NULL + +extra + + [TEXT] + +logical_date + + [TIMESTAMP] + +map_index + + [INTEGER] + +owner + + [VARCHAR(500)] + +owner_display_name + + [VARCHAR(500)] + +run_id + + [VARCHAR(250)] + +task_id + + [VARCHAR(250)] + +try_number + + [INTEGER] - + dag_priority_parsing_request - -dag_priority_parsing_request - -id - - [VARCHAR(32)] - NOT NULL - -bundle_name - - [VARCHAR(250)] - NOT NULL - -relative_fileloc - - [VARCHAR(2000)] - NOT NULL + +dag_priority_parsing_request + +id + + [VARCHAR(32)] + NOT NULL + +bundle_name + + [VARCHAR(250)] + NOT NULL + +relative_fileloc + + [VARCHAR(2000)] + NOT NULL - + import_error - -import_error - -id - - [INTEGER] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -filename - - [VARCHAR(1024)] - -stacktrace - - [TEXT] - -timestamp - - [TIMESTAMP] + +import_error + +id + + [INTEGER] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +filename + + [VARCHAR(1024)] + +stacktrace + + [TEXT] + +timestamp + + [TIMESTAMP] - + dag_bundle - -dag_bundle - -name - - [VARCHAR(250)] - NOT NULL - -active - - [BOOLEAN] - -last_refreshed - - [TIMESTAMP] - -signed_url_template - - [VARCHAR(200)] - -template_params - - [JSON] - -version - - [VARCHAR(200)] + +dag_bundle + +name + + [VARCHAR(250)] + NOT NULL + +active + + [BOOLEAN] + +last_refreshed + + [TIMESTAMP] + +signed_url_template + + [VARCHAR(200)] + +template_params + + [JSON] + +version + + [VARCHAR(200)] - + dag_bundle_team - -dag_bundle_team - -dag_bundle_name - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - NOT NULL + +dag_bundle_team + +dag_bundle_name + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + NOT NULL dag_bundle:name--dag_bundle_team:dag_bundle_name - -0..N -1 + +0..N +1 - + dag - -dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -asset_expression - - [JSON] - -bundle_name - - [VARCHAR(250)] - NOT NULL - -bundle_version - - [VARCHAR(200)] - -dag_display_name - - [VARCHAR(2000)] - -deadline - - [JSON] - -description - - [TEXT] - -exceeds_max_non_backfill - - [BOOLEAN] - NOT NULL - -fail_fast - - [BOOLEAN] - NOT NULL - -fileloc - - [VARCHAR(2000)] - -has_import_errors - - [BOOLEAN] - NOT NULL - -has_task_concurrency_limits - - [BOOLEAN] - NOT NULL - -is_paused - - [BOOLEAN] - NOT NULL - -is_stale - - [BOOLEAN] - NOT NULL - -last_expired - - [TIMESTAMP] - -last_parse_duration - - [DOUBLE PRECISION] - -last_parsed_time - - [TIMESTAMP] - -max_active_runs - - [INTEGER] - -max_active_tasks - - [INTEGER] - NOT NULL - -max_consecutive_failed_dag_runs - - [INTEGER] - NOT NULL - -next_dagrun - - [TIMESTAMP] - -next_dagrun_create_after - - [TIMESTAMP] - -next_dagrun_data_interval_end - - [TIMESTAMP] - -next_dagrun_data_interval_start - - [TIMESTAMP] - -owners - - [VARCHAR(2000)] - -relative_fileloc - - [VARCHAR(2000)] - -timetable_description - - [VARCHAR(1000)] - -timetable_summary - - [TEXT] - -timetable_type - - [VARCHAR(255)] - NOT NULL + +dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +asset_expression + + [JSON] + +bundle_name + + [VARCHAR(250)] + NOT NULL + +bundle_version + + [VARCHAR(200)] + +dag_display_name + + [VARCHAR(2000)] + +deadline + + [JSON] + +description + + [TEXT] + +exceeds_max_non_backfill + + [BOOLEAN] + NOT NULL + +fail_fast + + [BOOLEAN] + NOT NULL + +fileloc + + [VARCHAR(2000)] + +has_import_errors + + [BOOLEAN] + NOT NULL + +has_task_concurrency_limits + + [BOOLEAN] + NOT NULL + +is_paused + + [BOOLEAN] + NOT NULL + +is_stale + + [BOOLEAN] + NOT NULL + +last_expired + + [TIMESTAMP] + +last_parse_duration + + [DOUBLE PRECISION] + +last_parsed_time + + [TIMESTAMP] + +max_active_runs + + [INTEGER] + +max_active_tasks + + [INTEGER] + NOT NULL + +max_consecutive_failed_dag_runs + + [INTEGER] + NOT NULL + +next_dagrun + + [TIMESTAMP] + +next_dagrun_create_after + + [TIMESTAMP] + +next_dagrun_data_interval_end + + [TIMESTAMP] + +next_dagrun_data_interval_start + + [TIMESTAMP] + +owners + + [VARCHAR(2000)] + +relative_fileloc + + [VARCHAR(2000)] + +timetable_description + + [VARCHAR(1000)] + +timetable_summary + + [TEXT] + +timetable_type + + [VARCHAR(255)] + NOT NULL dag_bundle:name--dag:bundle_name - -0..N -1 + +0..N +1 - + team - -team - -name - - [VARCHAR(50)] - NOT NULL + +team + +name + + [VARCHAR(50)] + NOT NULL team:name--dag_bundle_team:team_name - -0..N -1 + +0..N +1 - + connection - -connection - -id - - [INTEGER] - NOT NULL - -conn_id - - [VARCHAR(250)] - NOT NULL - -conn_type - - [VARCHAR(500)] - NOT NULL - -description - - [TEXT] - -extra - - [TEXT] - -host - - [VARCHAR(500)] - -is_encrypted - - [BOOLEAN] - NOT NULL - -is_extra_encrypted - - [BOOLEAN] - NOT NULL - -login - - [TEXT] - -password - - [TEXT] - -port - - [INTEGER] - -schema - - [VARCHAR(500)] - -team_name - - [VARCHAR(50)] + +connection + +id + + [INTEGER] + NOT NULL + +conn_id + + [VARCHAR(250)] + NOT NULL + +conn_type + + [VARCHAR(500)] + NOT NULL + +description + + [TEXT] + +extra + + [TEXT] + +host + + [VARCHAR(500)] + +is_encrypted + + [BOOLEAN] + NOT NULL + +is_extra_encrypted + + [BOOLEAN] + NOT NULL + +login + + [TEXT] + +password + + [TEXT] + +port + + [INTEGER] + +schema + + [VARCHAR(500)] + +team_name + + [VARCHAR(50)] team:name--connection:team_name - -0..N -{0,1} + +0..N +{0,1} - + slot_pool - -slot_pool - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -include_deferred - - [BOOLEAN] - NOT NULL - -pool - - [VARCHAR(256)] - NOT NULL - -slots - - [INTEGER] - NOT NULL - -team_name - - [VARCHAR(50)] + +slot_pool + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +include_deferred + + [BOOLEAN] + NOT NULL + +pool + + [VARCHAR(256)] + NOT NULL + +slots + + [INTEGER] + NOT NULL + +team_name + + [VARCHAR(50)] team:name--slot_pool:team_name - -0..N -{0,1} + +0..N +{0,1} - + variable - -variable - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -is_encrypted - - [BOOLEAN] - NOT NULL - -key - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - -val - - [TEXT] - NOT NULL + +variable + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +is_encrypted + + [BOOLEAN] + NOT NULL + +key + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + +val + + [TEXT] + NOT NULL team:name--variable:team_name - -0..N -{0,1} + +0..N +{0,1} - + trigger - -trigger - -id - - [INTEGER] - NOT NULL - -classpath - - [VARCHAR(1000)] - NOT NULL - -created_date - - [TIMESTAMP] - NOT NULL - -kwargs - - [TEXT] - NOT NULL - -queue - - [VARCHAR(256)] - -triggerer_id - - [INTEGER] + +trigger + +id + + [INTEGER] + NOT NULL + +classpath + + [VARCHAR(1000)] + NOT NULL + +created_date + + [TIMESTAMP] + NOT NULL + +kwargs + + [TEXT] + NOT NULL + +queue + + [VARCHAR(256)] + +triggerer_id + + [INTEGER] - + callback - -callback - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -data - - [JSONB] - NOT NULL - -fetch_method - - [VARCHAR(20)] - NOT NULL - -output - - [TEXT] - -priority_weight - - [INTEGER] - NOT NULL - -state - - [VARCHAR(10)] - -trigger_id - - [INTEGER] - -type - - [VARCHAR(20)] - NOT NULL + +callback + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +data + + [JSONB] + NOT NULL + +fetch_method + + [VARCHAR(20)] + NOT NULL + +output + + [TEXT] + +priority_weight + + [INTEGER] + NOT NULL + +state + + [VARCHAR(10)] + +trigger_id + + [INTEGER] + +type + + [VARCHAR(20)] + NOT NULL trigger:id--callback:trigger_id - -0..N -{0,1} + +0..N +{0,1} - + asset_watcher - -asset_watcher - -asset_id - - [INTEGER] - NOT NULL - -trigger_id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL + +asset_watcher + +asset_id + + [INTEGER] + NOT NULL + +trigger_id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL - + trigger:id--asset_watcher:trigger_id - -0..N -1 + +0..N +1 - + task_instance - -task_instance - -id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - NOT NULL - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - NOT NULL - -last_heartbeat_at - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - NOT NULL - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - NOT NULL - -queue - - [VARCHAR(256)] - NOT NULL - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - NOT NULL - -updated_at - - [TIMESTAMP] + +task_instance + +id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + NOT NULL + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + NOT NULL + +last_heartbeat_at + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + NOT NULL + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + NOT NULL + +queue + + [VARCHAR(256)] + NOT NULL + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + NOT NULL + +updated_at + + [TIMESTAMP] - + trigger:id--task_instance:trigger_id - -0..N -{0,1} + +0..N +{0,1} - + deadline - -deadline - -id - - [UUID] - NOT NULL - -callback_id - - [UUID] - NOT NULL - -dagrun_id - - [INTEGER] - -deadline_time - - [TIMESTAMP] - NOT NULL - -missed - - [BOOLEAN] - NOT NULL + +deadline + +id + + [UUID] + NOT NULL + +callback_id + + [UUID] + NOT NULL + +dagrun_id + + [INTEGER] + +deadline_time + + [TIMESTAMP] + NOT NULL + +missed + + [BOOLEAN] + NOT NULL callback:id--deadline:callback_id - -0..N -1 + +0..N +1 - + asset_alias - -asset_alias - -id - - [INTEGER] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL + +asset_alias + +id + + [INTEGER] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL - + asset_alias_asset - -asset_alias_asset - -alias_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL + +asset_alias_asset + +alias_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL asset_alias:id--asset_alias_asset:alias_id - -0..N -1 + +0..N +1 - + asset_alias_asset_event - -asset_alias_asset_event - -alias_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +asset_alias_asset_event + +alias_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL asset_alias:id--asset_alias_asset_event:alias_id - -0..N -1 + +0..N +1 - + dag_schedule_asset_alias_reference - -dag_schedule_asset_alias_reference - -alias_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_alias_reference + +alias_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset_alias:id--dag_schedule_asset_alias_reference:alias_id - -0..N -1 + +0..N +1 - + asset - -asset - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -extra - - [JSON] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +extra + + [JSON] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL asset:id--asset_alias_asset:asset_id - -0..N -1 + +0..N +1 - + asset:id--asset_watcher:asset_id - -0..N -1 + +0..N +1 - + asset_active - -asset_active - -name - - [VARCHAR(1500)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset_active + +name + + [VARCHAR(1500)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL -asset:name--asset_active:name - -1 -1 +asset:uri--asset_active:uri + +1 +1 -asset:uri--asset_active:uri - -1 -1 +asset:name--asset_active:name + +1 +1 - + dag_schedule_asset_reference - -dag_schedule_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset:id--dag_schedule_asset_reference:asset_id - -0..N -1 + +0..N +1 - + task_outlet_asset_reference - -task_outlet_asset_reference + +task_outlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL +updated_at + + [TIMESTAMP] + NOT NULL asset:id--task_outlet_asset_reference:asset_id - -0..N -1 + +0..N +1 - + task_inlet_asset_reference - -task_inlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_inlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL asset:id--task_inlet_asset_reference:asset_id - -0..N -1 + +0..N +1 - + asset_dag_run_queue - -asset_dag_run_queue - -asset_id - - [INTEGER] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +asset_dag_run_queue + +asset_id + + [INTEGER] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL - + asset:id--asset_dag_run_queue:asset_id - -0..N -1 + +0..N +1 - + asset_event - -asset_event - -id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -extra - - [JSON] - NOT NULL - -partition_key - - [VARCHAR(250)] - -source_dag_id - - [VARCHAR(250)] - -source_map_index - - [INTEGER] - -source_run_id - - [VARCHAR(250)] - -source_task_id - - [VARCHAR(250)] - -timestamp - - [TIMESTAMP] - NOT NULL + +asset_event + +id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +extra + + [JSON] + NOT NULL + +partition_key + + [VARCHAR(250)] + +source_dag_id + + [VARCHAR(250)] + +source_map_index + + [INTEGER] + +source_run_id + + [VARCHAR(250)] + +source_task_id + + [VARCHAR(250)] + +timestamp + + [TIMESTAMP] + NOT NULL asset_event:id--asset_alias_asset_event:event_id - -0..N -1 + +0..N +1 - + dagrun_asset_event - -dagrun_asset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +dagrun_asset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL - + asset_event:id--dagrun_asset_event:event_id - -0..N -1 + +0..N +1 - + dag_schedule_asset_name_reference - -dag_schedule_asset_name_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_name_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_schedule_asset_name_reference:dag_id - -0..N -1 + +0..N +1 - + dag_schedule_asset_uri_reference - -dag_schedule_asset_uri_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_uri_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_schedule_asset_uri_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--dag_schedule_asset_alias_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--dag_schedule_asset_reference:dag_id - -0..N -1 + +0..N +1 dag:dag_id--task_outlet_asset_reference:dag_id - -0..N -1 + +0..N +1 dag:dag_id--task_inlet_asset_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--asset_dag_run_queue:target_dag_id - -0..N -1 + +0..N +1 - + dag_version - -dag_version - -id - - [UUID] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -bundle_version - - [VARCHAR(250)] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -version_number - - [INTEGER] - NOT NULL + +dag_version + +id + + [UUID] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +bundle_version + + [VARCHAR(250)] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +version_number + + [INTEGER] + NOT NULL dag:dag_id--dag_version:dag_id - -0..N -1 + +0..N +1 - + dag_tag - -dag_tag - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL + +dag_tag + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL dag:dag_id--dag_tag:dag_id - -0..N -1 + +0..N +1 - + dag_owner_attributes - -dag_owner_attributes - -dag_id - - [VARCHAR(250)] - NOT NULL - -owner - - [VARCHAR(500)] - NOT NULL - -link - - [VARCHAR(500)] - NOT NULL + +dag_owner_attributes + +dag_id + + [VARCHAR(250)] + NOT NULL + +owner + + [VARCHAR(500)] + NOT NULL + +link + + [VARCHAR(500)] + NOT NULL dag:dag_id--dag_owner_attributes:dag_id - -0..N -1 + +0..N +1 - + dag_warning - -dag_warning - -dag_id - - [VARCHAR(250)] - NOT NULL - -warning_type - - [VARCHAR(50)] - NOT NULL - -message - - [TEXT] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL + +dag_warning + +dag_id + + [VARCHAR(250)] + NOT NULL + +warning_type + + [VARCHAR(50)] + NOT NULL + +message + + [TEXT] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_warning:dag_id - -0..N -1 + +0..N +1 - + dag_favorite - -dag_favorite - -dag_id - - [VARCHAR(250)] - NOT NULL - -user_id - - [VARCHAR(250)] - NOT NULL + +dag_favorite + +dag_id + + [VARCHAR(250)] + NOT NULL + +user_id + + [VARCHAR(250)] + NOT NULL dag:dag_id--dag_favorite:dag_id - -0..N -1 + +0..N +1 - + dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - -bundle_version - - [VARCHAR(250)] - -clear_number - - [INTEGER] - NOT NULL - -conf - - [JSONB] - -context_carrier - - [JSONB] - -created_dag_version_id - - [UUID] - -creating_job_id - - [INTEGER] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - NOT NULL - -logical_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -queued_at - - [TIMESTAMP] - -run_after - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -scheduled_by_job_id - - [INTEGER] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - NOT NULL - -triggered_by - - [VARCHAR(50)] - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + +bundle_version + + [VARCHAR(250)] + +clear_number + + [INTEGER] + NOT NULL + +conf + + [JSONB] + +context_carrier + + [JSONB] + +created_dag_version_id + + [UUID] + +creating_job_id + + [INTEGER] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + NOT NULL + +logical_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +queued_at + + [TIMESTAMP] + +run_after + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +scheduled_by_job_id + + [INTEGER] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + NOT NULL + +triggered_by + + [VARCHAR(50)] + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL dag_version:id--dag_run:created_dag_version_id - -0..N -{0,1} + +0..N +{0,1} - + dag_code - -dag_code - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -fileloc - - [VARCHAR(2000)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -source_code - - [TEXT] - NOT NULL - -source_code_hash - - [VARCHAR(32)] - NOT NULL + +dag_code + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +fileloc + + [VARCHAR(2000)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +source_code + + [TEXT] + NOT NULL + +source_code_hash + + [VARCHAR(32)] + NOT NULL dag_version:id--dag_code:dag_version_id - -0..N -1 + +0..N +1 - + serialized_dag - -serialized_dag - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_hash - - [VARCHAR(32)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -data - - [JSONB] - -data_compressed - - [BYTEA] - -last_updated - - [TIMESTAMP] - NOT NULL + +serialized_dag + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_hash + + [VARCHAR(32)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +data + + [JSONB] + +data_compressed + + [BYTEA] + +last_updated + + [TIMESTAMP] + NOT NULL dag_version:id--serialized_dag:dag_version_id - -0..N -1 + +0..N +1 - + dag_version:id--task_instance:dag_version_id - -0..N -{0,1} + +0..N +{0,1} - + log_template - -log_template - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -elasticsearch_id - - [TEXT] - NOT NULL - -filename - - [TEXT] - NOT NULL + +log_template + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +elasticsearch_id + + [TEXT] + NOT NULL + +filename + + [TEXT] + NOT NULL log_template:id--dag_run:log_template_id - -0..N -1 + +0..N +1 dag_run:id--deadline:dagrun_id - -0..N -{0,1} + +0..N +{0,1} - + dag_run:id--dagrun_asset_event:dag_run_id - -0..N -1 + +0..N +1 - + asset_partition_dag_run - -asset_partition_dag_run - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_run_id - - [INTEGER] - -partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +asset_partition_dag_run + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_run_id + + [INTEGER] + +partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL dag_run:id--asset_partition_dag_run:created_dag_run_id - -0..N -{0,1} + +0..N +{0,1} dag_run:run_id--task_instance:run_id - -0..N -1 + +0..N +1 dag_run:dag_id--task_instance:dag_id - -0..N -1 + +0..N +1 - + backfill_dag_run - -backfill_dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - NOT NULL - -dag_run_id - - [INTEGER] - -exception_reason - - [VARCHAR(250)] - -logical_date - - [TIMESTAMP] - NOT NULL - -sort_ordinal - - [INTEGER] - NOT NULL + +backfill_dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + NOT NULL + +dag_run_id + + [INTEGER] + +exception_reason + + [VARCHAR(250)] + +logical_date + + [TIMESTAMP] + NOT NULL + +sort_ordinal + + [INTEGER] + NOT NULL dag_run:id--backfill_dag_run:dag_run_id - -0..N -{0,1} + +0..N +{0,1} - + dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] dag_run:id--dag_run_note:dag_run_id - -1 -1 + +1 +1 - + backfill - -backfill - -id - - [INTEGER] - NOT NULL - -completed_at - - [TIMESTAMP] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_run_conf - - [JSON] - NOT NULL - -from_date - - [TIMESTAMP] - NOT NULL - -is_paused - - [BOOLEAN] - -max_active_runs - - [INTEGER] - NOT NULL - -reprocess_behavior - - [VARCHAR(250)] - NOT NULL - -to_date - - [TIMESTAMP] - NOT NULL - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL + +backfill + +id + + [INTEGER] + NOT NULL + +completed_at + + [TIMESTAMP] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_run_conf + + [JSON] + NOT NULL + +from_date + + [TIMESTAMP] + NOT NULL + +is_paused + + [BOOLEAN] + +max_active_runs + + [INTEGER] + NOT NULL + +reprocess_behavior + + [VARCHAR(250)] + NOT NULL + +to_date + + [TIMESTAMP] + NOT NULL + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL backfill:id--dag_run:backfill_id - -0..N -{0,1} + +0..N +{0,1} backfill:id--backfill_dag_run:backfill_id - -0..N -1 + +0..N +1 - + hitl_detail - -hitl_detail - -ti_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL + +hitl_detail + +ti_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL task_instance:id--hitl_detail:ti_id - -1 -1 + +1 +1 - + task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSONB] - -length - - [INTEGER] - NOT NULL + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSONB] + +length + + [INTEGER] + NOT NULL -task_instance:dag_id--task_map:dag_id - -0..N -1 +task_instance:map_index--task_map:map_index + +0..N +1 task_instance:task_id--task_map:task_id - -0..N -1 + +0..N +1 -task_instance:map_index--task_map:map_index - -0..N -1 +task_instance:dag_id--task_map:dag_id + +0..N +1 task_instance:run_id--task_map:run_id - -0..N -1 + +0..N +1 - + task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -ti_id - - [UUID] - NOT NULL + +task_reschedule + +id + + [INTEGER] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +ti_id + + [UUID] + NOT NULL task_instance:id--task_reschedule:ti_id - -0..N -1 + +0..N +1 - + xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [JSONB] + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [JSONB] -task_instance:run_id--xcom:run_id - -0..N -1 +task_instance:dag_id--xcom:dag_id + +0..N +1 task_instance:task_id--xcom:task_id - -0..N -1 + +0..N +1 -task_instance:dag_id--xcom:dag_id - -0..N -1 +task_instance:run_id--xcom:run_id + +0..N +1 task_instance:map_index--xcom:map_index - -0..N -1 + +0..N +1 - + task_instance_note - -task_instance_note - -ti_id - - [UUID] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +task_instance_note + +ti_id + + [UUID] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] task_instance:id--task_instance_note:ti_id - -1 -1 + +1 +1 - + task_instance_history - -task_instance_history - -task_instance_id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] + +task_instance_history + +task_instance_id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] -task_instance:run_id--task_instance_history:run_id - -0..N -1 +task_instance:task_id--task_instance_history:task_id + +0..N +1 -task_instance:task_id--task_instance_history:task_id - -0..N -1 +task_instance:run_id--task_instance_history:run_id + +0..N +1 task_instance:map_index--task_instance_history:map_index - -0..N -1 + +0..N +1 task_instance:dag_id--task_instance_history:dag_id - -0..N -1 + +0..N +1 - + rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL -task_instance:map_index--rendered_task_instance_fields:map_index - -0..N -1 +task_instance:run_id--rendered_task_instance_fields:run_id + +0..N +1 -task_instance:task_id--rendered_task_instance_fields:task_id - -0..N -1 +task_instance:map_index--rendered_task_instance_fields:map_index + +0..N +1 -task_instance:run_id--rendered_task_instance_fields:run_id - -0..N -1 +task_instance:dag_id--rendered_task_instance_fields:dag_id + +0..N +1 -task_instance:dag_id--rendered_task_instance_fields:dag_id - -0..N -1 +task_instance:task_id--rendered_task_instance_fields:task_id + +0..N +1 - + hitl_detail_history - -hitl_detail_history - -ti_history_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL + +hitl_detail_history + +ti_history_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL task_instance_history:task_instance_id--hitl_detail_history:ti_history_id - -1 -1 + +1 +1 - + alembic_version - -alembic_version - -version_num - - [VARCHAR(32)] - NOT NULL + +alembic_version + +version_num + + [VARCHAR(32)] + NOT NULL diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 3b2080410e228..614a163c015ef 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``e79fc784f145`` (head) | ``0b112f49112d`` | ``3.2.0`` | add timetable_type to dag table for filtering. | +| ``9882c124ea54`` (head) | ``e79fc784f145`` | ``3.2.0`` | Add connection_test_request table for async connection | +| | | | testing on workers. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``e79fc784f145`` | ``0b112f49112d`` | ``3.2.0`` | add timetable_type to dag table for filtering. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``0b112f49112d`` | ``c47f2e1ab9d4`` | ``3.2.0`` | Add exceeds max runs flag to dag model. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connection_test.py new file mode 100644 index 0000000000000..bbb583806a985 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connection_test.py @@ -0,0 +1,44 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Datamodels for async connection testing.""" + +from __future__ import annotations + +from datetime import datetime + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ConnectionTestQueuedResponse(BaseModel): + """Response when a connection test is queued for async execution.""" + + request_id: str + state: str + message: str + + +class ConnectionTestStatusResponse(BaseModel): + """Response with the full status of a connection test request.""" + + request_id: str + state: str + result_status: bool | None + result_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 7d841f3b46b56..718ff04f2ece9 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1297,6 +1297,112 @@ paths: security: - OAuth2PasswordBearer: [] - HTTPBearer: [] + /api/v2/connections/test: + post: + tags: + - Connection + summary: Test Connection + description: 'Queue a connection test for asynchronous execution on a worker. + + + This endpoint queues the connection test request for execution on a worker + node, + + which provides better security isolation (workers run in ephemeral environments) + + and network accessibility (workers can reach external systems that API servers + cannot). + + + Returns a request_id that can be used to poll for the test result via GET + /connections/test/{request_id}.' + operationId: test_connection + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionBody' + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionTestQueuedResponse' + '401': + description: Unauthorized + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '403': + description: Forbidden + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] + /api/v2/connections/test/{request_id}: + get: + tags: + - Connection + summary: Get Connection Test Status + description: 'Get the status of a connection test request. + + + Poll this endpoint to check if a connection test has completed and retrieve + the result.' + operationId: get_connection_test_status + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] + parameters: + - name: request_id + in: path + required: true + schema: + type: string + title: Request Id + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionTestStatusResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /api/v2/connections/{connection_id}: delete: tags: @@ -1631,56 +1737,6 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' - /api/v2/connections/test: - post: - tags: - - Connection - summary: Test Connection - description: 'Test an API connection. - - - This method first creates an in-memory transient conn_id & exports that to - an env var, - - as some hook classes tries to find out the `conn` from their __init__ method - & errors out if not found. - - It also deletes the conn id env connection after the test.' - operationId: test_connection - requestBody: - content: - application/json: - schema: - $ref: '#/components/schemas/ConnectionBody' - required: true - responses: - '200': - description: Successful Response - content: - application/json: - schema: - $ref: '#/components/schemas/ConnectionTestResponse' - '401': - description: Unauthorized - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - '403': - description: Forbidden - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' - security: - - OAuth2PasswordBearer: [] - - HTTPBearer: [] /api/v2/connections/defaults: post: tags: @@ -10066,20 +10122,69 @@ components: - team_name title: ConnectionResponse description: Connection serializer for responses. - ConnectionTestResponse: + ConnectionTestQueuedResponse: properties: - status: - type: boolean - title: Status + request_id: + type: string + title: Request Id + state: + type: string + title: State message: type: string title: Message type: object required: - - status + - request_id + - state - message - title: ConnectionTestResponse - description: Connection Test serializer for responses. + title: ConnectionTestQueuedResponse + description: Response when a connection test is queued for async execution. + ConnectionTestStatusResponse: + properties: + request_id: + type: string + title: Request Id + state: + type: string + title: State + result_status: + anyOf: + - type: boolean + - type: 'null' + title: Result Status + result_message: + anyOf: + - type: string + - type: 'null' + title: Result Message + created_at: + type: string + format: date-time + title: Created At + started_at: + anyOf: + - type: string + format: date-time + - type: 'null' + title: Started At + completed_at: + anyOf: + - type: string + format: date-time + - type: 'null' + title: Completed At + type: object + required: + - request_id + - state + - result_status + - result_message + - created_at + - started_at + - completed_at + title: ConnectionTestStatusResponse + description: Response with the full status of a connection test request. CreateAssetEventsBody: properties: asset_id: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 7dce8fd4ea806..c61a43f41eb04 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -import os +import uuid from typing import Annotated from fastapi import Depends, HTTPException, Query, status @@ -36,11 +36,14 @@ BulkBody, BulkResponse, ) +from airflow.api_fastapi.core_api.datamodels.connection_test import ( + ConnectionTestQueuedResponse, + ConnectionTestStatusResponse, +) from airflow.api_fastapi.core_api.datamodels.connections import ( ConnectionBody, ConnectionCollectionResponse, ConnectionResponse, - ConnectionTestResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.security import ( @@ -55,13 +58,137 @@ from airflow.api_fastapi.logging.decorators import action_logging from airflow.configuration import conf from airflow.models import Connection -from airflow.secrets.environment_variables import CONN_ENV_PREFIX +from airflow.models.connection_test import ConnectionTestRequest +from airflow.models.crypto import get_fernet from airflow.utils.db import create_default_connections as db_create_default_connections from airflow.utils.strings import get_random_string connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") +# NOTE: /test routes must be defined BEFORE /{connection_id} routes to avoid route conflicts +@connections_router.post( + "/test", + dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())], + responses=create_openapi_http_exception_doc([status.HTTP_403_FORBIDDEN]), +) +def test_connection( + test_body: ConnectionBody, + session: SessionDep, +) -> ConnectionTestQueuedResponse: + """ + Queue a connection test for asynchronous execution on a worker. + + This endpoint queues the connection test request for execution on a worker node, + which provides better security isolation (workers run in ephemeral environments) + and network accessibility (workers can reach external systems that API servers cannot). + + Returns a request_id that can be used to poll for the test result via GET /connections/test/{request_id}. + """ + if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": + raise HTTPException( + status.HTTP_403_FORBIDDEN, + "Testing connections is disabled in Airflow configuration. " + "Contact your deployment admin to enable it.", + ) + + # Create a transient connection to get its URI + transient_conn_id = get_random_string() + + # Check if we're testing an existing connection (connection_id provided) + # In this case, we need to merge masked fields (password, extra) with stored values + if test_body.connection_id: + existing_conn = session.scalar(select(Connection).filter_by(conn_id=test_body.connection_id)) + if existing_conn: + # Create a copy of the existing connection for testing + # This merges request data with stored credentials (handles masked passwords) + conn = Connection( + conn_id=transient_conn_id, + conn_type=existing_conn.conn_type, + description=existing_conn.description, + host=existing_conn.host, + schema=existing_conn.schema, + login=existing_conn.login, + port=existing_conn.port, + ) + # Copy password and extra (these are encrypted in db) + conn.set_password(existing_conn.password) + conn.set_extra(existing_conn.extra) + # Now apply updates from request body, merging masked fields + update_orm_from_pydantic(conn, test_body) + else: + # Connection ID provided but not found - use request body as-is + data = test_body.model_dump(by_alias=True) + data["conn_id"] = transient_conn_id + conn = Connection(**data) + else: + # New connection test - use request body directly + data = test_body.model_dump(by_alias=True) + data["conn_id"] = transient_conn_id + conn = Connection(**data) + + # Encrypt the connection URI for secure storage + fernet = get_fernet() + connection_uri = conn.get_uri() + encrypted_uri = fernet.encrypt(connection_uri.encode("utf-8")).decode("utf-8") + + # Create the test request + test_request = ConnectionTestRequest.create_request( + encrypted_connection_uri=encrypted_uri, + conn_type=test_body.conn_type, + session=session, + ) + session.commit() + + return ConnectionTestQueuedResponse( + request_id=str(test_request.id), + state=test_request.state, + message="Connection test request queued for execution on a worker.", + ) + + +@connections_router.get( + "/test/{request_id}", + dependencies=[Depends(requires_access_connection(method="GET"))], + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), +) +def get_connection_test_status( + request_id: str, + session: SessionDep, +) -> ConnectionTestStatusResponse: + """ + Get the status of a connection test request. + + Poll this endpoint to check if a connection test has completed and retrieve the result. + """ + # Validate that request_id is a valid UUID format + try: + uuid.UUID(request_id) + except (ValueError, TypeError): + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"Connection test request with id `{request_id}` was not found.", + ) + + test_request = session.scalar(select(ConnectionTestRequest).where(ConnectionTestRequest.id == request_id)) + + if test_request is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"Connection test request with id `{request_id}` was not found.", + ) + + return ConnectionTestStatusResponse( + request_id=str(test_request.id), + state=test_request.state, + result_status=test_request.result_status, + result_message=test_request.result_message, + created_at=test_request.created_at, + started_at=test_request.started_at, + completed_at=test_request.completed_at, + ) + + @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -211,37 +338,6 @@ def patch_connection( return connection -@connections_router.post("/test", dependencies=[Depends(requires_access_connection(method="POST"))]) -def test_connection( - test_body: ConnectionBody, -) -> ConnectionTestResponse: - """ - Test an API connection. - - This method first creates an in-memory transient conn_id & exports that to an env var, - as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found. - It also deletes the conn id env connection after the test. - """ - if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": - raise HTTPException( - status.HTTP_403_FORBIDDEN, - "Testing connections is disabled in Airflow configuration. " - "Contact your deployment admin to enable it.", - ) - - transient_conn_id = get_random_string() - conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}" - try: - data = test_body.model_dump(by_alias=True) - data["conn_id"] = transient_conn_id - conn = Connection(**data) - os.environ[conn_env_var] = conn.get_uri() - test_status, test_message = conn.test_connection() - return ConnectionTestResponse.model_validate({"status": test_status, "message": test_message}) - finally: - os.environ.pop(conn_env_var, None) - - @connections_router.post( "/defaults", status_code=status.HTTP_204_NO_CONTENT, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py new file mode 100644 index 0000000000000..9c925270b323f --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py @@ -0,0 +1,54 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Datamodels for worker-side connection test execution.""" + +from __future__ import annotations + +from typing import Literal + +from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel + + +class ConnectionTestWorkload(BaseModel): + """Workload data sent to worker for connection test execution.""" + + request_id: str + encrypted_connection_uri: str + conn_type: str + timeout: int + + +class ConnectionTestPendingResponse(BaseModel): + """Response containing pending connection test requests for workers.""" + + requests: list[ConnectionTestWorkload] + + +class ConnectionTestRunningPayload(StrictBaseModel): + """Payload for marking a connection test as running.""" + + state: Literal["running"] + hostname: str + + +class ConnectionTestResultPayload(StrictBaseModel): + """Payload for reporting connection test result.""" + + state: Literal["success", "failed"] + result_status: bool + result_message: str diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 562b8588fbf2c..f1cd194f405fd 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -23,6 +23,7 @@ from airflow.api_fastapi.execution_api.routes import ( asset_events, assets, + connection_tests, connections, dag_runs, health, @@ -41,6 +42,9 @@ authenticated_router.include_router(assets.router, prefix="/assets", tags=["Assets"]) authenticated_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"]) +authenticated_router.include_router( + connection_tests.router, prefix="/connection-tests", tags=["Connection Tests"] +) authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"]) authenticated_router.include_router(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"]) authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py new file mode 100644 index 0000000000000..af2d75ea20747 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -0,0 +1,126 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Execution API routes for worker-side connection test execution.""" + +from __future__ import annotations + +from typing import Annotated + +import structlog +from cadwyn import VersionedAPIRouter +from fastapi import Body, HTTPException, Query, status + +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.execution_api.datamodels.connection_test import ( + ConnectionTestPendingResponse, + ConnectionTestResultPayload, + ConnectionTestRunningPayload, + ConnectionTestWorkload, +) +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState + +router = VersionedAPIRouter() + +log = structlog.get_logger(__name__) + + +@router.get( + "/pending", + status_code=status.HTTP_200_OK, +) +def get_pending_connection_tests( + session: SessionDep, + hostname: Annotated[str, Query(description="Worker hostname requesting work")], + limit: Annotated[int, Query(description="Maximum number of requests to return", ge=1, le=100)] = 10, +) -> ConnectionTestPendingResponse: + """ + Get pending connection test requests for worker execution. + + Workers call this endpoint to fetch pending connection tests to execute. + """ + log.debug("Worker requesting pending connection tests", hostname=hostname, limit=limit) + + pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=limit) + + workloads = [] + for request in pending_requests: + request.mark_running(hostname) + workloads.append( + ConnectionTestWorkload( + request_id=str(request.id), + encrypted_connection_uri=request.encrypted_connection_uri, + conn_type=request.conn_type, + timeout=request.timeout, + ) + ) + + session.commit() + return ConnectionTestPendingResponse(requests=workloads) + + +@router.patch( + "/{request_id}/state", + status_code=status.HTTP_204_NO_CONTENT, + responses={ + status.HTTP_404_NOT_FOUND: {"description": "Connection test request not found"}, + status.HTTP_409_CONFLICT: {"description": "Invalid state transition"}, + }, +) +def update_connection_test_state( + request_id: str, + payload: Annotated[ConnectionTestRunningPayload | ConnectionTestResultPayload, Body()], + session: SessionDep, +) -> None: + """ + Update the state of a connection test request. + + Workers call this endpoint to report the result of a connection test. + """ + log.debug("Updating connection test state", request_id=request_id, payload=payload) + + test_request = session.get(ConnectionTestRequest, request_id) + + if test_request is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"Connection test request with id `{request_id}` was not found.", + ) + + if isinstance(payload, ConnectionTestRunningPayload): + if test_request.state not in (ConnectionTestState.PENDING.value, ConnectionTestState.RUNNING.value): + raise HTTPException( + status.HTTP_409_CONFLICT, + f"Cannot transition from state `{test_request.state}` to `running`.", + ) + test_request.mark_running(payload.hostname) + elif isinstance(payload, ConnectionTestResultPayload): + if test_request.state != ConnectionTestState.RUNNING.value: + raise HTTPException( + status.HTTP_409_CONFLICT, + f"Cannot report result when in state `{test_request.state}`. Expected `running`.", + ) + if payload.state == "success": + test_request.mark_success(payload.result_message) + else: + test_request.mark_failed(payload.result_message) + + log.info( + "Connection test state updated", + request_id=request_id, + new_state=test_request.state, + ) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 64fa475626797..1bacb68e64ad1 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -188,6 +188,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None) self.parallelism: int = parallelism self.team_name: str | None = team_name self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {} + self.queued_connection_tests: deque[workloads.TestConnection] = deque() self.running: set[TaskInstanceKey] = set() self.event_buffer: dict[TaskInstanceKey, EventBufferValueType] = {} self._task_event_logs: deque[Log] = deque() @@ -226,10 +227,27 @@ def log_task_event(self, *, event: str, extra: str, ti_key: TaskInstanceKey): self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra)) def queue_workload(self, workload: workloads.All, session: Session) -> None: - if not isinstance(workload, workloads.ExecuteTask): + if isinstance(workload, workloads.ExecuteTask): + ti = workload.ti + self.queued_tasks[ti.key] = workload + elif isinstance(workload, workloads.TestConnection): + # Queue connection test workloads separately + self._queue_connection_test(workload, session) + else: raise ValueError(f"Un-handled workload kind {type(workload).__name__!r} in {type(self).__name__}") - ti = workload.ti - self.queued_tasks[ti.key] = workload + + def _queue_connection_test(self, workload: workloads.TestConnection, session: Session) -> None: + """ + Queue a connection test workload for execution. + + This method can be overridden by subclasses to provide executor-specific + handling of connection test workloads. + + :param workload: The TestConnection workload to queue + :param session: SQLAlchemy session + """ + # Add to the queue initialized in __init__ + self.queued_connection_tests.append(workload) def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: """ @@ -402,7 +420,10 @@ def trigger_tasks(self, open_slots: int) -> None: ti.context_carrier = carrier workload_list.append(item) - if workload_list: + + # Process workloads and connection tests + # Call _process_workloads if we have task workloads OR connection tests queued + if workload_list or self.queued_connection_tests: self._process_workloads(workload_list) # TODO: This should not be using `TaskInstanceState` here, this is just "did the process complete, or did diff --git a/airflow-core/src/airflow/executors/connection_test_runner.py b/airflow-core/src/airflow/executors/connection_test_runner.py new file mode 100644 index 0000000000000..0b2c65fa500e9 --- /dev/null +++ b/airflow-core/src/airflow/executors/connection_test_runner.py @@ -0,0 +1,152 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Connection test runner for executing connection tests on workers.""" + +from __future__ import annotations + +import logging +import os +import random +import time +import traceback +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError +from typing import TYPE_CHECKING + +from airflow.models.connection import Connection +from airflow.models.crypto import get_fernet +from airflow.secrets.environment_variables import CONN_ENV_PREFIX +from airflow.utils.strings import get_random_string + +if TYPE_CHECKING: + from airflow.executors.workloads import TestConnection + +log = logging.getLogger(__name__) + +MAX_RETRIES = 3 +RETRY_DELAY_BASE = 2 + + +def _run_test_connection(conn: Connection, conn_type: str) -> tuple[bool, str]: + """Run the actual connection test.""" + log.info("Testing connection of type %s", conn_type) + status, message = conn.test_connection() + log.info("Connection test result: status=%s, message=%s", status, message) + return status, message + + +def execute_connection_test( + encrypted_connection_uri: str, + conn_type: str, + timeout: int = 60, +) -> tuple[bool, str]: + """Execute a connection test with timeout enforcement.""" + try: + fernet = get_fernet() + connection_uri = fernet.decrypt(encrypted_connection_uri.encode("utf-8")).decode("utf-8") + except Exception as e: + log.exception("Failed to decrypt connection URI") + return False, f"Failed to decrypt connection URI: {e}" + + # Some hooks look up the connection in __init__, so we need to export it + transient_conn_id = get_random_string() + conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}" + + try: + conn = Connection(conn_id=transient_conn_id, uri=connection_uri) + os.environ[conn_env_var] = connection_uri + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_run_test_connection, conn, conn_type) + try: + status, message = future.result(timeout=timeout) + return status, message + except FuturesTimeoutError: + log.warning("Connection test timed out after %s seconds", timeout) + return False, f"Connection test timed out after {timeout} seconds" + except Exception: + log.exception("Connection test failed with exception") + return False, f"Connection test failed: {traceback.format_exc()}" + finally: + os.environ.pop(conn_env_var, None) + + +def execute_connection_test_workload(workload: TestConnection) -> tuple[bool, str]: + """Execute a connection test from a TestConnection workload.""" + return execute_connection_test( + encrypted_connection_uri=workload.encrypted_connection_uri, + conn_type=workload.conn_type, + timeout=workload.timeout, + ) + + +def report_connection_test_result( + request_id: str, + success: bool, + message: str, + server_url: str, + token: str, +) -> bool: + """Report connection test result back to the API server with retry logic.""" + import httpx + + state = "success" if success else "failed" + result_url = f"{server_url.rstrip('/')}/connection-tests/{request_id}/state" + payload = {"state": state, "result_status": success, "result_message": message} + + log.info("Reporting connection test result: request_id=%s, state=%s", request_id, state) + + for attempt in range(MAX_RETRIES): + try: + with httpx.Client() as client: + response = client.patch( + result_url, + json=payload, + headers={"Authorization": f"Bearer {token}"}, + timeout=30.0, + ) + response.raise_for_status() + log.info("Connection test result reported: request_id=%s", request_id) + return True + except httpx.HTTPStatusError as e: + if 400 <= e.response.status_code < 500 and e.response.status_code != 429: + log.error( + "Failed to report result (client error): request_id=%s, status=%s", + request_id, + e.response.status_code, + ) + return False + log.warning( + "Failed to report result, retrying: request_id=%s, attempt=%s/%s", + request_id, + attempt + 1, + MAX_RETRIES, + ) + except Exception: + log.warning( + "Failed to report result, retrying: request_id=%s, attempt=%s/%s", + request_id, + attempt + 1, + MAX_RETRIES, + ) + + if attempt < MAX_RETRIES - 1: + delay = RETRY_DELAY_BASE * (2**attempt) * (0.5 + random.random()) + time.sleep(delay) + + log.error("Failed to report result after retries: request_id=%s", request_id) + return False diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 0724886132313..9f5a08f70579b 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -87,12 +87,22 @@ def _run_worker( # Received poison pill, no more tasks to run return - if not isinstance(workload, workloads.ExecuteTask): - raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") - # Decrement this as soon as we pick up a message off the queue with unread_messages: unread_messages.value -= 1 + + if isinstance(workload, workloads.TestConnection): + # Handle connection test workloads + try: + _execute_connection_test(log, workload, team_conf) + # Connection tests report their own results, no need to put in output queue + except Exception: + log.exception("Connection test execution failed") + continue + + if not isinstance(workload, workloads.ExecuteTask): + raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") + key = None if ti := getattr(workload, "ti", None): key = ti.key @@ -108,6 +118,39 @@ def _run_worker( output.put((key, TaskInstanceState.FAILED, e)) +def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None: + """Execute a connection test workload.""" + from airflow.executors.connection_test_runner import ( + execute_connection_test_workload, + report_connection_test_result, + ) + + team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" + setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: connection_test:{workload.request_id}", log) + log.info("Executing connection test", request_id=workload.request_id, conn_type=workload.conn_type) + + try: + success, message = execute_connection_test_workload(workload) + except Exception as e: + log.exception("Connection test execution failed") + success = False + message = f"Execution error: {e}" + + base_url = team_conf.get("api", "base_url", fallback="/") + if base_url.startswith("/"): + api_port = team_conf.get("api", "port", fallback="8080") + base_url = f"http://localhost:{api_port}{base_url}" + server = team_conf.get("core", "execution_api_server_url", fallback=f"{base_url.rstrip('/')}/execution/") + + report_connection_test_result( + request_id=workload.request_id, + success=success, + message=message, + server_url=server, + token=workload.token, + ) + + def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> None: """ Execute command received and stores result state in queue. @@ -123,9 +166,10 @@ def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> No setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: {workload.ti.id}", log) base_url = team_conf.get("api", "base_url", fallback="/") - # If it's a relative URL, use localhost:8080 as the default + # If it's a relative URL, use localhost with configured port as the default if base_url.startswith("/"): - base_url = f"http://localhost:8080{base_url}" + api_port = team_conf.get("api", "port", fallback="8080") + base_url = f"http://localhost:{api_port}{base_url}" default_execution_api_server = f"{base_url.rstrip('/')}/execution/" # This will return the exit code of the task process, but we don't care about that, just if the @@ -299,10 +343,24 @@ def end(self) -> None: def terminate(self): """Terminate the executor is not doing anything.""" - def _process_workloads(self, workloads): - for workload in workloads: + def _process_workloads(self, workloads_to_process): + from airflow.executors import workloads as workload_types + + for workload in workloads_to_process: self.activity_queue.put(workload) - del self.queued_tasks[workload.ti.key] + if isinstance(workload, workload_types.ExecuteTask): + del self.queued_tasks[workload.ti.key] + # Connection test workloads are not tracked in queued_tasks + with self._unread_messages: - self._unread_messages.value += len(workloads) + self._unread_messages.value += len(workloads_to_process) self._check_workers() + + # Also process any queued connection tests + if hasattr(self, "queued_connection_tests"): + while self.queued_connection_tests: + conn_test = self.queued_connection_tests.popleft() + self.activity_queue.put(conn_test) + with self._unread_messages: + self._unread_messages.value += 1 + self._check_workers() diff --git a/airflow-core/src/airflow/executors/workloads.py b/airflow-core/src/airflow/executors/workloads.py index 7cf1aae60ff21..42750aa4098ee 100644 --- a/airflow-core/src/airflow/executors/workloads.py +++ b/airflow-core/src/airflow/executors/workloads.py @@ -34,7 +34,7 @@ from airflow.models.taskinstancekey import TaskInstanceKey -__all__ = ["All", "ExecuteTask", "ExecuteCallback"] +__all__ = ["All", "ExecuteTask", "ExecuteCallback", "RunTrigger", "TestConnection"] log = structlog.get_logger(__name__) @@ -204,7 +204,48 @@ class RunTrigger(BaseModel): type: Literal["RunTrigger"] = Field(init=False, default="RunTrigger") +class TestConnection(BaseWorkload): + """ + Execute a connection test on a worker. + + This workload tests connectivity to an external system by decrypting the + connection URI and calling the connection's test_connection() method. + The result is reported back to the API server. + """ + + request_id: str + """Unique identifier for this connection test request""" + + encrypted_connection_uri: str + """Fernet-encrypted connection URI""" + + conn_type: str + """Connection type (e.g., 'postgres', 'mysql')""" + + timeout: int = 60 + """Timeout in seconds for the connection test""" + + type: Literal["TestConnection"] = Field(init=False, default="TestConnection") + + @classmethod + def make( + cls, + request_id: str, + encrypted_connection_uri: str, + conn_type: str, + timeout: int = 60, + generator: JWTGenerator | None = None, + ) -> TestConnection: + return cls( + request_id=request_id, + encrypted_connection_uri=encrypted_connection_uri, + conn_type=conn_type, + timeout=timeout, + token=cls.generate_token(request_id, generator), + ) + + All = Annotated[ - ExecuteTask | RunTrigger, + ExecuteTask | ExecuteCallback | RunTrigger | TestConnection, Field(discriminator="type"), ] diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 917b10efce435..f250d3bd8b3d9 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -70,6 +70,7 @@ ) from airflow.models.backfill import Backfill from airflow.models.callback import Callback +from airflow.models.connection_test import ConnectionTestRequest from airflow.models.dag import DagModel, get_next_data_interval from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -885,6 +886,41 @@ def _process_task_event_logs(log_records: deque[Log], session: Session): objects = (log_records.popleft() for _ in range(len(log_records))) session.bulk_save_objects(objects=objects, preserve_order=False) + def _dispatch_connection_tests(self, session: Session) -> int: + """Dispatch pending connection test requests to workers.""" + if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": + return 0 + + pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=10) + if not pending_requests: + return 0 + + dispatched = 0 + for request in pending_requests: + try: + if self.job.executors: + workload = workloads.TestConnection.make( + request_id=request.id, + encrypted_connection_uri=request.encrypted_connection_uri, + conn_type=request.conn_type, + timeout=request.timeout, + generator=self.job.executors[0].jwt_generator, + ) + request.mark_running(self.job.hostname or "scheduler") + self.job.executors[0].queue_workload(workload, session=session) + dispatched += 1 + self.log.info("Dispatched connection test %s to executor", request.id) + else: + self.log.warning("No executors available to dispatch connection test %s", request.id) + request.mark_failed("No executors available") + except Exception: + self.log.exception("Failed to dispatch connection test %s", request.id) + request.mark_failed("Failed to dispatch to executor") + + session.commit() + + return dispatched + @staticmethod def _is_metrics_enabled(): return any( @@ -1557,6 +1593,9 @@ def _run_scheduler_loop(self) -> None: ): deadline.handle_miss(session) + # Dispatch pending connection test requests to workers + self._dispatch_connection_tests(session) + # Heartbeat the scheduler periodically perform_heartbeat( job=self.job, heartbeat_callback=self.heartbeat_callback, only_if_necessary=True diff --git a/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py b/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py new file mode 100644 index 0000000000000..4bf57b14ad15e --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py @@ -0,0 +1,69 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add connection_test_request table for async connection testing on workers. + +Revision ID: 9882c124ea54 +Revises: e79fc784f145 +Create Date: 2026-01-14 10:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "9882c124ea54" +down_revision = "e79fc784f145" +branch_labels = None +depends_on = None +airflow_version = "3.2.0" + + +def upgrade(): + """Create connection_test_request table.""" + op.create_table( + "connection_test_request", + sa.Column("id", sa.String(36), nullable=False), + sa.Column("state", sa.String(10), nullable=False), + sa.Column("encrypted_connection_uri", sa.Text(), nullable=False), + sa.Column("conn_type", sa.String(500), nullable=False), + sa.Column("result_status", sa.Boolean(), nullable=True), + sa.Column("result_message", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("completed_at", sa.DateTime(), nullable=True), + sa.Column("timeout", sa.Integer(), nullable=False, default=60), + sa.Column("worker_hostname", sa.String(500), nullable=True), + sa.PrimaryKeyConstraint("id", name=op.f("connection_test_request_pkey")), + ) + op.create_index( + "idx_connection_test_state_created", + "connection_test_request", + ["state", "created_at"], + unique=False, + ) + + +def downgrade(): + """Drop connection_test_request table.""" + op.drop_index("idx_connection_test_state_created", table_name="connection_test_request") + op.drop_table("connection_test_request") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index 39a5a7350759e..dfe3e210333a5 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -100,6 +100,7 @@ def __getattr__(name): "BaseXCom": "airflow.sdk.bases.xcom", "Callback": "airflow.models.callback", "Connection": "airflow.models.connection", + "ConnectionTestRequest": "airflow.models.connection_test", "DagBag": "airflow.dag_processing.dagbag", "DagModel": "airflow.models.dag", "DagRun": "airflow.models.dagrun", @@ -130,6 +131,7 @@ def __getattr__(name): from airflow.models.base import ID_LEN, Base from airflow.models.callback import Callback from airflow.models.connection import Connection + from airflow.models.connection_test import ConnectionTestRequest from airflow.models.dag import DagModel, DagTag from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py new file mode 100644 index 0000000000000..c54fd98335f0a --- /dev/null +++ b/airflow-core/src/airflow/models/connection_test.py @@ -0,0 +1,189 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Connection test request model for async connection testing on workers.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING + +import uuid6 +from sqlalchemy import Boolean, Index, Integer, String, Text, select +from sqlalchemy.orm import Mapped + +from airflow._shared.timezones import timezone +from airflow.models.base import Base +from airflow.utils.sqlalchemy import UtcDateTime, mapped_column + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + +class ConnectionTestState(str, Enum): + """All possible states of a connection test request.""" + + PENDING = "pending" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + def __str__(self) -> str: + return self.value + + +TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) + +# Valid state transitions +VALID_STATE_TRANSITIONS: dict[ConnectionTestState, frozenset[ConnectionTestState]] = { + ConnectionTestState.PENDING: frozenset((ConnectionTestState.RUNNING, ConnectionTestState.FAILED)), + ConnectionTestState.RUNNING: frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)), + ConnectionTestState.SUCCESS: frozenset(), # Terminal state, no transitions allowed + ConnectionTestState.FAILED: frozenset(), # Terminal state, no transitions allowed +} + + +class ConnectionTestRequest(Base): + """ + Stores connection test requests for asynchronous execution on workers. + + This model supports moving the test_connection functionality from the API server + to workers for improved security and network isolation. + """ + + __tablename__ = "connection_test_request" + + id: Mapped[str] = mapped_column(String(36), primary_key=True) + + # State: PENDING -> RUNNING -> SUCCESS/FAILED + state: Mapped[str] = mapped_column(String(10), nullable=False, default=ConnectionTestState.PENDING.value) + + # Encrypted connection URI (using Fernet encryption) + encrypted_connection_uri: Mapped[str] = mapped_column(Text, nullable=False) + + # Connection type (needed to instantiate the hook on the worker) + conn_type: Mapped[str] = mapped_column(String(500), nullable=False) + + # Result fields - populated when test completes + result_status: Mapped[bool | None] = mapped_column(Boolean, nullable=True) + result_message: Mapped[str | None] = mapped_column(Text, nullable=True) + + # Timestamps + created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + started_at: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) + completed_at: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) + + # Timeout in seconds (default 60s) + timeout: Mapped[int] = mapped_column(Integer, default=60, nullable=False) + + # Which worker picked up this request + worker_hostname: Mapped[str | None] = mapped_column(String(500), nullable=True) + + __table_args__ = (Index("idx_connection_test_state_created", "state", "created_at"),) + + def __repr__(self) -> str: + return f"" + + def _validate_state_transition(self, new_state: ConnectionTestState) -> None: + """ + Validate that a state transition is allowed. + + :param new_state: The target state + :raises ValueError: If the transition is not allowed + """ + current_state = ConnectionTestState(self.state) + allowed_transitions = VALID_STATE_TRANSITIONS.get(current_state, frozenset()) + if new_state not in allowed_transitions: + raise ValueError( + f"Invalid state transition from {current_state.value!r} to {new_state.value!r}. " + f"Allowed transitions: {[s.value for s in allowed_transitions]}" + ) + + @classmethod + def create_request( + cls, + encrypted_connection_uri: str, + conn_type: str, + timeout: int = 60, + session: Session | None = None, + ) -> ConnectionTestRequest: + """ + Create a new connection test request. + + :param encrypted_connection_uri: The Fernet-encrypted connection URI + :param conn_type: The connection type (e.g., 'postgres', 'mysql') + :param timeout: Timeout in seconds for the test + :param session: Optional SQLAlchemy session to add the request to + :return: The created ConnectionTestRequest + """ + request = cls( + id=str(uuid6.uuid7()), + encrypted_connection_uri=encrypted_connection_uri, + conn_type=conn_type, + timeout=timeout, + state=ConnectionTestState.PENDING.value, + ) + if session: + session.add(request) + return request + + def mark_running(self, worker_hostname: str) -> None: + """Mark the request as running on a specific worker.""" + self._validate_state_transition(ConnectionTestState.RUNNING) + self.state = ConnectionTestState.RUNNING.value + self.worker_hostname = worker_hostname + self.started_at = timezone.utcnow() + + def mark_success(self, message: str) -> None: + """Mark the request as successfully completed.""" + self._validate_state_transition(ConnectionTestState.SUCCESS) + self.state = ConnectionTestState.SUCCESS.value + self.result_status = True + self.result_message = message + self.completed_at = timezone.utcnow() + + def mark_failed(self, message: str) -> None: + """Mark the request as failed.""" + self._validate_state_transition(ConnectionTestState.FAILED) + self.state = ConnectionTestState.FAILED.value + self.result_status = False + self.result_message = message + self.completed_at = timezone.utcnow() + + @property + def is_terminal(self) -> bool: + """Check if the request is in a terminal state.""" + return ConnectionTestState(self.state) in TERMINAL_STATES + + @classmethod + def get_pending_requests(cls, session: Session, limit: int = 10) -> list[ConnectionTestRequest]: + """ + Get pending connection test requests for worker processing. + + :param session: SQLAlchemy session + :param limit: Maximum number of requests to return + :return: List of pending ConnectionTestRequest objects + """ + return list( + session.scalars( + select(cls) + .where(cls.state == ConnectionTestState.PENDING.value) + .order_by(cls.created_at) + .limit(limit) + .with_for_update(skip_locked=True) + ) + ) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 4157583cd3f94..98335f093d4d4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -107,6 +107,12 @@ export const UseBackfillServiceListBackfillsUiKeyFn = ({ active, dagId, limit, o offset?: number; orderBy?: string[]; } = {}, queryKey?: Array) => [useBackfillServiceListBackfillsUiKey, ...(queryKey ?? [{ active, dagId, limit, offset, orderBy }])]; +export type ConnectionServiceGetConnectionTestStatusDefaultResponse = Awaited>; +export type ConnectionServiceGetConnectionTestStatusQueryResult = UseQueryResult; +export const useConnectionServiceGetConnectionTestStatusKey = "ConnectionServiceGetConnectionTestStatus"; +export const UseConnectionServiceGetConnectionTestStatusKeyFn = ({ requestId }: { + requestId: string; +}, queryKey?: Array) => [useConnectionServiceGetConnectionTestStatusKey, ...(queryKey ?? [{ requestId }])]; export type ConnectionServiceGetConnectionDefaultResponse = Awaited>; export type ConnectionServiceGetConnectionQueryResult = UseQueryResult; export const useConnectionServiceGetConnectionKey = "ConnectionServiceGetConnection"; @@ -883,8 +889,8 @@ export type AssetServiceCreateAssetEventMutationResult = Awaited>; export type BackfillServiceCreateBackfillMutationResult = Awaited>; export type BackfillServiceCreateBackfillDryRunMutationResult = Awaited>; -export type ConnectionServicePostConnectionMutationResult = Awaited>; export type ConnectionServiceTestConnectionMutationResult = Awaited>; +export type ConnectionServicePostConnectionMutationResult = Awaited>; export type ConnectionServiceCreateDefaultConnectionsMutationResult = Awaited>; export type DagRunServiceClearDagRunMutationResult = Awaited>; export type DagRunServiceTriggerDagRunMutationResult = Awaited>; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index 9501e4bb1a17c..87f06bac3cf52 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -197,6 +197,19 @@ export const ensureUseBackfillServiceListBackfillsUiData = (queryClient: QueryCl orderBy?: string[]; } = {}) => queryClient.ensureQueryData({ queryKey: Common.UseBackfillServiceListBackfillsUiKeyFn({ active, dagId, limit, offset, orderBy }), queryFn: () => BackfillService.listBackfillsUi({ active, dagId, limit, offset, orderBy }) }); /** +* Get Connection Test Status +* Get the status of a connection test request. +* +* Poll this endpoint to check if a connection test has completed and retrieve the result. +* @param data The data for the request. +* @param data.requestId +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const ensureUseConnectionServiceGetConnectionTestStatusData = (queryClient: QueryClient, { requestId }: { + requestId: string; +}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ requestId }), queryFn: () => ConnectionService.getConnectionTestStatus({ requestId }) }); +/** * Get Connection * Get a connection entry. * @param data The data for the request. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index bbd67f6be6259..dba4f87ccf518 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -197,6 +197,19 @@ export const prefetchUseBackfillServiceListBackfillsUi = (queryClient: QueryClie orderBy?: string[]; } = {}) => queryClient.prefetchQuery({ queryKey: Common.UseBackfillServiceListBackfillsUiKeyFn({ active, dagId, limit, offset, orderBy }), queryFn: () => BackfillService.listBackfillsUi({ active, dagId, limit, offset, orderBy }) }); /** +* Get Connection Test Status +* Get the status of a connection test request. +* +* Poll this endpoint to check if a connection test has completed and retrieve the result. +* @param data The data for the request. +* @param data.requestId +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const prefetchUseConnectionServiceGetConnectionTestStatus = (queryClient: QueryClient, { requestId }: { + requestId: string; +}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ requestId }), queryFn: () => ConnectionService.getConnectionTestStatus({ requestId }) }); +/** * Get Connection * Get a connection entry. * @param data The data for the request. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index f55804064d885..9eea89efb7f42 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -197,6 +197,19 @@ export const useBackfillServiceListBackfillsUi = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseBackfillServiceListBackfillsUiKeyFn({ active, dagId, limit, offset, orderBy }, queryKey), queryFn: () => BackfillService.listBackfillsUi({ active, dagId, limit, offset, orderBy }) as TData, ...options }); /** +* Get Connection Test Status +* Get the status of a connection test request. +* +* Poll this endpoint to check if a connection test has completed and retrieve the result. +* @param data The data for the request. +* @param data.requestId +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServiceGetConnectionTestStatus = = unknown[]>({ requestId }: { + requestId: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ requestId }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ requestId }) as TData, ...options }); +/** * Get Connection * Get a connection entry. * @param data The data for the request. @@ -1730,35 +1743,37 @@ export const useBackfillServiceCreateBackfillDryRun = ({ mutationFn: ({ requestBody }) => BackfillService.createBackfillDryRun({ requestBody }) as unknown as Promise, ...options }); /** -* Post Connection -* Create connection entry. +* Test Connection +* Queue a connection test for asynchronous execution on a worker. +* +* This endpoint queues the connection test request for execution on a worker node, +* which provides better security isolation (workers run in ephemeral environments) +* and network accessibility (workers can reach external systems that API servers cannot). +* +* Returns a request_id that can be used to poll for the test result via GET /connections/test/{request_id}. * @param data The data for the request. * @param data.requestBody -* @returns ConnectionResponse Successful Response +* @returns ConnectionTestQueuedResponse Successful Response * @throws ApiError */ -export const useConnectionServicePostConnection = (options?: Omit(options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ requestBody }) => ConnectionService.postConnection({ requestBody }) as unknown as Promise, ...options }); +}, TContext>({ mutationFn: ({ requestBody }) => ConnectionService.testConnection({ requestBody }) as unknown as Promise, ...options }); /** -* Test Connection -* Test an API connection. -* -* This method first creates an in-memory transient conn_id & exports that to an env var, -* as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found. -* It also deletes the conn id env connection after the test. +* Post Connection +* Create connection entry. * @param data The data for the request. * @param data.requestBody -* @returns ConnectionTestResponse Successful Response +* @returns ConnectionResponse Successful Response * @throws ApiError */ -export const useConnectionServiceTestConnection = (options?: Omit(options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ requestBody }) => ConnectionService.testConnection({ requestBody }) as unknown as Promise, ...options }); +}, TContext>({ mutationFn: ({ requestBody }) => ConnectionService.postConnection({ requestBody }) as unknown as Promise, ...options }); /** * Create Default Connections * Create default connections. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index 9c38554e4ab0d..ae670b5156b5c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -197,6 +197,19 @@ export const useBackfillServiceListBackfillsUiSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseBackfillServiceListBackfillsUiKeyFn({ active, dagId, limit, offset, orderBy }, queryKey), queryFn: () => BackfillService.listBackfillsUi({ active, dagId, limit, offset, orderBy }) as TData, ...options }); /** +* Get Connection Test Status +* Get the status of a connection test request. +* +* Poll this endpoint to check if a connection test has completed and retrieve the result. +* @param data The data for the request. +* @param data.requestId +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServiceGetConnectionTestStatusSuspense = = unknown[]>({ requestId }: { + requestId: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ requestId }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ requestId }) as TData, ...options }); +/** * Get Connection * Get a connection entry. * @param data The data for the request. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 06eae729f48b0..f7584b3c135b8 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1691,11 +1691,15 @@ export const $ConnectionResponse = { description: 'Connection serializer for responses.' } as const; -export const $ConnectionTestResponse = { +export const $ConnectionTestQueuedResponse = { properties: { - status: { - type: 'boolean', - title: 'Status' + request_id: { + type: 'string', + title: 'Request Id' + }, + state: { + type: 'string', + title: 'State' }, message: { type: 'string', @@ -1703,9 +1707,77 @@ export const $ConnectionTestResponse = { } }, type: 'object', - required: ['status', 'message'], - title: 'ConnectionTestResponse', - description: 'Connection Test serializer for responses.' + required: ['request_id', 'state', 'message'], + title: 'ConnectionTestQueuedResponse', + description: 'Response when a connection test is queued for async execution.' +} as const; + +export const $ConnectionTestStatusResponse = { + properties: { + request_id: { + type: 'string', + title: 'Request Id' + }, + state: { + type: 'string', + title: 'State' + }, + result_status: { + anyOf: [ + { + type: 'boolean' + }, + { + type: 'null' + } + ], + title: 'Result Status' + }, + result_message: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Result Message' + }, + created_at: { + type: 'string', + format: 'date-time', + title: 'Created At' + }, + started_at: { + anyOf: [ + { + type: 'string', + format: 'date-time' + }, + { + type: 'null' + } + ], + title: 'Started At' + }, + completed_at: { + anyOf: [ + { + type: 'string', + format: 'date-time' + }, + { + type: 'null' + } + ], + title: 'Completed At' + } + }, + type: 'object', + required: ['request_id', 'state', 'result_status', 'result_message', 'created_at', 'started_at', 'completed_at'], + title: 'ConnectionTestStatusResponse', + description: 'Response with the full status of a connection test request.' } as const; export const $CreateAssetEventsBody = { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 63e0f04f361a4..8a6ca908d46d6 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, TestConnectionData, TestConnectionResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; +import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, TestConnectionData, TestConnectionResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; export class AssetService { /** @@ -613,6 +613,60 @@ export class BackfillService { } export class ConnectionService { + /** + * Test Connection + * Queue a connection test for asynchronous execution on a worker. + * + * This endpoint queues the connection test request for execution on a worker node, + * which provides better security isolation (workers run in ephemeral environments) + * and network accessibility (workers can reach external systems that API servers cannot). + * + * Returns a request_id that can be used to poll for the test result via GET /connections/test/{request_id}. + * @param data The data for the request. + * @param data.requestBody + * @returns ConnectionTestQueuedResponse Successful Response + * @throws ApiError + */ + public static testConnection(data: TestConnectionData): CancelablePromise { + return __request(OpenAPI, { + method: 'POST', + url: '/api/v2/connections/test', + body: data.requestBody, + mediaType: 'application/json', + errors: { + 401: 'Unauthorized', + 403: 'Forbidden', + 422: 'Validation Error' + } + }); + } + + /** + * Get Connection Test Status + * Get the status of a connection test request. + * + * Poll this endpoint to check if a connection test has completed and retrieve the result. + * @param data The data for the request. + * @param data.requestId + * @returns ConnectionTestStatusResponse Successful Response + * @throws ApiError + */ + public static getConnectionTestStatus(data: GetConnectionTestStatusData): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v2/connections/test/{request_id}', + path: { + request_id: data.requestId + }, + errors: { + 401: 'Unauthorized', + 403: 'Forbidden', + 404: 'Not Found', + 422: 'Validation Error' + } + }); + } + /** * Delete Connection * Delete a connection entry. @@ -768,32 +822,6 @@ export class ConnectionService { }); } - /** - * Test Connection - * Test an API connection. - * - * This method first creates an in-memory transient conn_id & exports that to an env var, - * as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found. - * It also deletes the conn id env connection after the test. - * @param data The data for the request. - * @param data.requestBody - * @returns ConnectionTestResponse Successful Response - * @throws ApiError - */ - public static testConnection(data: TestConnectionData): CancelablePromise { - return __request(OpenAPI, { - method: 'POST', - url: '/api/v2/connections/test', - body: data.requestBody, - mediaType: 'application/json', - errors: { - 401: 'Unauthorized', - 403: 'Forbidden', - 422: 'Validation Error' - } - }); - } - /** * Create Default Connections * Create default connections. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 5523d61b8182a..32743c7d49f66 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -496,13 +496,27 @@ export type ConnectionResponse = { }; /** - * Connection Test serializer for responses. + * Response when a connection test is queued for async execution. */ -export type ConnectionTestResponse = { - status: boolean; +export type ConnectionTestQueuedResponse = { + request_id: string; + state: string; message: string; }; +/** + * Response with the full status of a connection test request. + */ +export type ConnectionTestStatusResponse = { + request_id: string; + state: string; + result_status: boolean | null; + result_message: string | null; + created_at: string; + started_at: string | null; + completed_at: string | null; +}; + /** * Create asset events request. */ @@ -2284,6 +2298,18 @@ export type ListBackfillsUiData = { export type ListBackfillsUiResponse = BackfillCollectionResponse; +export type TestConnectionData = { + requestBody: ConnectionBody; +}; + +export type TestConnectionResponse = ConnectionTestQueuedResponse; + +export type GetConnectionTestStatusData = { + requestId: string; +}; + +export type GetConnectionTestStatusResponse = ConnectionTestStatusResponse; + export type DeleteConnectionData = { connectionId: string; }; @@ -2331,12 +2357,6 @@ export type BulkConnectionsData = { export type BulkConnectionsResponse = BulkResponse; -export type TestConnectionData = { - requestBody: ConnectionBody; -}; - -export type TestConnectionResponse = ConnectionTestResponse; - export type CreateDefaultConnectionsResponse = void; export type HookMetaDataResponse = Array; @@ -4093,6 +4113,56 @@ export type $OpenApiTs = { }; }; }; + '/api/v2/connections/test': { + post: { + req: TestConnectionData; + res: { + /** + * Successful Response + */ + 200: ConnectionTestQueuedResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; + '/api/v2/connections/test/{request_id}': { + get: { + req: GetConnectionTestStatusData; + res: { + /** + * Successful Response + */ + 200: ConnectionTestStatusResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; '/api/v2/connections/{connection_id}': { delete: { req: DeleteConnectionData; @@ -4247,29 +4317,6 @@ export type $OpenApiTs = { }; }; }; - '/api/v2/connections/test': { - post: { - req: TestConnectionData; - res: { - /** - * Successful Response - */ - 200: ConnectionTestResponse; - /** - * Unauthorized - */ - 401: HTTPExceptionResponse; - /** - * Forbidden - */ - 403: HTTPExceptionResponse; - /** - * Validation Error - */ - 422: HTTPValidationError; - }; - }; - }; '/api/v2/connections/defaults': { post: { res: { diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/admin.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/admin.json index ad231ceb31059..c3b264b50b465 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/admin.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/admin.json @@ -50,11 +50,21 @@ "test": "Test Connection", "testDisabled": "Test connection feature is disabled. Please contact an administrator to enable it.", "testError": { + "description": "The connection test failed. Check the connection parameters.", "title": "Test Connection Failed" }, + "testQueued": { + "description": "The connection test has been queued and will be executed on a worker.", + "title": "Connection Test Queued" + }, "testSuccess": { + "description": "The connection test completed successfully.", "title": "Test Connection Successful" }, + "testTimeout": { + "description": "The connection test took too long to complete. Please try again.", + "title": "Connection Test Timeout" + }, "typeMeta": { "error": "Failed to retrieve Connection Type Meta", "standardFields": { diff --git a/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx b/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx index 7b27877d87037..68e9b888fade7 100644 --- a/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx @@ -75,8 +75,7 @@ const TestConnectionButton = ({ connection }: Props) => { actionName={ option === "Enabled" ? translate("connections.test") : translate("connections.testDisabled") } - disabled={option === "Disabled"} - display={option === "Hidden" ? "none" : "flex"} + display="flex" icon={icon} loading={isPending} onClick={() => { diff --git a/airflow-core/src/airflow/ui/src/queries/useTestConnection.ts b/airflow-core/src/airflow/ui/src/queries/useTestConnection.ts index 5a810a223ffb6..ffc77ab6a274f 100644 --- a/airflow-core/src/airflow/ui/src/queries/useTestConnection.ts +++ b/airflow-core/src/airflow/ui/src/queries/useTestConnection.ts @@ -16,39 +16,120 @@ * specific language governing permissions and limitations * under the License. */ +import { useMutation, useQuery } from "@tanstack/react-query"; +import { useCallback, useEffect, useRef, useState } from "react"; import type { Dispatch, SetStateAction } from "react"; import { useTranslation } from "react-i18next"; -import { useConnectionServiceTestConnection } from "openapi/queries"; -import type { ConnectionTestResponse } from "openapi/requests/types.gen"; +import { ConnectionService } from "openapi/requests"; +import type { ConnectionTestStatusResponse } from "openapi/requests/types.gen"; import { toaster } from "src/components/ui"; +const POLL_INTERVAL = 1000; +const MAX_POLL_TIME_MS = 60_000; + export const useTestConnection = (setConnected: Dispatch>) => { const { t: translate } = useTranslation("admin"); + const [requestId, setRequestId] = useState(undefined); + const [isPolling, setIsPolling] = useState(false); + const pollStartTimeRef = useRef(undefined); - const onSuccess = (res: ConnectionTestResponse) => { - setConnected(res.status); - if (res.status) { + const onQueueSuccess = useCallback( + (res: { request_id: string }) => { + setRequestId(res.request_id); + setIsPolling(true); + pollStartTimeRef.current = Date.now(); toaster.create({ - description: res.message, + description: translate("connections.testQueued.description"), + title: translate("connections.testQueued.title"), + type: "info", + }); + }, + [translate], + ); + + const onQueueError = useCallback(() => { + setConnected(false); + setIsPolling(false); + pollStartTimeRef.current = undefined; + toaster.create({ + description: translate("connections.testError.description"), + title: translate("connections.testError.title"), + type: "error", + }); + }, [setConnected, translate]); + + const queueMutation = useMutation({ + mutationFn: ConnectionService.testConnection, + onError: onQueueError, + onSuccess: onQueueSuccess, + }); + + const statusQuery = useQuery({ + enabled: isPolling && requestId !== undefined, + queryFn: () => ConnectionService.getConnectionTestStatus({ requestId: requestId as string }), + queryKey: ["connectionTestStatus", requestId], + refetchInterval: (query) => { + const { data } = query.state; + + if (data?.state === "success" || data?.state === "failed") { + return false; + } + + if ( + pollStartTimeRef.current !== undefined && + Date.now() - pollStartTimeRef.current > MAX_POLL_TIME_MS + ) { + return false; + } + + return POLL_INTERVAL; + }, + refetchOnWindowFocus: false, + }); + + useEffect(() => { + if (!statusQuery.data || !isPolling) { + return; + } + + const { data } = statusQuery; + + if (data.state === "success") { + setIsPolling(false); + pollStartTimeRef.current = undefined; + setConnected(data.result_status ?? false); + toaster.create({ + description: data.result_message ?? translate("connections.testSuccess.description"), title: translate("connections.testSuccess.title"), type: "success", }); - } else { + } else if (data.state === "failed") { + setIsPolling(false); + pollStartTimeRef.current = undefined; + setConnected(false); toaster.create({ - description: res.message, + description: data.result_message ?? translate("connections.testError.description"), title: translate("connections.testError.title"), type: "error", }); + } else if ( + pollStartTimeRef.current !== undefined && + Date.now() - pollStartTimeRef.current > MAX_POLL_TIME_MS + ) { + setIsPolling(false); + pollStartTimeRef.current = undefined; + setConnected(false); + toaster.create({ + description: translate("connections.testTimeout.description"), + title: translate("connections.testTimeout.title"), + type: "error", + }); } - }; + }, [statusQuery, isPolling, setConnected, translate]); - const onError = () => { - setConnected(false); + return { + isPending: queueMutation.isPending || isPolling, + mutate: queueMutation.mutate, }; - - return useConnectionServiceTestConnection({ - onError, - onSuccess, - }); }; diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index d444689e26b56..ba3060dc41345 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -112,7 +112,7 @@ class MappedClassProtocol(Protocol): "3.0.0": "29ce7909c52b", "3.0.3": "fe199e1abd77", "3.1.0": "cc92b33c6709", - "3.2.0": "e79fc784f145", + "3.2.0": "9882c124ea54", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index d2c48f945d480..a5356c624a326 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -17,7 +17,6 @@ from __future__ import annotations import os -from importlib.metadata import PackageNotFoundError, metadata from unittest import mock import pytest @@ -25,13 +24,11 @@ from sqlalchemy.orm import Session from airflow.models import Connection -from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.session import NEW_SESSION, provide_session from tests_common.test_utils.api_fastapi import _check_last_log from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_connections, clear_db_logs, clear_test_connections -from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker pytestmark = pytest.mark.db_test @@ -935,31 +932,63 @@ def test_patch_should_response_200_redacted_password( _check_last_log(session, dag_id=None, event="patch_connection", logical_date=None, check_masked=True) -class TestConnection(TestConnectionEndpoint): - def setup_method(self): - try: - metadata("apache-airflow-providers-sqlite") - except PackageNotFoundError: - pytest.skip("The SQlite distribution package is not installed.") +class TestQueueConnectionTest(TestConnectionEndpoint): + """Test the async connection test queue endpoint POST /connections/test.""" + + @pytest.fixture(autouse=True) + def setup_connection_test(self): + """Clean up connection test requests after tests.""" + from tests_common.test_utils.db import clear_db_connection_tests + + yield + clear_db_connection_tests() @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) @pytest.mark.parametrize( - ("body", "message"), + "body", [ - ({"connection_id": TEST_CONN_ID, "conn_type": "sqlite"}, "Connection successfully tested"), - ( - {"connection_id": TEST_CONN_ID, "conn_type": "fs", "extra": '{"path": "/"}'}, - "Path / is existing.", - ), + {"connection_id": TEST_CONN_ID, "conn_type": "sqlite"}, + {"connection_id": TEST_CONN_ID, "conn_type": "fs", "extra": '{"path": "/"}'}, + {"connection_id": TEST_CONN_ID, "conn_type": "postgres", "host": "localhost", "port": 5432}, ], ) - def test_should_respond_200(self, test_client, body, message): + def test_should_queue_connection_test(self, test_client, body, session): + """Test that POST /connections/test queues a connection test request.""" + from airflow.models.connection_test import ConnectionTestRequest + response = test_client.post("/connections/test", json=body) assert response.status_code == 200 - assert response.json() == { - "status": True, - "message": message, + + data = response.json() + assert "request_id" in data + assert data["state"] == "pending" + assert "message" in data + + # Verify request was persisted in database + request = session.get(ConnectionTestRequest, data["request_id"]) + assert request is not None + assert request.conn_type == body["conn_type"] + assert request.state == "pending" + assert request.encrypted_connection_uri is not None + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_connection_uri_is_encrypted(self, test_client, session): + """Test that connection URI is encrypted when stored.""" + from airflow.models.connection_test import ConnectionTestRequest + + body = { + "connection_id": TEST_CONN_ID, + "conn_type": "postgres", + "host": "myhost.example.com", + "password": "secret123", } + response = test_client.post("/connections/test", json=body) + assert response.status_code == 200 + + request = session.get(ConnectionTestRequest, response.json()["request_id"]) + # Encrypted URI should not contain plaintext password + assert "secret123" not in request.encrypted_connection_uri + assert "myhost.example.com" not in request.encrypted_connection_uri def test_should_respond_401(self, unauthenticated_test_client): response = unauthenticated_test_client.post( @@ -973,19 +1002,6 @@ def test_should_respond_403(self, unauthorized_test_client): ) assert response.status_code == 403 - @skip_if_force_lowest_dependencies_marker - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - @pytest.mark.parametrize( - "body", - [ - {"connection_id": TEST_CONN_ID, "conn_type": "sqlite"}, - {"connection_id": TEST_CONN_ID, "conn_type": "ftp"}, - ], - ) - def test_connection_env_is_cleaned_after_run(self, test_client, body): - test_client.post("/connections/test", json=body) - assert not any([key.startswith(CONN_ENV_PREFIX) for key in os.environ.keys()]) - @pytest.mark.parametrize( "body", [ @@ -1002,6 +1018,122 @@ def test_should_respond_403_by_default(self, test_client, body): } +class TestGetConnectionTestStatus(TestConnectionEndpoint): + """Test the connection test status endpoint GET /connections/test/{request_id}.""" + + @pytest.fixture(autouse=True) + def setup_connection_test(self): + """Clean up connection test requests after tests.""" + from tests_common.test_utils.db import clear_db_connection_tests + + yield + clear_db_connection_tests() + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_pending_status(self, test_client, session): + """Test getting status of a pending connection test.""" + from airflow.models.connection_test import ConnectionTestRequest + + # Create a pending request + request = ConnectionTestRequest.create_request( + encrypted_connection_uri="gAAAAABfakedata...", + conn_type="postgres", + session=session, + ) + session.commit() + + response = test_client.get(f"/connections/test/{request.id}") + assert response.status_code == 200 + + data = response.json() + assert data["request_id"] == request.id + assert data["state"] == "pending" + assert data["result_status"] is None + assert data["result_message"] is None + assert data["created_at"] is not None + assert data["started_at"] is None + assert data["completed_at"] is None + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_running_status(self, test_client, session): + """Test getting status of a running connection test.""" + from airflow.models.connection_test import ConnectionTestRequest + + request = ConnectionTestRequest.create_request( + encrypted_connection_uri="gAAAAABfakedata...", + conn_type="postgres", + session=session, + ) + request.mark_running("worker-1.example.com") + session.commit() + + response = test_client.get(f"/connections/test/{request.id}") + assert response.status_code == 200 + + data = response.json() + assert data["state"] == "running" + assert data["started_at"] is not None + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_success_status(self, test_client, session): + """Test getting status of a successful connection test.""" + from airflow.models.connection_test import ConnectionTestRequest + + request = ConnectionTestRequest.create_request( + encrypted_connection_uri="gAAAAABfakedata...", + conn_type="postgres", + session=session, + ) + request.mark_running("worker-1") + request.mark_success("Connection successfully tested") + session.commit() + + response = test_client.get(f"/connections/test/{request.id}") + assert response.status_code == 200 + + data = response.json() + assert data["state"] == "success" + assert data["result_status"] is True + assert data["result_message"] == "Connection successfully tested" + assert data["completed_at"] is not None + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_failed_status(self, test_client, session): + """Test getting status of a failed connection test.""" + from airflow.models.connection_test import ConnectionTestRequest + + request = ConnectionTestRequest.create_request( + encrypted_connection_uri="gAAAAABfakedata...", + conn_type="postgres", + session=session, + ) + request.mark_running("worker-1") + request.mark_failed("Connection refused: timeout") + session.commit() + + response = test_client.get(f"/connections/test/{request.id}") + assert response.status_code == 200 + + data = response.json() + assert data["state"] == "failed" + assert data["result_status"] is False + assert data["result_message"] == "Connection refused: timeout" + + def test_get_status_not_found(self, test_client): + """Test getting status of a non-existent request.""" + response = test_client.get("/connections/test/non-existent-id") + assert response.status_code == 404 + assert "was not found" in response.json()["detail"] + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/connections/test/some-id") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/connections/test/some-id") + assert response.status_code == 403 + + class TestCreateDefaultConnections(TestConnectionEndpoint): def test_should_respond_204(self, test_client, session): response = test_client.post("/connections/defaults") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py new file mode 100644 index 0000000000000..20b2913a5d240 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -0,0 +1,301 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest + +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState + +from tests_common.test_utils.db import clear_db_connection_tests + +pytestmark = pytest.mark.db_test + +TEST_ENCRYPTED_URI = "gAAAAABfakeencrypteddata..." +TEST_CONN_TYPE = "postgres" + + +@pytest.fixture(autouse=True) +def clean_connection_tests(): + """Clean up connection test requests after tests.""" + yield + clear_db_connection_tests() + + +class TestGetPendingConnectionTests: + """Tests for GET /execution/connection-tests/pending.""" + + def test_get_pending_no_requests(self, client, session): + """Test getting pending requests when none exist.""" + response = client.get( + "/execution/connection-tests/pending", + params={"hostname": "worker-1.example.com"}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["requests"] == [] + + def test_get_pending_returns_pending_requests(self, client, session): + """Test that pending requests are returned.""" + request1 = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + request2 = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="mysql", + session=session, + ) + session.commit() + + response = client.get( + "/execution/connection-tests/pending", + params={"hostname": "worker-1.example.com"}, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["requests"]) == 2 + + # Requests should contain the expected fields + request_ids = {r["request_id"] for r in data["requests"]} + assert request1.id in request_ids + assert request2.id in request_ids + + for req in data["requests"]: + assert "encrypted_connection_uri" in req + assert "conn_type" in req + assert "timeout" in req + + def test_get_pending_marks_requests_as_running(self, client, session): + """Test that fetched requests are marked as running.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + response = client.get( + "/execution/connection-tests/pending", + params={"hostname": "worker-1.example.com"}, + ) + + assert response.status_code == 200 + + # Refresh from database + session.expire_all() + updated_request = session.get(ConnectionTestRequest, request.id) + assert updated_request.state == ConnectionTestState.RUNNING.value + assert updated_request.worker_hostname == "worker-1.example.com" + assert updated_request.started_at is not None + + def test_get_pending_excludes_running_requests(self, client, session): + """Test that running requests are not returned.""" + pending = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="pending_type", + session=session, + ) + running = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="running_type", + session=session, + ) + running.mark_running("other-worker") + session.commit() + + response = client.get( + "/execution/connection-tests/pending", + params={"hostname": "worker-1.example.com"}, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["requests"]) == 1 + assert data["requests"][0]["request_id"] == pending.id + + def test_get_pending_respects_limit(self, client, session): + """Test that limit parameter is respected.""" + for i in range(5): + ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=f"type_{i}", + session=session, + ) + session.commit() + + response = client.get( + "/execution/connection-tests/pending", + params={"hostname": "worker-1.example.com", "limit": 2}, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["requests"]) == 2 + + def test_get_pending_requires_hostname(self, client): + """Test that hostname parameter is required.""" + response = client.get("/execution/connection-tests/pending") + + assert response.status_code == 422 + + +class TestUpdateConnectionTestState: + """Tests for PATCH /execution/connection-tests/{request_id}/state.""" + + def test_update_to_success(self, client, session): + """Test marking a connection test as successful.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + request.mark_running("worker-1") + session.commit() + + response = client.patch( + f"/execution/connection-tests/{request.id}/state", + json={ + "state": "success", + "result_status": True, + "result_message": "Connection successfully tested", + }, + ) + + assert response.status_code == 204 + + # Verify database state + session.expire_all() + updated = session.get(ConnectionTestRequest, request.id) + assert updated.state == ConnectionTestState.SUCCESS.value + assert updated.result_status is True + assert updated.result_message == "Connection successfully tested" + assert updated.completed_at is not None + + def test_update_to_failed(self, client, session): + """Test marking a connection test as failed.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + request.mark_running("worker-1") + session.commit() + + response = client.patch( + f"/execution/connection-tests/{request.id}/state", + json={ + "state": "failed", + "result_status": False, + "result_message": "Connection refused: timeout", + }, + ) + + assert response.status_code == 204 + + # Verify database state + session.expire_all() + updated = session.get(ConnectionTestRequest, request.id) + assert updated.state == ConnectionTestState.FAILED.value + assert updated.result_status is False + assert updated.result_message == "Connection refused: timeout" + + def test_update_to_running(self, client, session): + """Test updating state to running.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{request.id}/state", + json={ + "state": "running", + "hostname": "worker-2.example.com", + }, + ) + + assert response.status_code == 204 + + # Verify database state + session.expire_all() + updated = session.get(ConnectionTestRequest, request.id) + assert updated.state == ConnectionTestState.RUNNING.value + assert updated.worker_hostname == "worker-2.example.com" + + def test_update_not_found(self, client): + """Test updating a non-existent request.""" + response = client.patch( + "/execution/connection-tests/non-existent-id/state", + json={ + "state": "success", + "result_status": True, + "result_message": "OK", + }, + ) + + assert response.status_code == 404 + assert "was not found" in response.json()["detail"] + + def test_update_invalid_state_transition_result_from_pending(self, client, session): + """Test that result cannot be reported from pending state.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + # Try to report result from pending state (should fail) + response = client.patch( + f"/execution/connection-tests/{request.id}/state", + json={ + "state": "success", + "result_status": True, + "result_message": "OK", + }, + ) + + assert response.status_code == 409 + assert "Expected `running`" in response.json()["detail"] + + def test_update_invalid_state_transition_from_completed(self, client, session): + """Test that state cannot be changed after completion.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + request.mark_running("worker-1") + request.mark_success("OK") + session.commit() + + # Try to update state from completed (should fail) + response = client.patch( + f"/execution/connection-tests/{request.id}/state", + json={ + "state": "running", + "hostname": "worker-2", + }, + ) + + assert response.status_code == 409 diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py new file mode 100644 index 0000000000000..8d52b738b28b5 --- /dev/null +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -0,0 +1,392 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow.models.connection_test import ( + TERMINAL_STATES, + VALID_STATE_TRANSITIONS, + ConnectionTestRequest, + ConnectionTestState, +) +from airflow.utils.session import create_session + +from tests_common.test_utils.db import clear_db_connection_tests + +pytestmark = [pytest.mark.db_test] + +TEST_ENCRYPTED_URI = "gAAAAABfakeencrypteddata..." +TEST_CONN_TYPE = "postgres" + + +@pytest.fixture +def session(): + """Fixture that provides a SQLAlchemy session""" + with create_session() as session: + yield session + + +@pytest.fixture(autouse=True) +def clean_db(): + """Clean up connection test requests before and after each test.""" + clear_db_connection_tests() + yield + clear_db_connection_tests() + + +class TestConnectionTestState: + def test_state_values(self): + """Test that ConnectionTestState has expected values.""" + assert ConnectionTestState.PENDING.value == "pending" + assert ConnectionTestState.RUNNING.value == "running" + assert ConnectionTestState.SUCCESS.value == "success" + assert ConnectionTestState.FAILED.value == "failed" + + def test_terminal_states(self): + """Test that terminal states are correctly defined.""" + assert ConnectionTestState.SUCCESS in TERMINAL_STATES + assert ConnectionTestState.FAILED in TERMINAL_STATES + assert ConnectionTestState.PENDING not in TERMINAL_STATES + assert ConnectionTestState.RUNNING not in TERMINAL_STATES + + def test_state_str(self): + """Test that state __str__ returns the value.""" + assert str(ConnectionTestState.PENDING) == "pending" + assert str(ConnectionTestState.SUCCESS) == "success" + + +class TestConnectionTestRequest: + def test_create_request(self, session): + """Test creating a connection test request.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + timeout=120, + session=session, + ) + session.commit() + + assert request.id is not None + assert request.state == ConnectionTestState.PENDING + assert request.encrypted_connection_uri == TEST_ENCRYPTED_URI + assert request.conn_type == TEST_CONN_TYPE + assert request.timeout == 120 + assert request.result_status is None + assert request.result_message is None + assert request.created_at is not None + assert request.started_at is None + assert request.completed_at is None + assert request.worker_hostname is None + + def test_create_request_default_timeout(self, session): + """Test that default timeout is 60 seconds.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + assert request.timeout == 60 + + def test_create_request_without_session(self): + """Test creating a request without adding to session.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + ) + + assert request.id is not None + assert request.state == ConnectionTestState.PENDING + + def test_mark_running(self, session): + """Test marking a request as running.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + request.mark_running("worker-1.example.com") + + assert request.state == ConnectionTestState.RUNNING + assert request.worker_hostname == "worker-1.example.com" + assert request.started_at is not None + + def test_mark_success(self, session): + """Test marking a request as successful.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + request.mark_running("worker-1") + + request.mark_success("Connection successfully tested") + + assert request.state == ConnectionTestState.SUCCESS + assert request.result_status is True + assert request.result_message == "Connection successfully tested" + assert request.completed_at is not None + + def test_mark_failed(self, session): + """Test marking a request as failed.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + request.mark_running("worker-1") + + request.mark_failed("Connection refused: timeout") + + assert request.state == ConnectionTestState.FAILED + assert request.result_status is False + assert request.result_message == "Connection refused: timeout" + assert request.completed_at is not None + + def test_is_terminal(self, session): + """Test the is_terminal property.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + # Initially pending - not terminal + assert request.is_terminal is False + + # Running - not terminal + request.mark_running("worker-1") + assert request.is_terminal is False + + # Success - terminal + request.mark_success("OK") + assert request.is_terminal is True + + def test_is_terminal_failed(self, session): + """Test that failed state is terminal.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + request.mark_running("worker-1") + request.mark_failed("Error") + + assert request.is_terminal is True + + def test_get_pending_requests(self, session): + """Test getting pending requests.""" + # Create multiple requests in different states + pending1 = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + pending2 = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="mysql", + session=session, + ) + running = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="sqlite", + session=session, + ) + session.commit() + + running.mark_running("worker-1") + session.commit() + + # Get pending requests + pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=10) + + # Should only return pending requests + assert len(pending_requests) == 2 + pending_ids = {r.id for r in pending_requests} + assert pending1.id in pending_ids + assert pending2.id in pending_ids + assert running.id not in pending_ids + + def test_get_pending_requests_limit(self, session): + """Test that get_pending_requests respects the limit.""" + # Create 5 pending requests + for i in range(5): + ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=f"type_{i}", + session=session, + ) + session.commit() + + # Get only 2 + pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=2) + assert len(pending_requests) == 2 + + def test_get_pending_requests_ordered_by_created_at(self, session): + """Test that pending requests are ordered by created_at.""" + import time + + first = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="first", + session=session, + ) + session.commit() + + time.sleep(0.01) # Small delay to ensure different timestamps + + second = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type="second", + session=session, + ) + session.commit() + + pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=10) + + assert len(pending_requests) == 2 + assert pending_requests[0].id == first.id + assert pending_requests[1].id == second.id + + def test_repr(self, session): + """Test the __repr__ method.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + repr_str = repr(request) + assert "ConnectionTestRequest" in repr_str + assert str(request.id) in repr_str + assert "pending" in repr_str + assert TEST_CONN_TYPE in repr_str + + def test_persistence(self, session): + """Test that request is correctly persisted and retrieved.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + timeout=90, + session=session, + ) + session.commit() + + request_id = request.id + + # Clear session and retrieve + session.expire_all() + retrieved = session.get(ConnectionTestRequest, request_id) + + assert retrieved is not None + assert retrieved.id == request_id + assert retrieved.encrypted_connection_uri == TEST_ENCRYPTED_URI + assert retrieved.conn_type == TEST_CONN_TYPE + assert retrieved.timeout == 90 + assert retrieved.state == "pending" + + def test_valid_state_transitions_defined(self): + """Test that valid state transitions are correctly defined.""" + # PENDING can go to RUNNING or FAILED + assert ConnectionTestState.RUNNING in VALID_STATE_TRANSITIONS[ConnectionTestState.PENDING] + assert ConnectionTestState.FAILED in VALID_STATE_TRANSITIONS[ConnectionTestState.PENDING] + + # RUNNING can go to SUCCESS or FAILED + assert ConnectionTestState.SUCCESS in VALID_STATE_TRANSITIONS[ConnectionTestState.RUNNING] + assert ConnectionTestState.FAILED in VALID_STATE_TRANSITIONS[ConnectionTestState.RUNNING] + + # Terminal states have no valid transitions + assert len(VALID_STATE_TRANSITIONS[ConnectionTestState.SUCCESS]) == 0 + assert len(VALID_STATE_TRANSITIONS[ConnectionTestState.FAILED]) == 0 + + def test_invalid_transition_pending_to_success(self, session): + """Test that PENDING -> SUCCESS is rejected.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + with pytest.raises(ValueError, match="Invalid state transition"): + request.mark_success("Should fail") + + def test_invalid_transition_success_to_running(self, session): + """Test that SUCCESS -> RUNNING is rejected.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + request.mark_running("worker-1") + request.mark_success("OK") + + with pytest.raises(ValueError, match="Invalid state transition"): + request.mark_running("worker-2") + + def test_invalid_transition_failed_to_success(self, session): + """Test that FAILED -> SUCCESS is rejected.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + request.mark_running("worker-1") + request.mark_failed("Error") + + with pytest.raises(ValueError, match="Invalid state transition"): + request.mark_success("Should fail") + + def test_valid_transition_pending_to_failed(self, session): + """Test that PENDING -> FAILED is allowed (e.g., no executor available).""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + + # Should not raise + request.mark_failed("No executor available") + + assert request.state == ConnectionTestState.FAILED + assert request.result_status is False + + def test_invalid_transition_running_to_running(self, session): + """Test that RUNNING -> RUNNING is rejected.""" + request = ConnectionTestRequest.create_request( + encrypted_connection_uri=TEST_ENCRYPTED_URI, + conn_type=TEST_CONN_TYPE, + session=session, + ) + session.commit() + request.mark_running("worker-1") + + with pytest.raises(ValueError, match="Invalid state transition"): + request.mark_running("worker-2") diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 5e42be6c1bd87..f0a6e188be5c2 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -244,15 +244,30 @@ class ConnectionResponse(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None -class ConnectionTestResponse(BaseModel): +class ConnectionTestQueuedResponse(BaseModel): """ - Connection Test serializer for responses. + Response when a connection test is queued for async execution. """ - status: Annotated[bool, Field(title="Status")] + request_id: Annotated[str, Field(title="Request Id")] + state: Annotated[str, Field(title="State")] message: Annotated[str, Field(title="Message")] +class ConnectionTestStatusResponse(BaseModel): + """ + Response with the full status of a connection test request. + """ + + request_id: Annotated[str, Field(title="Request Id")] + state: Annotated[str, Field(title="State")] + result_status: Annotated[bool | None, Field(title="Result Status")] = None + result_message: Annotated[str | None, Field(title="Result Message")] = None + created_at: Annotated[datetime, Field(title="Created At")] + started_at: Annotated[datetime | None, Field(title="Started At")] = None + completed_at: Annotated[datetime | None, Field(title="Completed At")] = None + + class CreateAssetEventsBody(BaseModel): """ Create asset events request. diff --git a/devel-common/src/tests_common/test_utils/db.py b/devel-common/src/tests_common/test_utils/db.py index f154c601959d9..319d035715e3d 100644 --- a/devel-common/src/tests_common/test_utils/db.py +++ b/devel-common/src/tests_common/test_utils/db.py @@ -339,6 +339,15 @@ def clear_db_callbacks(): session.execute(delete(DbCallbackRequest)) +def clear_db_connection_tests(): + """Clear all connection test requests from the database.""" + with create_session() as session: + if AIRFLOW_V_3_2_PLUS: + from airflow.models.connection_test import ConnectionTestRequest + + session.execute(delete(ConnectionTestRequest)) + + def set_default_pool_slots(slots): with create_session() as session: default_pool = Pool.get_default_pool(session) diff --git a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py index c4dd0572f6ef8..31dcda470239b 100644 --- a/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py +++ b/providers/dbt/cloud/src/airflow/providers/dbt/cloud/hooks/dbt.py @@ -926,10 +926,11 @@ def retry_failed_job_run(self, job_id: int, account_id: int | None = None) -> Re """ return self._run_and_get_response(method="POST", endpoint=f"{account_id}/jobs/{job_id}/rerun/") - def test_connection(self) -> tuple[bool, str]: + @fallback_to_default_account + def test_connection(self, account_id: int | None = None) -> tuple[bool, str]: """Test dbt Cloud connection.""" try: - self._run_and_get_response() + self._run_and_get_response(endpoint=f"{account_id}/") return True, "Successfully connected to dbt Cloud." except Exception as e: return False, str(e) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 6a3a07f5e8a92..b40627561091e 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -78,6 +78,47 @@ class ConnectionResponse(BaseModel): extra: Annotated[str | None, Field(title="Extra")] = None +class State(str, Enum): + SUCCESS = "success" + FAILED = "failed" + + +class ConnectionTestResultPayload(BaseModel): + """ + Payload for reporting connection test result. + """ + + model_config = ConfigDict( + extra="forbid", + ) + state: Annotated[State, Field(title="State")] + result_status: Annotated[bool, Field(title="Result Status")] + result_message: Annotated[str, Field(title="Result Message")] + + +class ConnectionTestRunningPayload(BaseModel): + """ + Payload for marking a connection test as running. + """ + + model_config = ConfigDict( + extra="forbid", + ) + state: Annotated[Literal["running"], Field(title="State")] + hostname: Annotated[str, Field(title="Hostname")] + + +class ConnectionTestWorkload(BaseModel): + """ + Workload data sent to worker for connection test execution. + """ + + request_id: Annotated[str, Field(title="Request Id")] + encrypted_connection_uri: Annotated[str, Field(title="Encrypted Connection Uri")] + conn_type: Annotated[str, Field(title="Conn Type")] + timeout: Annotated[int, Field(title="Timeout")] + + class DagRunAssetReference(BaseModel): """ DagRun serializer for asset responses. @@ -509,6 +550,14 @@ class AssetResponse(BaseModel): extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None +class ConnectionTestPendingResponse(BaseModel): + """ + Response containing pending connection test requests for workers. + """ + + requests: Annotated[list[ConnectionTestWorkload], Field(title="Requests")] + + class HITLDetailRequest(BaseModel): """ Schema for the request part of a Human-in-the-loop detail for a specific task instance. diff --git a/task-sdk/src/airflow/sdk/execution_time/execute_workload.py b/task-sdk/src/airflow/sdk/execution_time/execute_workload.py index 410c676eeb913..cb2008268c8ce 100644 --- a/task-sdk/src/airflow/sdk/execution_time/execute_workload.py +++ b/task-sdk/src/airflow/sdk/execution_time/execute_workload.py @@ -15,15 +15,7 @@ # specific language governing permissions and limitations # under the License. -""" -Module for executing an Airflow task using the workload json provided by a input file. - -Usage: - python execute_workload.py - -Arguments: - input_file (str): Path to the JSON file containing the workload definition. -""" +"""Execute an Airflow workload from a JSON input file.""" from __future__ import annotations @@ -34,15 +26,14 @@ import structlog if TYPE_CHECKING: - from airflow.executors.workloads import ExecuteTask + from airflow.executors.workloads import ExecuteTask, TestConnection log = structlog.get_logger(logger_name=__name__) -def execute_workload(workload: ExecuteTask) -> None: +def execute_workload(workload: ExecuteTask | TestConnection) -> None: from airflow.executors import workloads from airflow.sdk.configuration import conf - from airflow.sdk.execution_time.supervisor import supervise from airflow.sdk.log import configure_logging from airflow.settings import dispose_orm @@ -50,6 +41,10 @@ def execute_workload(workload: ExecuteTask) -> None: configure_logging(output=sys.stdout.buffer, json_output=True) + if isinstance(workload, workloads.TestConnection): + _execute_connection_test(workload, conf) + return + if not isinstance(workload, workloads.ExecuteTask): raise ValueError(f"Executor does not know how to handle {type(workload)}") @@ -63,6 +58,8 @@ def execute_workload(workload: ExecuteTask) -> None: server = conf.get("core", "execution_api_server_url", fallback=default_execution_api_server) log.info("Connecting to server:", server=server) + from airflow.sdk.execution_time.supervisor import supervise + supervise( # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this. ti=workload.ti, # type: ignore[arg-type] @@ -78,6 +75,36 @@ def execute_workload(workload: ExecuteTask) -> None: ) +def _execute_connection_test(workload: TestConnection, conf) -> None: + """Execute a connection test workload.""" + from airflow.executors.connection_test_runner import ( + execute_connection_test_workload, + report_connection_test_result, + ) + + log.info("Executing connection test", request_id=workload.request_id, conn_type=workload.conn_type) + + base_url = conf.get("api", "base_url", fallback="/") + if base_url.startswith("/"): + base_url = f"http://localhost:8080{base_url}" + server = conf.get("core", "execution_api_server_url", fallback=f"{base_url.rstrip('/')}/execution/") + + try: + success, message = execute_connection_test_workload(workload) + except Exception as e: + log.exception("Connection test execution failed") + success = False + message = f"Execution error: {e}" + + report_connection_test_result( + request_id=workload.request_id, + success=success, + message=message, + server_url=server, + token=workload.token, + ) + + def main(): parser = argparse.ArgumentParser( description="Execute a workload in a Containerised executor using the task SDK." From 66b994a0b40663e8cba74e49fd119153bbb37308 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 15 Jan 2026 21:12:25 -0600 Subject: [PATCH 2/6] clean ups --- .../core_api/routes/public/connections.py | 41 +---- .../datamodels/connection_test.py | 26 +-- .../execution_api/routes/connection_tests.py | 82 ++------- .../executors/connection_test_runner.py | 19 +-- .../src/airflow/jobs/scheduler_job_runner.py | 5 +- .../versions/head/test_connection_tests.py | 157 +----------------- .../airflow/sdk/api/datamodels/_generated.py | 31 ---- 7 files changed, 39 insertions(+), 322 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index c61a43f41eb04..a2bf426e84bdc 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -57,11 +57,11 @@ ) from airflow.api_fastapi.logging.decorators import action_logging from airflow.configuration import conf +from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection from airflow.models.connection_test import ConnectionTestRequest from airflow.models.crypto import get_fernet from airflow.utils.db import create_default_connections as db_create_default_connections -from airflow.utils.strings import get_random_string connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") @@ -92,40 +92,11 @@ def test_connection( "Contact your deployment admin to enable it.", ) - # Create a transient connection to get its URI - transient_conn_id = get_random_string() - - # Check if we're testing an existing connection (connection_id provided) - # In this case, we need to merge masked fields (password, extra) with stored values - if test_body.connection_id: - existing_conn = session.scalar(select(Connection).filter_by(conn_id=test_body.connection_id)) - if existing_conn: - # Create a copy of the existing connection for testing - # This merges request data with stored credentials (handles masked passwords) - conn = Connection( - conn_id=transient_conn_id, - conn_type=existing_conn.conn_type, - description=existing_conn.description, - host=existing_conn.host, - schema=existing_conn.schema, - login=existing_conn.login, - port=existing_conn.port, - ) - # Copy password and extra (these are encrypted in db) - conn.set_password(existing_conn.password) - conn.set_extra(existing_conn.extra) - # Now apply updates from request body, merging masked fields - update_orm_from_pydantic(conn, test_body) - else: - # Connection ID provided but not found - use request body as-is - data = test_body.model_dump(by_alias=True) - data["conn_id"] = transient_conn_id - conn = Connection(**data) - else: - # New connection test - use request body directly - data = test_body.model_dump(by_alias=True) - data["conn_id"] = transient_conn_id - conn = Connection(**data) + try: + conn = Connection.get_connection_from_secrets(test_body.connection_id) + update_orm_from_pydantic(conn, test_body) + except AirflowNotFoundException: + conn = Connection(**test_body.model_dump(by_alias=True)) # Encrypt the connection URI for secure storage fernet = get_fernet() diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py index 9c925270b323f..e64d5eedad6e5 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py @@ -15,35 +15,13 @@ # specific language governing permissions and limitations # under the License. -"""Datamodels for worker-side connection test execution.""" +"""Datamodels for worker-side connection test result reporting.""" from __future__ import annotations from typing import Literal -from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel - - -class ConnectionTestWorkload(BaseModel): - """Workload data sent to worker for connection test execution.""" - - request_id: str - encrypted_connection_uri: str - conn_type: str - timeout: int - - -class ConnectionTestPendingResponse(BaseModel): - """Response containing pending connection test requests for workers.""" - - requests: list[ConnectionTestWorkload] - - -class ConnectionTestRunningPayload(StrictBaseModel): - """Payload for marking a connection test as running.""" - - state: Literal["running"] - hostname: str +from airflow.api_fastapi.core_api.base import StrictBaseModel class ConnectionTestResultPayload(StrictBaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index af2d75ea20747..c50164e6a769b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -15,23 +15,16 @@ # specific language governing permissions and limitations # under the License. -"""Execution API routes for worker-side connection test execution.""" +"""Execution API routes for worker-side connection test result reporting.""" from __future__ import annotations -from typing import Annotated - import structlog from cadwyn import VersionedAPIRouter -from fastapi import Body, HTTPException, Query, status +from fastapi import HTTPException, status from airflow.api_fastapi.common.db.common import SessionDep -from airflow.api_fastapi.execution_api.datamodels.connection_test import ( - ConnectionTestPendingResponse, - ConnectionTestResultPayload, - ConnectionTestRunningPayload, - ConnectionTestWorkload, -) +from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultPayload from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState router = VersionedAPIRouter() @@ -39,40 +32,6 @@ log = structlog.get_logger(__name__) -@router.get( - "/pending", - status_code=status.HTTP_200_OK, -) -def get_pending_connection_tests( - session: SessionDep, - hostname: Annotated[str, Query(description="Worker hostname requesting work")], - limit: Annotated[int, Query(description="Maximum number of requests to return", ge=1, le=100)] = 10, -) -> ConnectionTestPendingResponse: - """ - Get pending connection test requests for worker execution. - - Workers call this endpoint to fetch pending connection tests to execute. - """ - log.debug("Worker requesting pending connection tests", hostname=hostname, limit=limit) - - pending_requests = ConnectionTestRequest.get_pending_requests(session, limit=limit) - - workloads = [] - for request in pending_requests: - request.mark_running(hostname) - workloads.append( - ConnectionTestWorkload( - request_id=str(request.id), - encrypted_connection_uri=request.encrypted_connection_uri, - conn_type=request.conn_type, - timeout=request.timeout, - ) - ) - - session.commit() - return ConnectionTestPendingResponse(requests=workloads) - - @router.patch( "/{request_id}/state", status_code=status.HTTP_204_NO_CONTENT, @@ -83,7 +42,7 @@ def get_pending_connection_tests( ) def update_connection_test_state( request_id: str, - payload: Annotated[ConnectionTestRunningPayload | ConnectionTestResultPayload, Body()], + payload: ConnectionTestResultPayload, session: SessionDep, ) -> None: """ @@ -91,8 +50,6 @@ def update_connection_test_state( Workers call this endpoint to report the result of a connection test. """ - log.debug("Updating connection test state", request_id=request_id, payload=payload) - test_request = session.get(ConnectionTestRequest, request_id) if test_request is None: @@ -101,26 +58,19 @@ def update_connection_test_state( f"Connection test request with id `{request_id}` was not found.", ) - if isinstance(payload, ConnectionTestRunningPayload): - if test_request.state not in (ConnectionTestState.PENDING.value, ConnectionTestState.RUNNING.value): - raise HTTPException( - status.HTTP_409_CONFLICT, - f"Cannot transition from state `{test_request.state}` to `running`.", - ) - test_request.mark_running(payload.hostname) - elif isinstance(payload, ConnectionTestResultPayload): - if test_request.state != ConnectionTestState.RUNNING.value: - raise HTTPException( - status.HTTP_409_CONFLICT, - f"Cannot report result when in state `{test_request.state}`. Expected `running`.", - ) - if payload.state == "success": - test_request.mark_success(payload.result_message) - else: - test_request.mark_failed(payload.result_message) + if test_request.state != ConnectionTestState.RUNNING.value: + raise HTTPException( + status.HTTP_409_CONFLICT, + f"Cannot report result when in state `{test_request.state}`. Expected `running`.", + ) + + if payload.state == "success": + test_request.mark_success(payload.result_message) + else: + test_request.mark_failed(payload.result_message) log.info( - "Connection test state updated", + "Connection test result reported", request_id=request_id, - new_state=test_request.state, + state=payload.state, ) diff --git a/airflow-core/src/airflow/executors/connection_test_runner.py b/airflow-core/src/airflow/executors/connection_test_runner.py index 0b2c65fa500e9..dc394f429fe40 100644 --- a/airflow-core/src/airflow/executors/connection_test_runner.py +++ b/airflow-core/src/airflow/executors/connection_test_runner.py @@ -123,26 +123,17 @@ def report_connection_test_result( log.info("Connection test result reported: request_id=%s", request_id) return True except httpx.HTTPStatusError as e: - if 400 <= e.response.status_code < 500 and e.response.status_code != 429: - log.error( - "Failed to report result (client error): request_id=%s, status=%s", - request_id, - e.response.status_code, - ) - return False log.warning( - "Failed to report result, retrying: request_id=%s, attempt=%s/%s", + "HTTP error reporting result: request_id=%s, status=%s, attempt=%s/%s", request_id, + e.response.status_code, attempt + 1, MAX_RETRIES, ) + if 400 <= e.response.status_code < 500 and e.response.status_code != 429: + return False except Exception: - log.warning( - "Failed to report result, retrying: request_id=%s, attempt=%s/%s", - request_id, - attempt + 1, - MAX_RETRIES, - ) + log.exception("Failed to report result: request_id=%s", request_id) if attempt < MAX_RETRIES - 1: delay = RETRY_DELAY_BASE * (2**attempt) * (0.5 + random.random()) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f250d3bd8b3d9..1da2628c1b9a4 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1594,7 +1594,10 @@ def _run_scheduler_loop(self) -> None: deadline.handle_miss(session) # Dispatch pending connection test requests to workers - self._dispatch_connection_tests(session) + try: + self._dispatch_connection_tests(session) + except Exception: + self.log.exception("Error dispatching connection tests") # Heartbeat the scheduler periodically perform_heartbeat( diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index 20b2913a5d240..8aaddc7b0a0d1 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -36,127 +36,6 @@ def clean_connection_tests(): clear_db_connection_tests() -class TestGetPendingConnectionTests: - """Tests for GET /execution/connection-tests/pending.""" - - def test_get_pending_no_requests(self, client, session): - """Test getting pending requests when none exist.""" - response = client.get( - "/execution/connection-tests/pending", - params={"hostname": "worker-1.example.com"}, - ) - - assert response.status_code == 200 - data = response.json() - assert data["requests"] == [] - - def test_get_pending_returns_pending_requests(self, client, session): - """Test that pending requests are returned.""" - request1 = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type=TEST_CONN_TYPE, - session=session, - ) - request2 = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type="mysql", - session=session, - ) - session.commit() - - response = client.get( - "/execution/connection-tests/pending", - params={"hostname": "worker-1.example.com"}, - ) - - assert response.status_code == 200 - data = response.json() - assert len(data["requests"]) == 2 - - # Requests should contain the expected fields - request_ids = {r["request_id"] for r in data["requests"]} - assert request1.id in request_ids - assert request2.id in request_ids - - for req in data["requests"]: - assert "encrypted_connection_uri" in req - assert "conn_type" in req - assert "timeout" in req - - def test_get_pending_marks_requests_as_running(self, client, session): - """Test that fetched requests are marked as running.""" - request = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type=TEST_CONN_TYPE, - session=session, - ) - session.commit() - - response = client.get( - "/execution/connection-tests/pending", - params={"hostname": "worker-1.example.com"}, - ) - - assert response.status_code == 200 - - # Refresh from database - session.expire_all() - updated_request = session.get(ConnectionTestRequest, request.id) - assert updated_request.state == ConnectionTestState.RUNNING.value - assert updated_request.worker_hostname == "worker-1.example.com" - assert updated_request.started_at is not None - - def test_get_pending_excludes_running_requests(self, client, session): - """Test that running requests are not returned.""" - pending = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type="pending_type", - session=session, - ) - running = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type="running_type", - session=session, - ) - running.mark_running("other-worker") - session.commit() - - response = client.get( - "/execution/connection-tests/pending", - params={"hostname": "worker-1.example.com"}, - ) - - assert response.status_code == 200 - data = response.json() - assert len(data["requests"]) == 1 - assert data["requests"][0]["request_id"] == pending.id - - def test_get_pending_respects_limit(self, client, session): - """Test that limit parameter is respected.""" - for i in range(5): - ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type=f"type_{i}", - session=session, - ) - session.commit() - - response = client.get( - "/execution/connection-tests/pending", - params={"hostname": "worker-1.example.com", "limit": 2}, - ) - - assert response.status_code == 200 - data = response.json() - assert len(data["requests"]) == 2 - - def test_get_pending_requires_hostname(self, client): - """Test that hostname parameter is required.""" - response = client.get("/execution/connection-tests/pending") - - assert response.status_code == 422 - - class TestUpdateConnectionTestState: """Tests for PATCH /execution/connection-tests/{request_id}/state.""" @@ -217,31 +96,6 @@ def test_update_to_failed(self, client, session): assert updated.result_status is False assert updated.result_message == "Connection refused: timeout" - def test_update_to_running(self, client, session): - """Test updating state to running.""" - request = ConnectionTestRequest.create_request( - encrypted_connection_uri=TEST_ENCRYPTED_URI, - conn_type=TEST_CONN_TYPE, - session=session, - ) - session.commit() - - response = client.patch( - f"/execution/connection-tests/{request.id}/state", - json={ - "state": "running", - "hostname": "worker-2.example.com", - }, - ) - - assert response.status_code == 204 - - # Verify database state - session.expire_all() - updated = session.get(ConnectionTestRequest, request.id) - assert updated.state == ConnectionTestState.RUNNING.value - assert updated.worker_hostname == "worker-2.example.com" - def test_update_not_found(self, client): """Test updating a non-existent request.""" response = client.patch( @@ -256,7 +110,7 @@ def test_update_not_found(self, client): assert response.status_code == 404 assert "was not found" in response.json()["detail"] - def test_update_invalid_state_transition_result_from_pending(self, client, session): + def test_update_invalid_state_transition_from_pending(self, client, session): """Test that result cannot be reported from pending state.""" request = ConnectionTestRequest.create_request( encrypted_connection_uri=TEST_ENCRYPTED_URI, @@ -279,7 +133,7 @@ def test_update_invalid_state_transition_result_from_pending(self, client, sessi assert "Expected `running`" in response.json()["detail"] def test_update_invalid_state_transition_from_completed(self, client, session): - """Test that state cannot be changed after completion.""" + """Test that result cannot be reported after completion.""" request = ConnectionTestRequest.create_request( encrypted_connection_uri=TEST_ENCRYPTED_URI, conn_type=TEST_CONN_TYPE, @@ -289,12 +143,13 @@ def test_update_invalid_state_transition_from_completed(self, client, session): request.mark_success("OK") session.commit() - # Try to update state from completed (should fail) + # Try to report another result from completed state (should fail) response = client.patch( f"/execution/connection-tests/{request.id}/state", json={ - "state": "running", - "hostname": "worker-2", + "state": "failed", + "result_status": False, + "result_message": "Should not work", }, ) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index b40627561091e..5185341252eee 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -96,29 +96,6 @@ class ConnectionTestResultPayload(BaseModel): result_message: Annotated[str, Field(title="Result Message")] -class ConnectionTestRunningPayload(BaseModel): - """ - Payload for marking a connection test as running. - """ - - model_config = ConfigDict( - extra="forbid", - ) - state: Annotated[Literal["running"], Field(title="State")] - hostname: Annotated[str, Field(title="Hostname")] - - -class ConnectionTestWorkload(BaseModel): - """ - Workload data sent to worker for connection test execution. - """ - - request_id: Annotated[str, Field(title="Request Id")] - encrypted_connection_uri: Annotated[str, Field(title="Encrypted Connection Uri")] - conn_type: Annotated[str, Field(title="Conn Type")] - timeout: Annotated[int, Field(title="Timeout")] - - class DagRunAssetReference(BaseModel): """ DagRun serializer for asset responses. @@ -550,14 +527,6 @@ class AssetResponse(BaseModel): extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None -class ConnectionTestPendingResponse(BaseModel): - """ - Response containing pending connection test requests for workers. - """ - - requests: Annotated[list[ConnectionTestWorkload], Field(title="Requests")] - - class HITLDetailRequest(BaseModel): """ Schema for the request part of a Human-in-the-loop detail for a specific task instance. From 7aa11042d5d310afa8ffdfb9c94659973a49095b Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 15 Jan 2026 21:42:57 -0600 Subject: [PATCH 3/6] add test for dbt --- .../tests/unit/dbt/cloud/hooks/test_dbt.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py b/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py index 45f940236b737..27bdcc02c520e 100644 --- a/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py +++ b/providers/dbt/cloud/tests/unit/dbt/cloud/hooks/test_dbt.py @@ -1055,25 +1055,33 @@ def test_get_job_run_artifact_with_payload(self, mock_http_run, mock_paginate, c hook._paginate.assert_not_called() @pytest.mark.parametrize( - argnames="conn_id", - argvalues=[ACCOUNT_ID_CONN, NO_ACCOUNT_ID_CONN], + argnames=("conn_id", "account_id"), + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], ids=["default_account", "explicit_account"], ) - def test_connection_success(self, requests_mock, conn_id): - requests_mock.get(BASE_URL, status_code=200) - status, msg = DbtCloudHook(conn_id).test_connection() + @patch.object( + DbtCloudHook, + "run_with_advanced_retry", + return_value={"data": "success"}, + ) + def test_connection_success(self, _, conn_id, account_id): + status, msg = DbtCloudHook(conn_id).test_connection(account_id=account_id) assert status is True assert msg == "Successfully connected to dbt Cloud." @pytest.mark.parametrize( - argnames="conn_id", - argvalues=[ACCOUNT_ID_CONN, NO_ACCOUNT_ID_CONN], + argnames=("conn_id", "account_id"), + argvalues=[(ACCOUNT_ID_CONN, None), (NO_ACCOUNT_ID_CONN, ACCOUNT_ID)], ids=["default_account", "explicit_account"], ) - def test_connection_failure(self, requests_mock, conn_id): - requests_mock.get(BASE_URL, status_code=403, reason="Authentication credentials were not provided") - status, msg = DbtCloudHook(conn_id).test_connection() + @patch.object( + DbtCloudHook, + "run_with_advanced_retry", + side_effect=Exception("403:Authentication credentials were not provided"), + ) + def test_connection_failure(self, _, conn_id, account_id): + status, msg = DbtCloudHook(conn_id).test_connection(account_id=account_id) assert status is False assert msg == "403:Authentication credentials were not provided" From 004fda6d74802f45e43ad89344d2d1907511df0f Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 15 Jan 2026 21:46:54 -0600 Subject: [PATCH 4/6] revert unwanted changes --- .../airflow/ui/src/pages/Connections/TestConnectionButton.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx b/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx index 68e9b888fade7..7b27877d87037 100644 --- a/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Connections/TestConnectionButton.tsx @@ -75,7 +75,8 @@ const TestConnectionButton = ({ connection }: Props) => { actionName={ option === "Enabled" ? translate("connections.test") : translate("connections.testDisabled") } - display="flex" + disabled={option === "Disabled"} + display={option === "Hidden" ? "none" : "flex"} icon={icon} loading={isPending} onClick={() => { From 1bebe433360aa3e19316c2e468b59d7d6e2bd175 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 15 Jan 2026 22:35:37 -0600 Subject: [PATCH 5/6] fix failig tests --- airflow-ctl/src/airflowctl/api/operations.py | 8 ++++---- airflow-ctl/tests/airflow_ctl/api/test_operations.py | 9 +++++---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index 31e0298a19d82..af1464199cae5 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -42,7 +42,7 @@ ConnectionBody, ConnectionCollectionResponse, ConnectionResponse, - ConnectionTestResponse, + ConnectionTestQueuedResponse, CreateAssetEventsBody, DAGCollectionResponse, DAGDetailsResponse, @@ -439,11 +439,11 @@ def update( def test( self, connection: ConnectionBody, - ) -> ConnectionTestResponse | ServerResponseError: - """Test a connection.""" + ) -> ConnectionTestQueuedResponse | ServerResponseError: + """Queue a connection test for async execution on a worker.""" try: self.response = self.client.post("connections/test", json=connection.model_dump(mode="json")) - return ConnectionTestResponse.model_validate_json(self.response.content) + return ConnectionTestQueuedResponse.model_validate_json(self.response.content) except ServerResponseError as e: raise e diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index f0a638475c480..697cd1f15d845 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -54,7 +54,7 @@ ConnectionBody, ConnectionCollectionResponse, ConnectionResponse, - ConnectionTestResponse, + ConnectionTestQueuedResponse, CreateAssetEventsBody, DAGCollectionResponse, DAGDetailsResponse, @@ -626,9 +626,10 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert response == self.connection_response def test_test(self): - connection_test_response = ConnectionTestResponse( - status=True, - message="message", + connection_test_response = ConnectionTestQueuedResponse( + request_id="test-request-id", + state="pending", + message="Connection test request queued for execution on a worker.", ) def handle_request(request: httpx.Request) -> httpx.Response: From 3f47b478ae771bc18a8a348a969882ef707e740a Mon Sep 17 00:00:00 2001 From: Anish Date: Fri, 16 Jan 2026 00:44:18 -0600 Subject: [PATCH 6/6] fix migration script --- airflow-core/docs/img/airflow_erd.sha256 | 2 +- airflow-core/docs/img/airflow_erd.svg | 4494 ++++++++--------- ...3_2_0_add_connection_test_request_table.py | 8 +- 3 files changed, 2253 insertions(+), 2251 deletions(-) diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 index 55e734468a288..d4be25bf916a4 100644 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -1 +1 @@ -a2c06c083c41afde4bce7352ed0e510df75a3b2584ca49c4e1f5e602448f0e03 \ No newline at end of file +66d288a8d55d9d4b06e3c4ff10d5ebff298c8eb22756c7da1297fe5e55b805fe \ No newline at end of file diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg index 3068f7d3f1600..c4a3dd2cc764a 100644 --- a/airflow-core/docs/img/airflow_erd.svg +++ b/airflow-core/docs/img/airflow_erd.svg @@ -4,2718 +4,2718 @@ - - + + %3 - + job - -job - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -end_date - - [TIMESTAMP] - -executor_class - - [VARCHAR(500)] - -hostname - - [VARCHAR(500)] - -job_type - - [VARCHAR(30)] - -latest_heartbeat - - [TIMESTAMP] - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -unixname - - [VARCHAR(1000)] + +job + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +end_date + + [TIMESTAMP] + +executor_class + + [VARCHAR(500)] + +hostname + + [VARCHAR(500)] + +job_type + + [VARCHAR(30)] + +latest_heartbeat + + [TIMESTAMP] + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +unixname + + [VARCHAR(1000)] connection_test_request - -connection_test_request - -id - - [VARCHAR(36)] - NOT NULL - -completed_at - - [TIMESTAMP] - -conn_type - - [VARCHAR(500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -encrypted_connection_uri - - [TEXT] - NOT NULL - -result_message - - [TEXT] - -result_status - - [BOOLEAN] - -started_at - - [TIMESTAMP] - -state - - [VARCHAR(10)] - NOT NULL - -timeout - - [INTEGER] - NOT NULL - -worker_hostname - - [VARCHAR(500)] + +connection_test_request + +id + + [VARCHAR(36)] + NOT NULL + +completed_at + + [TIMESTAMP] + +conn_type + + [VARCHAR(500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +encrypted_connection_uri + + [TEXT] + NOT NULL + +result_message + + [TEXT] + +result_status + + [BOOLEAN] + +started_at + + [TIMESTAMP] + +state + + [VARCHAR(10)] + NOT NULL + +timeout + + [INTEGER] + NOT NULL + +worker_hostname + + [VARCHAR(500)] partitioned_asset_key_log - -partitioned_asset_key_log - -id - - [INTEGER] - NOT NULL - -asset_event_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -asset_partition_dag_run_id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -source_partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -target_partition_key - - [VARCHAR(250)] - NOT NULL + +partitioned_asset_key_log + +id + + [INTEGER] + NOT NULL + +asset_event_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +asset_partition_dag_run_id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +source_partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +target_partition_key + + [VARCHAR(250)] + NOT NULL log - -log - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -dttm - - [TIMESTAMP] - NOT NULL - -event - - [VARCHAR(60)] - NOT NULL - -extra - - [TEXT] - -logical_date - - [TIMESTAMP] - -map_index - - [INTEGER] - -owner - - [VARCHAR(500)] - -owner_display_name - - [VARCHAR(500)] - -run_id - - [VARCHAR(250)] - -task_id - - [VARCHAR(250)] - -try_number - - [INTEGER] + +log + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +dttm + + [TIMESTAMP] + NOT NULL + +event + + [VARCHAR(60)] + NOT NULL + +extra + + [TEXT] + +logical_date + + [TIMESTAMP] + +map_index + + [INTEGER] + +owner + + [VARCHAR(500)] + +owner_display_name + + [VARCHAR(500)] + +run_id + + [VARCHAR(250)] + +task_id + + [VARCHAR(250)] + +try_number + + [INTEGER] dag_priority_parsing_request - -dag_priority_parsing_request - -id - - [VARCHAR(32)] - NOT NULL - -bundle_name - - [VARCHAR(250)] - NOT NULL - -relative_fileloc - - [VARCHAR(2000)] - NOT NULL + +dag_priority_parsing_request + +id + + [VARCHAR(32)] + NOT NULL + +bundle_name + + [VARCHAR(250)] + NOT NULL + +relative_fileloc + + [VARCHAR(2000)] + NOT NULL import_error - -import_error - -id - - [INTEGER] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -filename - - [VARCHAR(1024)] - -stacktrace - - [TEXT] - -timestamp - - [TIMESTAMP] + +import_error + +id + + [INTEGER] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +filename + + [VARCHAR(1024)] + +stacktrace + + [TEXT] + +timestamp + + [TIMESTAMP] dag_bundle - -dag_bundle - -name - - [VARCHAR(250)] - NOT NULL - -active - - [BOOLEAN] - -last_refreshed - - [TIMESTAMP] - -signed_url_template - - [VARCHAR(200)] - -template_params - - [JSON] - -version - - [VARCHAR(200)] + +dag_bundle + +name + + [VARCHAR(250)] + NOT NULL + +active + + [BOOLEAN] + +last_refreshed + + [TIMESTAMP] + +signed_url_template + + [VARCHAR(200)] + +template_params + + [JSON] + +version + + [VARCHAR(200)] dag_bundle_team - -dag_bundle_team - -dag_bundle_name - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - NOT NULL + +dag_bundle_team + +dag_bundle_name + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + NOT NULL dag_bundle:name--dag_bundle_team:dag_bundle_name - -0..N -1 + +0..N +1 dag - -dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -asset_expression - - [JSON] - -bundle_name - - [VARCHAR(250)] - NOT NULL - -bundle_version - - [VARCHAR(200)] - -dag_display_name - - [VARCHAR(2000)] - -deadline - - [JSON] - -description - - [TEXT] - -exceeds_max_non_backfill - - [BOOLEAN] - NOT NULL - -fail_fast - - [BOOLEAN] - NOT NULL - -fileloc - - [VARCHAR(2000)] - -has_import_errors - - [BOOLEAN] - NOT NULL - -has_task_concurrency_limits - - [BOOLEAN] - NOT NULL - -is_paused - - [BOOLEAN] - NOT NULL - -is_stale - - [BOOLEAN] - NOT NULL - -last_expired - - [TIMESTAMP] - -last_parse_duration - - [DOUBLE PRECISION] - -last_parsed_time - - [TIMESTAMP] - -max_active_runs - - [INTEGER] - -max_active_tasks - - [INTEGER] - NOT NULL - -max_consecutive_failed_dag_runs - - [INTEGER] - NOT NULL - -next_dagrun - - [TIMESTAMP] - -next_dagrun_create_after - - [TIMESTAMP] - -next_dagrun_data_interval_end - - [TIMESTAMP] - -next_dagrun_data_interval_start - - [TIMESTAMP] - -owners - - [VARCHAR(2000)] - -relative_fileloc - - [VARCHAR(2000)] - -timetable_description - - [VARCHAR(1000)] - -timetable_summary - - [TEXT] - -timetable_type - - [VARCHAR(255)] - NOT NULL + +dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +asset_expression + + [JSON] + +bundle_name + + [VARCHAR(250)] + NOT NULL + +bundle_version + + [VARCHAR(200)] + +dag_display_name + + [VARCHAR(2000)] + +deadline + + [JSON] + +description + + [TEXT] + +exceeds_max_non_backfill + + [BOOLEAN] + NOT NULL + +fail_fast + + [BOOLEAN] + NOT NULL + +fileloc + + [VARCHAR(2000)] + +has_import_errors + + [BOOLEAN] + NOT NULL + +has_task_concurrency_limits + + [BOOLEAN] + NOT NULL + +is_paused + + [BOOLEAN] + NOT NULL + +is_stale + + [BOOLEAN] + NOT NULL + +last_expired + + [TIMESTAMP] + +last_parse_duration + + [DOUBLE PRECISION] + +last_parsed_time + + [TIMESTAMP] + +max_active_runs + + [INTEGER] + +max_active_tasks + + [INTEGER] + NOT NULL + +max_consecutive_failed_dag_runs + + [INTEGER] + NOT NULL + +next_dagrun + + [TIMESTAMP] + +next_dagrun_create_after + + [TIMESTAMP] + +next_dagrun_data_interval_end + + [TIMESTAMP] + +next_dagrun_data_interval_start + + [TIMESTAMP] + +owners + + [VARCHAR(2000)] + +relative_fileloc + + [VARCHAR(2000)] + +timetable_description + + [VARCHAR(1000)] + +timetable_summary + + [TEXT] + +timetable_type + + [VARCHAR(255)] + NOT NULL dag_bundle:name--dag:bundle_name - -0..N -1 + +0..N +1 team - -team - -name - - [VARCHAR(50)] - NOT NULL + +team + +name + + [VARCHAR(50)] + NOT NULL team:name--dag_bundle_team:team_name - -0..N -1 + +0..N +1 connection - -connection - -id - - [INTEGER] - NOT NULL - -conn_id - - [VARCHAR(250)] - NOT NULL - -conn_type - - [VARCHAR(500)] - NOT NULL - -description - - [TEXT] - -extra - - [TEXT] - -host - - [VARCHAR(500)] - -is_encrypted - - [BOOLEAN] - NOT NULL - -is_extra_encrypted - - [BOOLEAN] - NOT NULL - -login - - [TEXT] - -password - - [TEXT] - -port - - [INTEGER] - -schema - - [VARCHAR(500)] - -team_name - - [VARCHAR(50)] + +connection + +id + + [INTEGER] + NOT NULL + +conn_id + + [VARCHAR(250)] + NOT NULL + +conn_type + + [VARCHAR(500)] + NOT NULL + +description + + [TEXT] + +extra + + [TEXT] + +host + + [VARCHAR(500)] + +is_encrypted + + [BOOLEAN] + NOT NULL + +is_extra_encrypted + + [BOOLEAN] + NOT NULL + +login + + [TEXT] + +password + + [TEXT] + +port + + [INTEGER] + +schema + + [VARCHAR(500)] + +team_name + + [VARCHAR(50)] team:name--connection:team_name - -0..N -{0,1} + +0..N +{0,1} slot_pool - -slot_pool - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -include_deferred - - [BOOLEAN] - NOT NULL - -pool - - [VARCHAR(256)] - NOT NULL - -slots - - [INTEGER] - NOT NULL - -team_name - - [VARCHAR(50)] + +slot_pool + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +include_deferred + + [BOOLEAN] + NOT NULL + +pool + + [VARCHAR(256)] + NOT NULL + +slots + + [INTEGER] + NOT NULL + +team_name + + [VARCHAR(50)] team:name--slot_pool:team_name - -0..N -{0,1} + +0..N +{0,1} variable - -variable - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -is_encrypted - - [BOOLEAN] - NOT NULL - -key - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - -val - - [TEXT] - NOT NULL + +variable + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +is_encrypted + + [BOOLEAN] + NOT NULL + +key + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + +val + + [TEXT] + NOT NULL team:name--variable:team_name - -0..N -{0,1} + +0..N +{0,1} trigger - -trigger - -id - - [INTEGER] - NOT NULL - -classpath - - [VARCHAR(1000)] - NOT NULL - -created_date - - [TIMESTAMP] - NOT NULL - -kwargs - - [TEXT] - NOT NULL - -queue - - [VARCHAR(256)] - -triggerer_id - - [INTEGER] + +trigger + +id + + [INTEGER] + NOT NULL + +classpath + + [VARCHAR(1000)] + NOT NULL + +created_date + + [TIMESTAMP] + NOT NULL + +kwargs + + [TEXT] + NOT NULL + +queue + + [VARCHAR(256)] + +triggerer_id + + [INTEGER] callback - -callback - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -data - - [JSONB] - NOT NULL - -fetch_method - - [VARCHAR(20)] - NOT NULL - -output - - [TEXT] - -priority_weight - - [INTEGER] - NOT NULL - -state - - [VARCHAR(10)] - -trigger_id - - [INTEGER] - -type - - [VARCHAR(20)] - NOT NULL + +callback + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +data + + [JSONB] + NOT NULL + +fetch_method + + [VARCHAR(20)] + NOT NULL + +output + + [TEXT] + +priority_weight + + [INTEGER] + NOT NULL + +state + + [VARCHAR(10)] + +trigger_id + + [INTEGER] + +type + + [VARCHAR(20)] + NOT NULL trigger:id--callback:trigger_id - -0..N -{0,1} + +0..N +{0,1} asset_watcher - -asset_watcher - -asset_id - - [INTEGER] - NOT NULL - -trigger_id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL + +asset_watcher + +asset_id + + [INTEGER] + NOT NULL + +trigger_id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL trigger:id--asset_watcher:trigger_id - -0..N -1 + +0..N +1 task_instance - -task_instance - -id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - NOT NULL - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - NOT NULL - -last_heartbeat_at - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - NOT NULL - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - NOT NULL - -queue - - [VARCHAR(256)] - NOT NULL - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - NOT NULL - -updated_at - - [TIMESTAMP] + +task_instance + +id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + NOT NULL + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + NOT NULL + +last_heartbeat_at + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + NOT NULL + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + NOT NULL + +queue + + [VARCHAR(256)] + NOT NULL + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + NOT NULL + +updated_at + + [TIMESTAMP] - + trigger:id--task_instance:trigger_id - -0..N -{0,1} + +0..N +{0,1} deadline - -deadline - -id - - [UUID] - NOT NULL - -callback_id - - [UUID] - NOT NULL - -dagrun_id - - [INTEGER] - -deadline_time - - [TIMESTAMP] - NOT NULL - -missed - - [BOOLEAN] - NOT NULL + +deadline + +id + + [UUID] + NOT NULL + +callback_id + + [UUID] + NOT NULL + +dagrun_id + + [INTEGER] + +deadline_time + + [TIMESTAMP] + NOT NULL + +missed + + [BOOLEAN] + NOT NULL callback:id--deadline:callback_id - -0..N -1 + +0..N +1 asset_alias - -asset_alias - -id - - [INTEGER] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL + +asset_alias + +id + + [INTEGER] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL asset_alias_asset - -asset_alias_asset - -alias_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL + +asset_alias_asset + +alias_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL asset_alias:id--asset_alias_asset:alias_id - -0..N -1 + +0..N +1 asset_alias_asset_event - -asset_alias_asset_event - -alias_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +asset_alias_asset_event + +alias_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL asset_alias:id--asset_alias_asset_event:alias_id - -0..N -1 + +0..N +1 dag_schedule_asset_alias_reference - -dag_schedule_asset_alias_reference - -alias_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_alias_reference + +alias_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL asset_alias:id--dag_schedule_asset_alias_reference:alias_id - -0..N -1 + +0..N +1 asset - -asset - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -extra - - [JSON] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +extra + + [JSON] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL asset:id--asset_alias_asset:asset_id - -0..N -1 + +0..N +1 asset:id--asset_watcher:asset_id - -0..N -1 + +0..N +1 asset_active - -asset_active - -name - - [VARCHAR(1500)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset_active + +name + + [VARCHAR(1500)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL asset:uri--asset_active:uri - -1 -1 + +1 +1 asset:name--asset_active:name - -1 -1 + +1 +1 dag_schedule_asset_reference - -dag_schedule_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset:id--dag_schedule_asset_reference:asset_id - -0..N -1 + +0..N +1 task_outlet_asset_reference - -task_outlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_outlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset:id--task_outlet_asset_reference:asset_id - -0..N -1 + +0..N +1 task_inlet_asset_reference - -task_inlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_inlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset:id--task_inlet_asset_reference:asset_id - -0..N -1 + +0..N +1 asset_dag_run_queue - -asset_dag_run_queue - -asset_id - - [INTEGER] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +asset_dag_run_queue + +asset_id + + [INTEGER] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL - + asset:id--asset_dag_run_queue:asset_id - -0..N -1 + +0..N +1 asset_event - -asset_event - -id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -extra - - [JSON] - NOT NULL - -partition_key - - [VARCHAR(250)] - -source_dag_id - - [VARCHAR(250)] - -source_map_index - - [INTEGER] - -source_run_id - - [VARCHAR(250)] - -source_task_id - - [VARCHAR(250)] - -timestamp - - [TIMESTAMP] - NOT NULL + +asset_event + +id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +extra + + [JSON] + NOT NULL + +partition_key + + [VARCHAR(250)] + +source_dag_id + + [VARCHAR(250)] + +source_map_index + + [INTEGER] + +source_run_id + + [VARCHAR(250)] + +source_task_id + + [VARCHAR(250)] + +timestamp + + [TIMESTAMP] + NOT NULL asset_event:id--asset_alias_asset_event:event_id - -0..N -1 + +0..N +1 dagrun_asset_event - -dagrun_asset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +dagrun_asset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL asset_event:id--dagrun_asset_event:event_id - -0..N -1 + +0..N +1 dag_schedule_asset_name_reference - -dag_schedule_asset_name_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_name_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_schedule_asset_name_reference:dag_id - -0..N -1 + +0..N +1 dag_schedule_asset_uri_reference - -dag_schedule_asset_uri_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_uri_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_schedule_asset_uri_reference:dag_id - -0..N -1 + +0..N +1 dag:dag_id--dag_schedule_asset_alias_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--dag_schedule_asset_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--task_outlet_asset_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--task_inlet_asset_reference:dag_id - -0..N -1 + +0..N +1 - + dag:dag_id--asset_dag_run_queue:target_dag_id - -0..N -1 + +0..N +1 dag_version - -dag_version - -id - - [UUID] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -bundle_version - - [VARCHAR(250)] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -version_number - - [INTEGER] - NOT NULL + +dag_version + +id + + [UUID] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +bundle_version + + [VARCHAR(250)] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +version_number + + [INTEGER] + NOT NULL dag:dag_id--dag_version:dag_id - -0..N -1 + +0..N +1 dag_tag - -dag_tag - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL + +dag_tag + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL dag:dag_id--dag_tag:dag_id - -0..N -1 + +0..N +1 dag_owner_attributes - -dag_owner_attributes - -dag_id - - [VARCHAR(250)] - NOT NULL - -owner - - [VARCHAR(500)] - NOT NULL - -link - - [VARCHAR(500)] - NOT NULL + +dag_owner_attributes + +dag_id + + [VARCHAR(250)] + NOT NULL + +owner + + [VARCHAR(500)] + NOT NULL + +link + + [VARCHAR(500)] + NOT NULL dag:dag_id--dag_owner_attributes:dag_id - -0..N -1 + +0..N +1 dag_warning - -dag_warning - -dag_id - - [VARCHAR(250)] - NOT NULL - -warning_type - - [VARCHAR(50)] - NOT NULL - -message - - [TEXT] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL + +dag_warning + +dag_id + + [VARCHAR(250)] + NOT NULL + +warning_type + + [VARCHAR(50)] + NOT NULL + +message + + [TEXT] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL dag:dag_id--dag_warning:dag_id - -0..N -1 + +0..N +1 dag_favorite - -dag_favorite - -dag_id - - [VARCHAR(250)] - NOT NULL - -user_id - - [VARCHAR(250)] - NOT NULL + +dag_favorite + +dag_id + + [VARCHAR(250)] + NOT NULL + +user_id + + [VARCHAR(250)] + NOT NULL dag:dag_id--dag_favorite:dag_id - -0..N -1 + +0..N +1 dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - -bundle_version - - [VARCHAR(250)] - -clear_number - - [INTEGER] - NOT NULL - -conf - - [JSONB] - -context_carrier - - [JSONB] - -created_dag_version_id - - [UUID] - -creating_job_id - - [INTEGER] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - NOT NULL - -logical_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -queued_at - - [TIMESTAMP] - -run_after - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -scheduled_by_job_id - - [INTEGER] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - NOT NULL - -triggered_by - - [VARCHAR(50)] - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + +bundle_version + + [VARCHAR(250)] + +clear_number + + [INTEGER] + NOT NULL + +conf + + [JSONB] + +context_carrier + + [JSONB] + +created_dag_version_id + + [UUID] + +creating_job_id + + [INTEGER] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + NOT NULL + +logical_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +queued_at + + [TIMESTAMP] + +run_after + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +scheduled_by_job_id + + [INTEGER] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + NOT NULL + +triggered_by + + [VARCHAR(50)] + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL - + dag_version:id--dag_run:created_dag_version_id - -0..N -{0,1} + +0..N +{0,1} dag_code - -dag_code - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -fileloc - - [VARCHAR(2000)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -source_code - - [TEXT] - NOT NULL - -source_code_hash - - [VARCHAR(32)] - NOT NULL + +dag_code + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +fileloc + + [VARCHAR(2000)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +source_code + + [TEXT] + NOT NULL + +source_code_hash + + [VARCHAR(32)] + NOT NULL dag_version:id--dag_code:dag_version_id - -0..N -1 + +0..N +1 serialized_dag - -serialized_dag - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_hash - - [VARCHAR(32)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -data - - [JSONB] - -data_compressed - - [BYTEA] - -last_updated - - [TIMESTAMP] - NOT NULL + +serialized_dag + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_hash + + [VARCHAR(32)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +data + + [JSONB] + +data_compressed + + [BYTEA] + +last_updated + + [TIMESTAMP] + NOT NULL dag_version:id--serialized_dag:dag_version_id - -0..N -1 + +0..N +1 - + dag_version:id--task_instance:dag_version_id - -0..N -{0,1} + +0..N +{0,1} log_template - -log_template - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -elasticsearch_id - - [TEXT] - NOT NULL - -filename - - [TEXT] - NOT NULL + +log_template + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +elasticsearch_id + + [TEXT] + NOT NULL + +filename + + [TEXT] + NOT NULL - + log_template:id--dag_run:log_template_id - -0..N -1 + +0..N +1 dag_run:id--deadline:dagrun_id - -0..N -{0,1} + +0..N +{0,1} dag_run:id--dagrun_asset_event:dag_run_id - -0..N -1 + +0..N +1 asset_partition_dag_run - -asset_partition_dag_run - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_run_id - - [INTEGER] - -partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +asset_partition_dag_run + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_run_id + + [INTEGER] + +partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL dag_run:id--asset_partition_dag_run:created_dag_run_id - -0..N -{0,1} + +0..N +{0,1} - + dag_run:run_id--task_instance:run_id - -0..N -1 + +0..N +1 - + dag_run:dag_id--task_instance:dag_id - -0..N -1 + +0..N +1 backfill_dag_run - -backfill_dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - NOT NULL - -dag_run_id - - [INTEGER] - -exception_reason - - [VARCHAR(250)] - -logical_date - - [TIMESTAMP] - NOT NULL - -sort_ordinal - - [INTEGER] - NOT NULL + +backfill_dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + NOT NULL + +dag_run_id + + [INTEGER] + +exception_reason + + [VARCHAR(250)] + +logical_date + + [TIMESTAMP] + NOT NULL + +sort_ordinal + + [INTEGER] + NOT NULL dag_run:id--backfill_dag_run:dag_run_id - -0..N -{0,1} + +0..N +{0,1} dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] dag_run:id--dag_run_note:dag_run_id - -1 -1 + +1 +1 backfill - -backfill - -id - - [INTEGER] - NOT NULL - -completed_at - - [TIMESTAMP] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_run_conf - - [JSON] - NOT NULL - -from_date - - [TIMESTAMP] - NOT NULL - -is_paused - - [BOOLEAN] - -max_active_runs - - [INTEGER] - NOT NULL - -reprocess_behavior - - [VARCHAR(250)] - NOT NULL - -to_date - - [TIMESTAMP] - NOT NULL - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL + +backfill + +id + + [INTEGER] + NOT NULL + +completed_at + + [TIMESTAMP] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_run_conf + + [JSON] + NOT NULL + +from_date + + [TIMESTAMP] + NOT NULL + +is_paused + + [BOOLEAN] + +max_active_runs + + [INTEGER] + NOT NULL + +reprocess_behavior + + [VARCHAR(250)] + NOT NULL + +to_date + + [TIMESTAMP] + NOT NULL + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL - + backfill:id--dag_run:backfill_id - -0..N -{0,1} + +0..N +{0,1} backfill:id--backfill_dag_run:backfill_id - -0..N -1 + +0..N +1 hitl_detail - -hitl_detail - -ti_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL + +hitl_detail + +ti_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL task_instance:id--hitl_detail:ti_id - -1 -1 + +1 +1 task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSONB] - -length - - [INTEGER] - NOT NULL + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSONB] + +length + + [INTEGER] + NOT NULL -task_instance:map_index--task_map:map_index - -0..N -1 +task_instance:task_id--task_map:task_id + +0..N +1 -task_instance:task_id--task_map:task_id - -0..N -1 +task_instance:map_index--task_map:map_index + +0..N +1 -task_instance:dag_id--task_map:dag_id - -0..N -1 +task_instance:run_id--task_map:run_id + +0..N +1 -task_instance:run_id--task_map:run_id - -0..N -1 +task_instance:dag_id--task_map:dag_id + +0..N +1 task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -ti_id - - [UUID] - NOT NULL + +task_reschedule + +id + + [INTEGER] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +ti_id + + [UUID] + NOT NULL task_instance:id--task_reschedule:ti_id - -0..N -1 + +0..N +1 xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [JSONB] + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [JSONB] -task_instance:dag_id--xcom:dag_id - -0..N -1 +task_instance:task_id--xcom:task_id + +0..N +1 -task_instance:task_id--xcom:task_id - -0..N -1 +task_instance:dag_id--xcom:dag_id + +0..N +1 -task_instance:run_id--xcom:run_id - -0..N -1 +task_instance:map_index--xcom:map_index + +0..N +1 -task_instance:map_index--xcom:map_index - -0..N -1 +task_instance:run_id--xcom:run_id + +0..N +1 task_instance_note - -task_instance_note - -ti_id - - [UUID] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +task_instance_note + +ti_id + + [UUID] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] task_instance:id--task_instance_note:ti_id - -1 -1 + +1 +1 task_instance_history - -task_instance_history - -task_instance_id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] + +task_instance_history + +task_instance_id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] -task_instance:task_id--task_instance_history:task_id - -0..N -1 +task_instance:dag_id--task_instance_history:dag_id + +0..N +1 task_instance:run_id--task_instance_history:run_id - -0..N -1 + +0..N +1 task_instance:map_index--task_instance_history:map_index - -0..N -1 + +0..N +1 -task_instance:dag_id--task_instance_history:dag_id - -0..N -1 +task_instance:task_id--task_instance_history:task_id + +0..N +1 rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL -task_instance:run_id--rendered_task_instance_fields:run_id - -0..N -1 +task_instance:task_id--rendered_task_instance_fields:task_id + +0..N +1 task_instance:map_index--rendered_task_instance_fields:map_index - -0..N -1 + +0..N +1 task_instance:dag_id--rendered_task_instance_fields:dag_id - -0..N -1 + +0..N +1 -task_instance:task_id--rendered_task_instance_fields:task_id - -0..N -1 +task_instance:run_id--rendered_task_instance_fields:run_id + +0..N +1 hitl_detail_history - -hitl_detail_history - -ti_history_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL + +hitl_detail_history + +ti_history_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL task_instance_history:task_instance_id--hitl_detail_history:ti_history_id - -1 -1 + +1 +1 alembic_version - -alembic_version - -version_num - - [VARCHAR(32)] - NOT NULL + +alembic_version + +version_num + + [VARCHAR(32)] + NOT NULL diff --git a/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py b/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py index 4bf57b14ad15e..bf56f5f03749b 100644 --- a/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py +++ b/airflow-core/src/airflow/migrations/versions/0099_3_2_0_add_connection_test_request_table.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.utils.sqlalchemy import UtcDateTime + # revision identifiers, used by Alembic. revision = "9882c124ea54" down_revision = "e79fc784f145" @@ -48,9 +50,9 @@ def upgrade(): sa.Column("conn_type", sa.String(500), nullable=False), sa.Column("result_status", sa.Boolean(), nullable=True), sa.Column("result_message", sa.Text(), nullable=True), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.Column("started_at", sa.DateTime(), nullable=True), - sa.Column("completed_at", sa.DateTime(), nullable=True), + sa.Column("created_at", UtcDateTime(timezone=True), nullable=False), + sa.Column("started_at", UtcDateTime(timezone=True), nullable=True), + sa.Column("completed_at", UtcDateTime(timezone=True), nullable=True), sa.Column("timeout", sa.Integer(), nullable=False, default=60), sa.Column("worker_hostname", sa.String(500), nullable=True), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_request_pkey")),