def set_connection(self, conn: Connection):
self.delete_connection(conn)
with set_airflow_db(self.db_path, self.fernet_key) as db:
db.set_connection(
conn_id=typhoon_airflow_conn_name(conn.conn_id),
**conn.get_connection_params().__dict__
)
def build_all_dags_airflow(remote: Optional[str],
matching: Optional[str] = None):
airflow_home = Path(os.environ['AIRFLOW_HOME'])
target_folder: Path = airflow_home / 'dags/typhoon_managed/'
target_folder.mkdir(parents=True, exist_ok=True)
rmtree(str(target_folder), ignore_errors=True)
target_folder.mkdir()
print('Build all DAGs...')
dags = load_dag_definitions(ignore_errors=True)
for dag, _ in dags:
print(f'Found DAG {dag.name}')
if not matching or re.match(matching, dag.name):
dag_target_folder = target_folder / dag.name
dag_target_folder.mkdir(parents=True, exist_ok=True)
init_path = (dag_target_folder / '__init__.py')
init_path.write_text('')
store: AirflowMetadataStore = Settings.metadata_store()
start_date = None
with set_airflow_db(store.db_path, store.fernet_key) as db:
dag_run = db.get_first_dag_run(dag.name)
if dag_run:
start_date = dag_run.execution_date.replace(tzinfo=None)
compiled_dag = AirflowDagFile(dag,
start_date=start_date,
debug_mode=remote is None).render()
(dag_target_folder / f'{dag.name}.py').write_text(compiled_dag)
tasks_code = TasksFile(dag.tasks).render()
(dag_target_folder / 'tasks.py').write_text(tasks_code)
for component, _ in load_components(ignore_errors=False, kind='all'):
print(f'Building component {component.name}...')
build_component_for_airflow(component, target_folder)
def delete_variable(self, variable: Union[str, Variable]):
var_name = variable.id if isinstance(variable, Variable) else variable
info_var = typhoon_airflow_variable_name_for_info(var_name)
value_var = typhoon_airflow_variable_name_for_value(var_name)
with set_airflow_db(self.db_path, self.fernet_key) as db:
db.delete_variable(info_var)
db.delete_variable(value_var)
def delete_connection(self, conn: Union[str, Connection]):
conn_id = conn.conn_id if isinstance(conn, Connection) else conn
af_name = typhoon_airflow_conn_name(conn_id)
with set_airflow_db(self.db_path, self.fernet_key) as db:
try:
db.delete_connection(af_name)
except UnmappedInstanceError:
pass
def get_variable(self, variable_id: str) -> Variable:
info_var = typhoon_airflow_variable_name_for_info(variable_id)
value_var = typhoon_airflow_variable_name_for_value(variable_id)
with set_airflow_db(self.db_path, self.fernet_key) as db:
info = db.get_variable(info_var)
if info is None:
raise MetadataObjectNotFound(f'Variable "{variable_id}" is not set')
info = json.loads(info)
value = db.get_variable(value_var)
return Variable(id=info['id'], type=VariableType(info['type']), contents=value)
def get_variables(self, to_dict: bool = False) -> List[Union[dict, Variable]]:
result = []
with set_airflow_db(self.db_path, self.fernet_key) as db:
for af_var in db.get_variables():
if af_var.key.startswith('typhoon:info#'):
info = json.loads(af_var.val)
var_id = info['id']
value_var = typhoon_airflow_variable_name_for_value(var_id)
contents = db.get_variable(value_var)
if contents is None:
raise MetadataObjectNotFound(f'Variable "{var_id}" does not have a value set')
var = Variable(id=var_id, type=VariableType(info['type']), contents=contents)
if to_dict:
var = var.dict_contents()
result.append(var)
return result
def get_connection(self, conn_id: str) -> Connection:
af_name = typhoon_airflow_conn_name(conn_id)
with set_airflow_db(self.db_path, self.fernet_key) as db:
af_conn = db.get_connection(af_name)
if af_conn is None:
raise MetadataObjectNotFound(f'Connection "{conn_id}" is not set')
conn = Connection(
conn_id=conn_id,
conn_type=af_conn.conn_type,
host=af_conn.host,
port=af_conn.port,
login=af_conn.login,
password=af_conn.password,
schema=af_conn.schema,
extra=af_conn.extra_dejson,
)
return conn
def set_variable(self, variable: Variable):
"""
Sets Airflow variables in JSON like:
{
"id": "variable_name",
"type": "variable_type",
"contents": "contents",
}
Following the same schema as the typhoon variable
"""
info_var = typhoon_airflow_variable_name_for_info(variable.id)
value_var = typhoon_airflow_variable_name_for_value(variable.id)
info = variable.dict_contents()
del info['contents']
with set_airflow_db(self.db_path, self.fernet_key) as db:
db.set_variable(info_var, json.dumps(info))
db.set_variable(value_var, variable.contents)
def get_connections(self, to_dict: bool = False) -> List[Union[dict, Connection]]:
result = []
with set_airflow_db(self.db_path, self.fernet_key) as db:
for af_conn in db.get_connections():
if af_conn.conn_id.startswith('typhoon#'):
conn = Connection(
conn_id=af_conn.conn_id.split('#')[1],
conn_type=af_conn.conn_type,
host=af_conn.host,
port=af_conn.port,
login=af_conn.login,
password=af_conn.password,
schema=af_conn.schema,
extra=af_conn.extra_dejson,
)
if to_dict:
conn = conn.__dict__
result.append(conn)
return result