airflow_tutorial/dag_code_generator.py at main · HWALIMLEE/airflow_tutorial · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import os
import click
from PyInquirer import prompt, Validator, ValidationError, style_from_dict, Token


class DagNameValidator(Validator):
    def validate(self, document):
        if len(document):
            return True
        else:
            raise ValidationError


style = style_from_dict({
    Token.QuestionMark: '#fac731 bold',
    Token.Answer: '#4688f1 bold',
    Token.Instruction: '',  # default
    Token.Separator: '#cc5454',
    Token.Selected: '#0abf5b',  # default
    Token.Pointer: '#673ab7 bold',
    Token.Question: '',
})


def ask_default_args():
    questions = [
        {
            'type': 'input',
            'name': 'dag_name',
            'message': 'dag name',
        },
        {
            'type': 'input',
            'name': 'description',
            'message': 'description',
            'default': ''
        },
        {
            'type': 'input',
            'name': 'owner',
            'message': 'owner',
            'default': 'hwalim'
        },
        {
            'type': 'input',
            'name': 'retries',
            'message': 'retries',
            'default': '3'
        },
        {
            'type': 'input',
            'name': 'retry_delay',
            'message': 'retry delay minutes',
            'default': '1'
        },
        {
            'type': 'input',
            'name': 'start_date',
            'message': 'start date (ex. 2023, 1, 1, 17, 30)',
            'default': '2023, 1, 1, 17, 30'
        },
        {
            'type': 'input',
            'name': 'schedule_interval',
            'message': 'schedule interval'
        },
        {
            'type': 'input',
            'name': 'catchup',
            'message': 'catchup',
            'default': 'False'
        },
        {
            'type': 'input',
            'name': 'tags',
            'message': 'tags',
            'default': 'etl'
        }
    ]
    default_args = prompt(questions, style=style)
    return default_args


def create_dag_file(conf):
    conf["catchup"] = conf.get("catchup") == "True"
    conf["tags"] = str(conf.get("tags").split(','))
    os.system(
        'bash '+os.path.join(os.path.dirname(os.path.abspath(__file__)), 'support.sh')+' "{dag_name}" "{owner}" "{retries}" "{retry_delay}" "{description}" "{schedule_interval}" "{start_date}" "{catchup}" "{tags}"'.format(**conf)
        )


@click.command()
def main():
    default_info = ask_default_args()
    create_dag_file(default_info)


if __name__ == '__main__':
    main()