style: lint and format

Author: yangkyeongmo

Makefile

@@ -8,9 +8,9 @@ run:
 run-sse:
 	$(PYTHON) src --transport sse $(PARAMS)
 build:
-	$(PYTHON) -m build
+	uv build
 publish:
-	$(PYTHON) -m twine upload --config-file .pypirc dist/*
+	uv publish
 lint:
 	$(PYTHON) -m ruff check . --fix
 format:

pyproject.toml

@@ -57,3 +57,8 @@ include = [
 [tool.ruff]
 select = ["E", "W", "F", "B", "I"]
 line-length = 120
+
+[dependency-groups]
+dev = [
+    "ruff>=0.11.0",
+]

src/__init__.py

@@ -1,2 +0,0 @@
-# To load all tools
-from src.airflow import *

src/airflow/config.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.config_api import ConfigApi
@@ -22,7 +22,7 @@ async def get_config(
     kwargs: Dict[str, Any] = {}
     if section is not None:
         kwargs["section"] = section
-    
+
     response = config_api.get_config(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 

src/airflow/connection.py

@@ -1,7 +1,6 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
-
 from airflow_client.client.api.connection_api import ConnectionApi
 
 from src.airflow.airflow_client import api_client
@@ -33,7 +32,7 @@ async def list_connections(
         kwargs["offset"] = offset
     if order_by is not None:
         kwargs["order_by"] = order_by
-    
+
     response = connection_api.get_connections(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -64,14 +63,12 @@ async def create_connection(
         connection_request["schema"] = schema
     if extra is not None:
         connection_request["extra"] = extra
-    
+
     response = connection_api.post_connection(connection_request=connection_request)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
 
-async def get_connection(
-    conn_id: str
-) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
+async def get_connection(conn_id: str) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     response = connection_api.get_connection(connection_id=conn_id)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -101,18 +98,14 @@ async def update_connection(
         update_request["schema"] = schema
     if extra is not None:
         update_request["extra"] = extra
-    
+
     response = connection_api.patch_connection(
-        connection_id=conn_id,
-        update_mask=list(update_request.keys()),
-        connection_request=update_request
+        connection_id=conn_id, update_mask=list(update_request.keys()), connection_request=update_request
     )
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
 
-async def delete_connection(
-    conn_id: str
-) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
+async def delete_connection(conn_id: str) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     response = connection_api.delete_connection(connection_id=conn_id)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -141,6 +134,6 @@ async def test_connection(
         connection_request["schema"] = schema
     if extra is not None:
         connection_request["extra"] = extra
-    
+
     response = connection_api.test_connection(connection_request=connection_request)
     return [types.TextContent(type="text", text=str(response.to_dict()))]

src/airflow/dag.py

@@ -1,10 +1,11 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.dag_api import DAGApi
 from airflow_client.client.model.clear_task_instances import ClearTaskInstances
-from airflow_client.client.model.update_task_instances_state import UpdateTaskInstancesState
 from airflow_client.client.model.dag import DAG
+from airflow_client.client.model.update_task_instances_state import UpdateTaskInstancesState
+
 from src.airflow.airflow_client import api_client
 from src.envs import AIRFLOW_HOST
 
@@ -86,12 +87,14 @@ async def get_dag(dag_id: str) -> List[Union[types.TextContent, types.ImageConte
     return [types.TextContent(type="text", text=str(response_dict))]
 
 
-async def get_dag_details(dag_id: str, fields: Optional[List[str]] = None) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
+async def get_dag_details(
+    dag_id: str, fields: Optional[List[str]] = None
+) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     # Build parameters dictionary
     kwargs: Dict[str, Any] = {}
     if fields is not None:
         kwargs["fields"] = fields
-    
+
     response = dag_api.get_dag_details(dag_id=dag_id, **kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -136,11 +139,13 @@ async def patch_dag(
 
 
 async def patch_dags(
-    dag_id_pattern: Optional[str] = None, is_paused: Optional[bool] = None, tags: Optional[List[str]] = None, 
+    dag_id_pattern: Optional[str] = None,
+    is_paused: Optional[bool] = None,
+    tags: Optional[List[str]] = None,
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     update_request = {}
     update_mask = []
-    
+
     if is_paused is not None:
         update_request["is_paused"] = is_paused
         update_mask.append("is_paused")
@@ -153,7 +158,7 @@ async def patch_dags(
     kwargs = {}
     if dag_id_pattern is not None:
         kwargs["dag_id_pattern"] = dag_id_pattern
-    
+
     response = dag_api.patch_dags(dag_id_pattern=dag_id_pattern, dag=dag, update_mask=update_mask, **kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -170,7 +175,9 @@ async def get_task(
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
 
-async def get_tasks(dag_id: str, order_by: Optional[str] = None) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
+async def get_tasks(
+    dag_id: str, order_by: Optional[str] = None
+) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     kwargs = {}
     if order_by is not None:
         kwargs["order_by"] = order_by
@@ -218,7 +225,7 @@ async def clear_task_instances(
         clear_request["reset_dag_runs"] = reset_dag_runs
 
     clear_task_instances = ClearTaskInstances(**clear_request)
-    
+
     response = dag_api.post_clear_task_instances(dag_id=dag_id, clear_task_instances=clear_task_instances)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -234,9 +241,7 @@ async def set_task_instances_state(
     include_past: Optional[bool] = None,
     dry_run: Optional[bool] = None,
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
-    state_request = {
-        "state": state
-    }
+    state_request = {"state": state}
     if task_ids is not None:
         state_request["task_ids"] = task_ids
     if execution_date is not None:
@@ -254,12 +259,15 @@ async def set_task_instances_state(
 
     update_task_instances_state = UpdateTaskInstancesState(**state_request)
 
-    response = dag_api.post_set_task_instances_state(dag_id=dag_id, update_task_instances_state=update_task_instances_state)
+    response = dag_api.post_set_task_instances_state(
+        dag_id=dag_id,
+        update_task_instances_state=update_task_instances_state,
+    )
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
 
 async def reparse_dag_file(
-    file_token: str
+    file_token: str,
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     response = dag_api.reparse_dag_file(file_token=file_token)
     return [types.TextContent(type="text", text=str(response.to_dict()))]

src/airflow/dagrun.py

@@ -1,5 +1,5 @@
-from typing import Any, Dict, List, Optional, Union, Callable
 from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.dag_run_api import DAGRunApi
@@ -27,6 +27,7 @@ def get_all_functions() -> list[tuple[Callable, str, str]]:
         (get_upstream_dataset_events, "get_upstream_dataset_events", "Get dataset events for a DAG run"),
     ]
 
+
 def get_dag_run_url(dag_id: str, dag_run_id: str) -> str:
     return f"{AIRFLOW_HOST}/dags/{dag_id}/grid?dag_run_id={dag_run_id}"
 
@@ -158,36 +159,40 @@ async def get_dag_runs_batch(
         request["page_limit"] = page_limit
 
     response = dag_run_api.get_dag_runs_batch(list_dag_runs_form=request)
-    
+
     # Convert response to dictionary for easier manipulation
     response_dict = response.to_dict()
-    
+
     # Add UI links to each DAG run
     for dag_run in response_dict.get("dag_runs", []):
         dag_run["ui_url"] = get_dag_run_url(dag_run["dag_id"], dag_run["dag_run_id"])
-    
+
     return [types.TextContent(type="text", text=str(response_dict))]
 
 
 async def get_dag_run(
     dag_id: str, dag_run_id: str
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     response = dag_run_api.get_dag_run(dag_id=dag_id, dag_run_id=dag_run_id)
-    
+
     # Convert response to dictionary for easier manipulation
     response_dict = response.to_dict()
-    
+
     # Add UI link to DAG run
     response_dict["ui_url"] = get_dag_run_url(dag_id, dag_run_id)
-    
+
     return [types.TextContent(type="text", text=str(response_dict))]
 
 
 async def update_dag_run_state(
     dag_id: str, dag_run_id: str, state: Optional[str] = None
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     update_dag_run_state = UpdateDagRunState(state=state)
-    response = dag_run_api.update_dag_run_state(dag_id=dag_id, dag_run_id=dag_run_id, update_dag_run_state=update_dag_run_state)
+    response = dag_run_api.update_dag_run_state(
+        dag_id=dag_id,
+        dag_run_id=dag_run_id,
+        update_dag_run_state=update_dag_run_state,
+    )
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
 

src/airflow/dagstats.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.dag_stats_api import DagStatsApi
@@ -21,6 +21,6 @@ async def get_dag_stats(
     kwargs: Dict[str, Any] = {}
     if dag_ids is not None:
         kwargs["dag_ids"] = dag_ids
-    
+
     response = dag_stats_api.get_dag_stats(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]

src/airflow/dataset.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.dataset_api import DatasetApi
@@ -17,7 +17,11 @@ def get_all_functions() -> list[tuple[Callable, str, str]]:
         (get_dag_dataset_queued_event, "get_dag_dataset_queued_event", "Get a queued Dataset event for a DAG"),
         (get_dag_dataset_queued_events, "get_dag_dataset_queued_events", "Get queued Dataset events for a DAG"),
         (delete_dag_dataset_queued_event, "delete_dag_dataset_queued_event", "Delete a queued Dataset event for a DAG"),
-        (delete_dag_dataset_queued_events, "delete_dag_dataset_queued_events", "Delete queued Dataset events for a DAG"),
+        (
+            delete_dag_dataset_queued_events,
+            "delete_dag_dataset_queued_events",
+            "Delete queued Dataset events for a DAG",
+        ),
         (get_dataset_queued_events, "get_dataset_queued_events", "Get queued Dataset events for a Dataset"),
         (delete_dataset_queued_events, "delete_dataset_queued_events", "Delete queued Dataset events for a Dataset"),
     ]
@@ -42,7 +46,7 @@ async def get_datasets(
         kwargs["uri_pattern"] = uri_pattern
     if dag_ids is not None:
         kwargs["dag_ids"] = dag_ids
-    
+
     response = dataset_api.get_datasets(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -82,7 +86,7 @@ async def get_dataset_events(
         kwargs["source_run_id"] = source_run_id
     if source_map_index is not None:
         kwargs["source_map_index"] = source_map_index
-    
+
     response = dataset_api.get_dataset_events(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -96,7 +100,7 @@ async def create_dataset_event(
     }
     if extra is not None:
         event_request["extra"] = extra
-    
+
     response = dataset_api.create_dataset_event(create_dataset_event=event_request)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -131,7 +135,7 @@ async def delete_dag_dataset_queued_events(
     kwargs: Dict[str, Any] = {}
     if before is not None:
         kwargs["before"] = before
-    
+
     response = dataset_api.delete_dag_dataset_queued_events(dag_id=dag_id, **kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -150,7 +154,6 @@ async def delete_dataset_queued_events(
     kwargs: Dict[str, Any] = {}
     if before is not None:
         kwargs["before"] = before
-    
+
     response = dataset_api.delete_dataset_queued_events(uri=uri, **kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
-

src/airflow/eventlog.py

@@ -1,5 +1,5 @@
-from typing import Any, Dict, List, Optional, Union, Callable
 from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.event_log_api import EventLogApi
@@ -62,7 +62,7 @@ async def get_event_logs(
         kwargs["included_events"] = included_events
     if excluded_events is not None:
         kwargs["excluded_events"] = excluded_events
-    
+
     response = event_log_api.get_event_logs(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 
@@ -72,4 +72,3 @@ async def get_event_log(
 ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
     response = event_log_api.get_event_log(event_log_id=event_log_id)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
-

src/airflow/importerror.py

@@ -1,11 +1,10 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import mcp.types as types
 from airflow_client.client.api.import_error_api import ImportErrorApi
 
 from src.airflow.airflow_client import api_client
 
-
 import_error_api = ImportErrorApi(api_client)
 
 
@@ -29,7 +28,7 @@ async def get_import_errors(
         kwargs["offset"] = offset
     if order_by is not None:
         kwargs["order_by"] = order_by
-    
+
     response = import_error_api.get_import_errors(**kwargs)
     return [types.TextContent(type="text", text=str(response.to_dict()))]
 

src/airflow/monitoring.py

@@ -1,4 +1,4 @@
-from typing import List, Union, Callable
+from typing import Callable, List, Union
 
 import mcp.types as types
 from airflow_client.client.api.monitoring_api import MonitoringApi