...
This commit is contained in:
155
_archive/lib/dagu/client.py
Normal file
155
_archive/lib/dagu/client.py
Normal file
@@ -0,0 +1,155 @@
|
||||
import os
|
||||
import requests
|
||||
from requests.auth import HTTPBasicAuth
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
import time
|
||||
|
||||
@dataclass
|
||||
class DAGStatus:
|
||||
name: str
|
||||
status: str
|
||||
group: Optional[str] = None
|
||||
schedule: Optional[str] = None
|
||||
lastRun: Optional[str] = None
|
||||
nextRun: Optional[str] = None
|
||||
pid: Optional[int] = None
|
||||
log: Optional[str] = None
|
||||
requestId: Optional[str] = None
|
||||
params: Optional[str] = None
|
||||
startedAt: Optional[str] = None
|
||||
finishedAt: Optional[str] = None
|
||||
suspended: Optional[bool] = None
|
||||
|
||||
def get_last_run_epoch(self) -> Optional[int]:
|
||||
"""Convert lastRun to epoch time."""
|
||||
return self._convert_to_epoch(self.lastRun)
|
||||
|
||||
def get_next_run_epoch(self) -> Optional[int]:
|
||||
"""Convert nextRun to epoch time."""
|
||||
return self._convert_to_epoch(self.nextRun)
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_epoch(timestamp: Optional[str]) -> Optional[int]:
|
||||
"""Helper method to convert an ISO 8601 timestamp to epoch time."""
|
||||
if timestamp:
|
||||
dt = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))
|
||||
return int(time.mktime(dt.timetuple()))
|
||||
return None
|
||||
|
||||
class DAGuClient:
|
||||
def __init__(self, base_url: str = "http://localhost:8888"):
|
||||
self.base_url = base_url
|
||||
self.auth = self._get_basic_auth()
|
||||
|
||||
def _get_basic_auth(self) -> HTTPBasicAuth:
|
||||
"""Retrieve the Basic Auth credentials from environment variables."""
|
||||
username = os.getenv('DAGU_BASICAUTH_USERNAME')
|
||||
password = os.getenv('DAGU_BASICAUTH_PASSWORD')
|
||||
|
||||
if not username or not password:
|
||||
raise EnvironmentError("Please set the DAGU_BASICAUTH_USERNAME and DAGU_BASICAUTH_PASSWORD environment variables.")
|
||||
|
||||
return HTTPBasicAuth(username, password)
|
||||
|
||||
def list_dags(self) -> List[DAGStatus]:
|
||||
"""Fetches the list of DAGs with their statuses from the DAGu REST API."""
|
||||
try:
|
||||
response = requests.get(f"{self.base_url}/api/v1/dags", auth=self.auth)
|
||||
response.raise_for_status() # Raises an HTTPError for bad responses (4xx or 5xx)
|
||||
dags_data = response.json().get('DAGs', [])
|
||||
|
||||
if isinstance(dags_data, list):
|
||||
return [self._parse_dag(dag) for dag in dags_data]
|
||||
else:
|
||||
print(f"Unexpected response format: {dags_data}")
|
||||
return []
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"Error during request: {e}")
|
||||
return []
|
||||
|
||||
def _parse_dag(self, dag_entry: dict) -> DAGStatus:
|
||||
"""Helper function to parse a DAG's JSON data into a DAGStatus object."""
|
||||
try:
|
||||
dag_data = dag_entry.get("DAG", {})
|
||||
status_data = dag_entry.get("Status", {})
|
||||
|
||||
return DAGStatus(
|
||||
name=dag_data.get("Name"),
|
||||
status=status_data.get("StatusText"),
|
||||
group=dag_data.get("Group"),
|
||||
schedule=(dag_data.get("Schedule", [{}])[0].get("Expression")
|
||||
if dag_data.get("Schedule") else None),
|
||||
lastRun=status_data.get("FinishedAt"),
|
||||
nextRun=None, # Adjust as needed based on your API's response format
|
||||
pid=status_data.get("Pid"),
|
||||
log=status_data.get("Log"),
|
||||
requestId=status_data.get("RequestId"),
|
||||
params=status_data.get("Params"),
|
||||
startedAt=status_data.get("StartedAt"),
|
||||
finishedAt=status_data.get("FinishedAt"),
|
||||
suspended=dag_entry.get("Suspended")
|
||||
)
|
||||
except AttributeError as e:
|
||||
print(f"Error parsing DAG data: {dag_entry}, Error: {e}")
|
||||
return None
|
||||
|
||||
def submit_dag_action(self, name: str, action: str, request_id: Optional[str] = None, params: Optional[str] = None) -> dict:
|
||||
"""Submit an action to a specified DAG.
|
||||
|
||||
Args:
|
||||
name (str): Name of the DAG.
|
||||
action (str): Action to be performed ('start', 'stop', or 'retry').
|
||||
request_id (Optional[str]): Required if action is 'retry'.
|
||||
params (Optional[str]): Parameters for the DAG execution.
|
||||
|
||||
Returns:
|
||||
dict: Response from the API.
|
||||
"""
|
||||
url = f"{self.base_url}/api/v1/dags/{name}"
|
||||
payload = {
|
||||
"action": action,
|
||||
**({"request-id": request_id} if request_id else {}),
|
||||
**({"params": params} if params else {}),
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=payload, auth=self.auth)
|
||||
response.raise_for_status() # Raises an HTTPError for bad responses (4xx or 5xx)
|
||||
return response.json()
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"Error during request: {e}")
|
||||
print(f"Response content: {response.content}")
|
||||
return {}
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
client = DAGuClient()
|
||||
|
||||
# List DAGs
|
||||
try:
|
||||
dags = client.list_dags()
|
||||
for dag in dags:
|
||||
if dag:
|
||||
print(f"DAG Name: {dag.name}, Status: {dag.status}, Group: {dag.group}, "
|
||||
f"Schedule: {dag.schedule}, Last Run: {dag.lastRun}, "
|
||||
f"Next Run: {dag.nextRun}, PID: {dag.pid}, Log: {dag.log}, "
|
||||
f"Request ID: {dag.requestId}, Params: {dag.params}, "
|
||||
f"Started At: {dag.startedAt}, Finished At: {dag.finishedAt}, "
|
||||
f"Suspended: {dag.suspended}")
|
||||
# Example of using helper methods to get epoch times
|
||||
if dag.get_last_run_epoch():
|
||||
print(f"Last Run Epoch: {dag.get_last_run_epoch()}")
|
||||
if dag.get_next_run_epoch():
|
||||
print(f"Next Run Epoch: {dag.get_next_run_epoch()}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
# Submit an action to a DAG (example: start a DAG)
|
||||
try:
|
||||
dag_name = "test11" # Replace with your actual DAG name
|
||||
action_response = client.submit_dag_action(name=dag_name, action="start")
|
||||
print(f"Action Response: {action_response}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
184
_archive/lib/dagu/dag.py
Normal file
184
_archive/lib/dagu/dag.py
Normal file
@@ -0,0 +1,184 @@
|
||||
import os
|
||||
import yaml
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Dict, Optional
|
||||
from server import *
|
||||
|
||||
@dataclass
|
||||
class EnvVariable:
|
||||
key: str
|
||||
value: str
|
||||
|
||||
@dataclass
|
||||
class HandlerCommand:
|
||||
command: str
|
||||
|
||||
@dataclass
|
||||
class Handlers:
|
||||
success: Optional[HandlerCommand] = None
|
||||
failure: Optional[HandlerCommand] = None
|
||||
cancel: Optional[HandlerCommand] = None
|
||||
exit: Optional[HandlerCommand] = None
|
||||
|
||||
@dataclass
|
||||
class RepeatPolicy:
|
||||
repeat: bool
|
||||
intervalSec: int
|
||||
|
||||
@dataclass
|
||||
class Precondition:
|
||||
condition: str
|
||||
expected: str
|
||||
|
||||
@dataclass
|
||||
class Step:
|
||||
name: str
|
||||
command: str
|
||||
script: Optional[str] = None
|
||||
depends: List[str] = field(default_factory=list)
|
||||
description: Optional[str] = None
|
||||
repeatPolicy: Optional[RepeatPolicy] = None
|
||||
|
||||
@dataclass
|
||||
class DAG:
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
schedule: Optional[str] = None
|
||||
group: Optional[str] = None
|
||||
tags: Optional[str] = None # This should be a single string
|
||||
env: Dict[str, str] = field(default_factory=dict)
|
||||
logDir: Optional[str] = None
|
||||
restartWaitSec: Optional[int] = None
|
||||
histRetentionDays: Optional[int] = None
|
||||
delaySec: Optional[int] = None
|
||||
maxActiveRuns: Optional[int] = None
|
||||
params: Optional[List[str]] = field(default_factory=list)
|
||||
preconditions: List[Precondition] = field(default_factory=list)
|
||||
mailOn: Dict[str, bool] = field(default_factory=dict)
|
||||
handlerOn: Handlers = field(default_factory=Handlers)
|
||||
MaxCleanUpTimeSec: Optional[int] = None
|
||||
steps: List[Step] = field(default_factory=list)
|
||||
|
||||
def add_step(self, step: Step):
|
||||
"""Add a step to the DAG."""
|
||||
self.steps.append(step)
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
return {
|
||||
"name": self.name,
|
||||
**({"description": self.description} if self.description else {}),
|
||||
**({"schedule": self.schedule} if self.schedule else {}),
|
||||
**({"group": self.group} if self.group else {}),
|
||||
**({"tags": self.tags} if self.tags else {}),
|
||||
**({"env": [{"key": k, "value": v} for k, v in self.env.items()]} if self.env else {}),
|
||||
**({"logDir": self.logDir} if self.logDir else {}),
|
||||
**({"restartWaitSec": self.restartWaitSec} if self.restartWaitSec else {}),
|
||||
**({"histRetentionDays": self.histRetentionDays} if self.histRetentionDays else {}),
|
||||
**({"delaySec": self.delaySec} if self.delaySec else {}),
|
||||
**({"maxActiveRuns": self.maxActiveRuns} if self.maxActiveRuns else {}),
|
||||
**({"params": " ".join(self.params)} if self.params else {}),
|
||||
**({"preconditions": [{"condition": pc.condition, "expected": pc.expected} for pc in self.preconditions]} if self.preconditions else {}),
|
||||
**({"mailOn": self.mailOn} if self.mailOn else {}),
|
||||
**({"MaxCleanUpTimeSec": self.MaxCleanUpTimeSec} if self.MaxCleanUpTimeSec else {}),
|
||||
**({"handlerOn": {
|
||||
"success": {"command": self.handlerOn.success.command} if self.handlerOn.success else None,
|
||||
"failure": {"command": self.handlerOn.failure.command} if self.handlerOn.failure else None,
|
||||
"cancel": {"command": self.handlerOn.cancel.command} if self.handlerOn.cancel else None,
|
||||
"exit": {"command": self.handlerOn.exit.command} if self.handlerOn.exit else None,
|
||||
}} if any(vars(self.handlerOn).values()) else {}),
|
||||
"steps": [
|
||||
{
|
||||
"name": step.name,
|
||||
"command": step.command,
|
||||
**({"script": step.script} if step.script else {}),
|
||||
**({"depends": step.depends} if step.depends else {}), # Change this back to depends_on if needed
|
||||
**({"description": step.description} if step.description else {}),
|
||||
**({"repeatPolicy": {
|
||||
"repeat": step.repeatPolicy.repeat,
|
||||
"intervalSec": step.repeatPolicy.intervalSec
|
||||
}} if step.repeatPolicy else {}),
|
||||
} for step in self.steps
|
||||
],
|
||||
}
|
||||
|
||||
def to_yaml(self) -> str:
|
||||
return yaml.dump(self.to_dict(), sort_keys=False)
|
||||
|
||||
def new(**kwargs) -> DAG:
|
||||
return DAG(**kwargs)
|
||||
|
||||
# Example usage to create a new DAG
|
||||
if __name__ == "__main__":
|
||||
# Initialize the server with the default DAG directory
|
||||
server = Server()
|
||||
|
||||
# List existing DAGs
|
||||
print("Listing existing DAGs:")
|
||||
dags = server.list_dags()
|
||||
for dag_name in dags:
|
||||
print(f" - {dag_name}")
|
||||
|
||||
# Create a new DAG
|
||||
dag = new(
|
||||
name="example_dag",
|
||||
description="Example DAG to demonstrate functionality",
|
||||
schedule="0 * * * *",
|
||||
group="ExampleGroup",
|
||||
tags="example", # Convert tags to a comma-separated string
|
||||
env={
|
||||
"LOG_DIR": "${HOME}/logs",
|
||||
"PATH": "/usr/local/bin:${PATH}"
|
||||
},
|
||||
logDir="${LOG_DIR}",
|
||||
restartWaitSec=60,
|
||||
histRetentionDays=3,
|
||||
delaySec=1,
|
||||
maxActiveRuns=1,
|
||||
params=["param1", "param2"],
|
||||
preconditions=[
|
||||
Precondition(condition="`echo $2`", expected="param2")
|
||||
],
|
||||
mailOn={"failure": True, "success": True},
|
||||
MaxCleanUpTimeSec=300,
|
||||
handlerOn=Handlers(
|
||||
success=HandlerCommand(command="echo succeed"), # Convert to map structure
|
||||
failure=HandlerCommand(command="echo failed"), # Convert to map structure
|
||||
cancel=HandlerCommand(command="echo canceled"), # Convert to map structure
|
||||
exit=HandlerCommand(command="echo finished") # Convert to map structure
|
||||
)
|
||||
)
|
||||
|
||||
# Add steps to the DAG
|
||||
dag.add_step(Step(
|
||||
name="pull_data",
|
||||
command="sh",
|
||||
script="echo `date '+%Y-%m-%d'`",
|
||||
))
|
||||
|
||||
dag.add_step(Step(
|
||||
name="cleanse_data",
|
||||
command="echo cleansing ${DATA_DIR}/${DATE}.csv",
|
||||
depends=["pull_data"] # Ensure this is the correct key
|
||||
))
|
||||
|
||||
dag.add_step(Step(
|
||||
name="transform_data",
|
||||
command="echo transforming ${DATA_DIR}/${DATE}_clean.csv",
|
||||
depends=["cleanse_data"] # Ensure this is the correct key
|
||||
))
|
||||
|
||||
dag.add_step(Step(
|
||||
name="A task",
|
||||
command="main.sh",
|
||||
repeatPolicy=RepeatPolicy(repeat=True, intervalSec=60)
|
||||
))
|
||||
|
||||
# Save the new DAG as a YAML file
|
||||
server.create_dag(dag)
|
||||
print(f"DAG '{dag.name}' created and saved and started.")
|
||||
|
||||
# List DAGs again to see the newly created one
|
||||
print("\nListing updated DAGs:")
|
||||
dags = server.list_dags()
|
||||
for dag_name in dags:
|
||||
print(f" - {dag_name}")
|
||||
51
_archive/lib/dagu/server.py
Normal file
51
_archive/lib/dagu/server.py
Normal file
@@ -0,0 +1,51 @@
|
||||
import os
|
||||
import yaml
|
||||
import glob
|
||||
from typing import List
|
||||
from dag import DAG
|
||||
from client import *
|
||||
|
||||
# Assuming the following classes have already been defined:
|
||||
# - DAG (for creating and managing DAG structures)
|
||||
# - Step
|
||||
# - Handlers
|
||||
# - RepeatPolicy
|
||||
# - Precondition
|
||||
|
||||
class Server:
|
||||
def __init__(self, dag_dir: str = "~/hero/var/dagu/dags/"):
|
||||
self.dag_dir = os.path.expanduser(dag_dir)
|
||||
os.makedirs(self.dag_dir, exist_ok=True) # Ensure the directory exists
|
||||
|
||||
def list_dags(self) -> List[str]:
|
||||
"""Lists the DAGs in the directory."""
|
||||
dag_files = glob.glob(os.path.join(self.dag_dir, "*.yaml"))
|
||||
return [os.path.splitext(os.path.basename(dag_file))[0] for dag_file in dag_files]
|
||||
|
||||
def delete_dag(self, name: str) -> bool:
|
||||
"""Deletes a DAG file based on its name."""
|
||||
dag_file = os.path.join(self.dag_dir, f"{name}.yaml")
|
||||
if os.path.exists(dag_file):
|
||||
os.remove(dag_file)
|
||||
return True
|
||||
else:
|
||||
print(f"DAG '{name}' does not exist.")
|
||||
return False
|
||||
|
||||
def create_dag(self, dag:DAG, start:bool = True) -> bool:
|
||||
"""Creates a new DAG and saves it as a YAML file."""
|
||||
dag_file = os.path.join(self.dag_dir, f"{dag.name}.yaml")
|
||||
with open(dag_file, 'w') as file:
|
||||
yaml.dump(dag.to_dict(), file, sort_keys=False)
|
||||
if start:
|
||||
self.start_dag(dag.name)
|
||||
return True
|
||||
|
||||
def start_dag(self,dag_name:str) -> bool:
|
||||
client = DAGuClient()
|
||||
action_response = client.submit_dag_action(name=dag_name, action="start")
|
||||
|
||||
def stop_dag(self,dag_name:str) -> bool:
|
||||
client = DAGuClient()
|
||||
action_response = client.submit_dag_action(name=dag_name, action="stop")
|
||||
|
||||
Reference in New Issue
Block a user