How to create shard&index in airflow mongohook? - airflow

I want to run mongo command with mongohook of airflow. How can I do it?
sh.shardCollection(db_name +, { _id: "hashed" }, false, { numInitialChunks: 128 });
db.collection.createIndex({ "field": 1 }, { field: true });

The pymongo client which the Mongohook provided in Airflow uses doesn't support the sh.shardCollection command in your script.
Though the createIndex collection method is supported in the pymongo client.
I recommend anyway to install the mongosh CLI binary and bake it into your container image for your workers.
You can write your shell command to a script such as /dags/templates/mongo-admin-create-index.js or some other location that it can be found.
Then can implement a custom operator using the SubprocessHook to run mongosh CLI command such as:
mongosh -f {mongosh_script} {db_address}
This custom operator would be along these lines
from airflow.compat.functools import cached_property
from airflow.hooks.subprocess import SubprocessHook
from airflow.providers.mongo.hooks import MongoHook
class MongoshScriptOperator(BaseOperator):
template_fields: Sequence[str] = ('mongosh_script')
def __init__(
self,
*,
mongosh_script: str,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.mongosh_script = mongosh_script
#cached_property
def subprocess_hook(self):
"""Returns hook for running the shell command"""
return SubprocessHook()
def execute(self):
"""Executes a mongosh script"""
mh = MongoHook(self.conn_id)
self.subprocess_hook.run_command(
command=['mongosh', '-f', self.mongosh_script, mh.uri],
)
When creating the DagNode, you can pass the location of the script to your custom operator.

Related

How to dynamically generate tasks via an XCOM dictionary for a a non TaskInstance operator?

I have a dag that outputs a dictionary (map), passed through an XCom.
I want to generate a number of tasks according to the keys of said XCom dictionary.
This is how the dictionary looks in the xcom :
{
"F1": {
"source": {
"project": "legacy_project",
"bucket": "legacy_bucket",
"prefix": "prefix/{{ds_nodash}}/F1",
"files": [
"file_1.csv"
]
},
"destination": {
"project_id": "new_project_1",
"bucket": "new_bucket_1",
"prefix": "DTM/F1/{{ds_nodash}}"
}
},
"F2": {
"source": {
"project": "legacy_project",
"bucket": "legacy_bucket",
"prefix": "prefix/{{ds_nodash}}/F2",
"files": [
"file_1.csv"
]
},
"destination": {
"project_id": "new_project_2",
"bucket": "new_bucket_2",
"prefix": "DTM/F2/{{ds_nodash}}"
}
}
// ...
}
Notice that I used {{ds_nodash}} in the generated xcom in order for it to be replaced by the date of execution of the dag.
What I want, is to create a task for each of F1,F2,...,Fn and input the F1["source"]["bucket"], F1["source"]["prefix"] into a GCSObjectsWithPrefixExistenceSensor, with a custom task_id for each task..
Something like this :
for f_key in <xcom_output_dict>:
GCSObjectsWithPrefixExistenceSensor(
task_id = f"{f_key}_sensor",
bucket = f_key["source"]["bucket"]
prefix = f_key["source"]["prefix"]
)
I tried reading up on expand and partial, but it's really unclear how one can pass the dictionary contents such as bucket and prefix to the task operator...
EDIT :
One of my trials :
from airflow import DAG, XComArg
from airflow.contrib.sensors.gcs_sensor import (
GCSObjectsWithPrefixExistenceSensor,
)
...
generate_prefix_existance_sensor_kwargs = PythonOperator(
task_id="generate_prefix_existance_sensor_kwargs",
python_callable=gen_prefix_existance_sensor_kwargs,
provide_context=True,
op_kwargs={"effective_migration_map": "effective_migration_map"},
)
sensor_files = GCSObjectsWithPrefixExistenceSensor.expand_kwargs(
input=XComArg(generate_prefix_existance_sensor_kwargs),
)
(
... # Some other tasks
>> generate_prefix_existance_sensor_kwargs
>> sensor_files
)
With the code to the kwargs generation function
def gen_prefix_existance_sensor_kwargs(ti, **kwargs) -> List[Dict[str, str]]:
effective_migration_map = load_xcom_via_kwarg(
ti=ti,
key="effective_migration_map",
default_value="effective_migration_map",
kwargs=kwargs,
)
args_dicts = []
for f_key in effective_migration_map:
args_dicts.append(
{
"task_id": f_key,
"bucket": effective_migration_map[f_key]["source"]["bucket"],
"prefix": effective_migration_map[f_key]["source"]["prefix"],
}
)
return args_dicts
But I get :
AttributeError: type object 'GCSObjectsWithPrefixExistenceSensor' has no attribute 'expand_kwargs'
Which is weird because I saw an Astronomer.io video where it is used. I guess it isn't implemented yet in Airflow 2.3.4?
EDIT 1 :
I create the wrapper for the operator :
class GCSObjectsWithPrefixExistenceSensorWrapper(
GCSObjectsWithPrefixExistenceSensor
):
"""This class is a temporary work around to using expand_kwargs(),
as expand() can only take one argument,
we must create a wrapper around every operator we use
"""
def __init__(
self,
src_or_dest: str,
inp_parameters: Tuple[str, Dict[str, Dict[str, str]]],
**kwargs
):
if src_or_dest not in ["source", "destination"]:
raise TypeError(
"Bad argument for src_or_dest, must be either 'source' or 'destination'"
)
else:
bucket = inp_parameters[1][src_or_dest]["bucket"]
prefix = inp_parameters[1][src_or_dest]["prefix"]
super().__init__(
bucket=bucket,
prefix=prefix,
**kwargs,
)
I add it to my dag ..
updated_map_2 = match_data_with_migration_map(
src_mig_map=updated_map_1,
files_and_prefixes="{{ti.xcom_pull('list_files_and_prefixes')}}",
)
GCSObjectsWithPrefixExistenceSensorWrapper.partial(
task_id="sensor_files",
src_or_dest = "source",
impersonation_chain=IMPERSONATED_SERVICE_ACCOUNT,
).expand(inp_parameters=updated_map_2.output)
I get this error :
AttributeError: 'dict' object has no attribute 'output'
EDIT 2 :
The previous task to this operator was not decorated correctly, so it created the "EDIT 1" problems
I've tried to solve your issue with dynamic task mapping. However this solution has the down side that if you pass more parameters into the expand() it will make a cross-product
So my solution is to create a custom class that inherits from GCSObjectsWithPrefixExistenceSensor and the init function assigns correct values:
class CustomExistenceSensor(GCSObjectsWithPrefixExistenceSensor):
def __init__(self, inp_params, **kwargs):
bucket , prefix = inp_params
super().__init__(bucket=bucket, prefix=prefix, **kwargs)
By using this class we can now expand the sensor and pass an iterable of parameters (bucket, prefix):
from airflow.decorators import dag, task
from datetime import datetime
from airflow.providers.google.cloud.sensors.gcs import GCSObjectsWithPrefixExistenceSensor
class CustomExistenceSensor(GCSObjectsWithPrefixExistenceSensor):
def __init__(self, inp_params, **kwargs):
bucket , prefix = inp_params
super().__init__(bucket=bucket, prefix=prefix, **kwargs)
#dag(
schedule=None,
start_date=datetime(2022, 10, 21, hour=8),
catchup=False,
tags=['demo'],
)
def template_dag():
"""
### Template dag"""
#task()
def example_func():
pass
t1 = example_func()
# Below just to test the sensor works
# t3 = GCSObjectsWithPrefixExistenceSensor(
# task_id="test_check",
# bucket="text_stack_bucket",
# prefix="test_prefix"
# prefix="prefix",
# timeout=3)
# t4 works but the issue here is that we can pass only array with one parameter if we pass more airflow will make a cross product
# t4 = GCSObjectsWithPrefixExistenceSensor.partial(task_id="test_check_dynamic", bucket="text_stack_bucket", timeout=3).expand(prefix=["test_prefix", "prefix"])
t5 = CustomExistenceSensor.partial(task_id="test_custom_class", timeout=2).expand(inp_params=[("text_stack_bucket", "test_prefix"), ("next_existing", "prefix")])
t1 >> t2 >> t5
dag = template_dag()
So if you can amend a bit your code and make it return a list in xcom so that you can invoke it like so:
task = CustomExistenceSensor.partial(task_id='dynamic_tasks').expand(inp_params=previous_task.output)

How to pass parameters from google composer to dataflow template

I'm trying to pass parameter from google composer into a dataflow template as following way, but it does not work.
# composer code
trigger_dataflow = DataflowTemplateOperator(
task_id="trigger_dataflow",
template="gs://mybucket/my_template",
dag=dag,
job_name='appsflyer_events_daily',
parameters={
"input": f'gs://my_bucket/' + "{{ ds }}" + "/*.gz"
}
)
# template code
class UserOptions(PipelineOptions):
#classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
'--input',
default='gs://my_bucket/*.gz',
help='path of input file')
def main():
pipeline_options = PipelineOptions()
user_options = pipeline_options.view_as(UserOptions)
p = beam.Pipeline(options=pipeline_options)
lines = (
p
| MatchFiles(user_options.input)
)
You can pass like following.
DataflowTemplateOperator(,
task_id="task1",
template=get_variable_value("template"),
on_failure_callback=update_job_message,
parameters={
"fileBucket": get_variable_value("file_bucket"),
"basePath": get_variable_value("path_input"),
"Day": "{{ json.loads(ti.xcom_pull(key=run_id))['day'] }}",
},
)
We are using Java and in Dataflow jobs we have option class get and set like following
public interface MyOptions extends CommonOptions {
#Description("The output bucket")
#Validation.Required
ValueProvider<String> getFileBucket();
void setFileBucket(ValueProvider<String> value);
}
We need to create template for this dataflow jobs and that template will be trigger by composer dag.
Moving from dataflow classic template to flex template fixed the issue.

How to access Xcom value in a non airflow operator python function

I have a stored XCom value that I wanted to pass to another python function which is not called using PythonOperator.
def sql_file_template():
<some code which uses xcom variable>
def call_stored_proc(**kwargs):
#project = kwargs['row_id']
print("INSIDE CALL STORE PROC ------------")
query = """CALL `{0}.dataset_name.store_proc`(
'{1}' # source table
, ['{2}'] # row_ids
, '{3}' # pivot_col_name
, '{4}' # pivot_col_value
, 100 # max_columns
, 'MAX' # aggregation
);"""
query = query.format(kwargs['project'],kwargs['source_tbl'] ,kwargs['row_id'],kwargs['pivot_col'],kwargs['pivot_val'])
job = client.query(query, location="US")
for result in job.result():
task_instance = kwargs['task_instance']
task_instance.xcom_push(key='query_string', value=result)
print result
return result
bq_cmd = PythonOperator (
task_id= 'task1'
provide_context= True,
python_callable= call_stored_proc,
op_kwargs= {'project' : project,
'source_tbl' : source_tbl,
'row_id' : row_id,
'pivot_col' : pivot_col,
'pivot_val' : pivot_val
},
dag= dag
)
dummy_operator >> bq_cmd
sql_file_template()
The output of stored proc is a string which is captured using xcom.
Now I would like to pass this value to some python function sql_file_template without using PythonOperator.
As per Airflow documentation xcom can be accessed only between tasks.
Can anyone help on this?
If you have access to the Airflow installation you'd like to query (configuration, database access, and code) you can use Airflow's airflow.models.XCom:get_one class method:
from datetime import datetime
from airflow.models import XCom
execution_date = datetime(2020, 8, 28)
xcom_value = XCom.get_one(execution_date=execution_date,
task_id="the_task_id",
dag_id="the_dag_id")
So you want to access XCOM outside Airflow (probably a different project / module, without creating any Airflow DAGs / tasks)?
Airflow uses SQLAlchemy for mapping all it's models (including XCOM) to corresponding SQLAlchemy backend (meta-db) tables
Therefore this can be done in two ways
Leverage Airflow's SQLAlchemy model
(without having to create a task or DAG). Here's an untested code snippet for reference
from typing import List
from airflow.models import XCom
from airflow.settings import Session
from airflow.utils.db import provide_session
from pendulum import Pendulum
#provide_session
def read_xcom_values(dag_id: str,
task_id: str,
execution_date: Pendulum,
session: Optional[Session]) -> List[str]:
"""
Function that reads and returns 'values' of XCOMs with given filters
:param dag_id:
:param task_id:
:param execution_date: datetime object
:param session: Airflow's SQLAlchemy Session (this param must not be passed, it will be automatically supplied by
'#provide_session' decorator)
:return:
"""
# read XCOMs
xcoms: List[XCom] = session.query(XCom).filter(
XCom.dag_id == dag_id, XCom.task_id == task_id,
XCom.execution_date == execution_date).all()
# retrive 'value' fields from XCOMs
xcom_values: List[str] = list(map(lambda xcom: xcom.value, xcoms))
return xcom_values
Do note that since it is importing airflow packages, it still requires working airflow installation on python classpath (as well as connection to backend-db), but here we are not creating any tasks or dags (this snippet can be run in a standalone python file)
For this snippet, I have referred to views.py which is my favorite place to peek into Airflow's SQLAlchemy magic
Directly query Airflow's SQLAlchemy backend meta-db
Connect to meta db and run this query
SELECT value FROM xcom WHERE dag_id='' AND task_id='' AND ..

What is the Deno equivalent of process.argv in Node.js?

When working with NodeJS, I can pass the arguments to a Node script like this:
$ node node-server.js arg1 arg2=arg2-val arg3
And can get the arguments like so:
// print process.argv
process.argv.forEach(function (val, index, array) {
console.log(index + ': ' + val);
});
//Output
0: node
1: /Users/umar/work/node/node-server.js
2: arg1
3: arg2=arg2-val
4: arg3
How to get the command-line arguments in Deno?
Some experts suggested me to solve the problem by answers to the question
Deno executable path ~ process.argv[0]:
Deno.execPath()
File URL of executed script ~ process.argv[1]:
Deno.mainModule
You can use path.fromFileUrl for conversions of URL to path string:
import { fromFileUrl } from "https://deno.land/std#0.55.0/path/mod.ts";
const modPath = fromFileUrl(import.meta.url)
Command-line arguments ~ process.argv.slice(2):
Deno.args
Example
deno run --allow-read test.ts -foo -bar=baz 42
Sample output (Windows):
Deno.execPath(): <scoop path>\apps\deno\current\deno.exe
import.meta.url: file:///C:/path/to/project/test.ts
as path: C:\path\to\project\test.ts
Deno.args: [ "-foo", "-bar=baz", "42" ]
To get your script’s CLI arguments in Deno, just use Deno.args:
> deno run ./code.ts foo bar
console.log(Deno.args); // ['foo', 'bar']
If you need something identical to Node's process.argv for compatibility reasons, use the official 'node' shim:
import process from 'https://deno.land/std#0.120.0/node/process.ts'
console.log(process.argv); // ['/path/to/deno', '/path/to/code.ts', 'foo', 'bar']
For illustrative purposes, if you wanted to manually construct a process.argv-style array (without using the official 'node' shim) you could do this:
import { fromFileUrl } from "https://deno.land/std#0.120.0/path/mod.ts";
const argv = [
Deno.execPath(),
fromFileUrl(Deno.mainModule),
...Deno.args,
]
console.log(argv); // ['/path/to/deno', '/path/to/code.ts', 'foo', 'bar']

Removing Airflow task logs

I'm running 5 DAG's which have generated a total of about 6GB of log data in the base_log_folder over a months period. I just added a remote_base_log_folder but it seems it does not exclude logging to the base_log_folder.
Is there anyway to automatically remove old log files, rotate them or force airflow to not log on disk (base_log_folder) only in remote storage?
Please refer https://github.com/teamclairvoyant/airflow-maintenance-dags
This plugin has DAGs that can kill halted tasks and log-cleanups.
You can grab the concepts and can come up with a new DAG that can cleanup as per your requirement.
We remove the Task logs by implementing our own FileTaskHandler, and then pointing to it in the airflow.cfg. So, we overwrite the default LogHandler to keep only N task logs, without scheduling additional DAGs.
We are using Airflow==1.10.1.
[core]
logging_config_class = log_config.LOGGING_CONFIG
log_config.LOGGING_CONFIG
BASE_LOG_FOLDER = conf.get('core', 'BASE_LOG_FOLDER')
FOLDER_TASK_TEMPLATE = '{{ ti.dag_id }}/{{ ti.task_id }}'
FILENAME_TEMPLATE = '{{ ti.dag_id }}/{{ ti.task_id }}/{{ ts }}/{{ try_number }}.log'
LOGGING_CONFIG = {
'formatters': {},
'handlers': {
'...': {},
'task': {
'class': 'file_task_handler.FileTaskRotationHandler',
'formatter': 'airflow.job',
'base_log_folder': os.path.expanduser(BASE_LOG_FOLDER),
'filename_template': FILENAME_TEMPLATE,
'folder_task_template': FOLDER_TASK_TEMPLATE,
'retention': 20
},
'...': {}
},
'loggers': {
'airflow.task': {
'handlers': ['task'],
'level': JOB_LOG_LEVEL,
'propagate': False,
},
'airflow.task_runner': {
'handlers': ['task'],
'level': LOG_LEVEL,
'propagate': True,
},
'...': {}
}
}
file_task_handler.FileTaskRotationHandler
import os
import shutil
from airflow.utils.helpers import parse_template_string
from airflow.utils.log.file_task_handler import FileTaskHandler
class FileTaskRotationHandler(FileTaskHandler):
def __init__(self, base_log_folder, filename_template, folder_task_template, retention):
"""
:param base_log_folder: Base log folder to place logs.
:param filename_template: template filename string.
:param folder_task_template: template folder task path.
:param retention: Number of folder logs to keep
"""
super(FileTaskRotationHandler, self).__init__(base_log_folder, filename_template)
self.retention = retention
self.folder_task_template, self.folder_task_template_jinja_template = \
parse_template_string(folder_task_template)
#staticmethod
def _get_directories(path='.'):
return next(os.walk(path))[1]
def _render_folder_task_path(self, ti):
if self.folder_task_template_jinja_template:
jinja_context = ti.get_template_context()
return self.folder_task_template_jinja_template.render(**jinja_context)
return self.folder_task_template.format(dag_id=ti.dag_id, task_id=ti.task_id)
def _init_file(self, ti):
relative_path = self._render_folder_task_path(ti)
folder_task_path = os.path.join(self.local_base, relative_path)
subfolders = self._get_directories(folder_task_path)
to_remove = set(subfolders) - set(subfolders[-self.retention:])
for dir_to_remove in to_remove:
full_dir_to_remove = os.path.join(folder_task_path, dir_to_remove)
print('Removing', full_dir_to_remove)
shutil.rmtree(full_dir_to_remove)
return FileTaskHandler._init_file(self, ti)
Airflow maintainers don't think truncating logs is a part of airflow core logic, to see this, and then in this issue, maintainers suggest to change LOG_LEVEL avoid too many log data.
And in this PR, we can learn how to change log level in airflow.cfg.
good luck.
I know it sounds savage, but have you tried pointing base_log_folder to /dev/null? I use Airflow as a part of a container, so I don't care about the files either, as long as the logger pipe to STDOUT as well.
Not sure how well this plays with S3 though.
For your concrete problems, I have some suggestions.
For those, you would always need a specialized logging config as described in this answer: https://stackoverflow.com/a/54195537/2668430
automatically remove old log files and rotate them
I don't have any practical experience with the TimedRotatingFileHandler from the Python standard library yet, but you might give it a try:
https://docs.python.org/3/library/logging.handlers.html#timedrotatingfilehandler
It not only offers to rotate your files based on a time interval, but if you specify the backupCount parameter, it even deletes your old log files:
If backupCount is nonzero, at most backupCount files will be kept, and if more would be created when rollover occurs, the oldest one is deleted. The deletion logic uses the interval to determine which files to delete, so changing the interval may leave old files lying around.
Which sounds pretty much like the best solution for your first problem.
force airflow to not log on disk (base_log_folder), but only in remote storage?
In this case you should specify the logging config in such a way that you do not have any logging handlers that write to a file, i.e. remove all FileHandlers.
Rather, try to find logging handlers that send the output directly to a remote address.
E.g. CMRESHandler which logs directly to ElasticSearch but needs some extra fields in the log calls.
Alternatively, write your own handler class and let it inherit from the Python standard library's HTTPHandler.
A final suggestion would be to combine both the TimedRotatingFileHandler and setup ElasticSearch together with FileBeat, so you would be able to store your logs inside ElasticSearch (i.e. remote), but you wouldn't store a huge amount of logs on your Airflow disk since they will be removed by the backupCount retention policy of your TimedRotatingFileHandler.
Usually apache airflow grab the disk space due to 3 reasons
1. airflow scheduler logs files
2. mysql binaly logs [Major]
3. xcom table records.
To make it clean up on regular basis I have set up a dag which run on daily basis and cleans the binary logs and truncate the xcom table to make the disk space free
You also might need to install [pip install mysql-connector-python].
To clean up scheduler log files I do delete them manually two times in a week to avoid the risk of logs deleted which needs to be required for some reasons.
I clean the logs files by [sudo rm -rd airflow/logs/] command.
Below is my python code for reference
'
"""Example DAG demonstrating the usage of the PythonOperator."""
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
from airflow.utils.dates import days_ago
from airflow.operators.bash import BashOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
args = {
'owner': 'airflow',
'email_on_failure':True,
'retries': 1,
'email':['Your Email Id'],
'retry_delay': timedelta(minutes=5)
}
dag = DAG(
dag_id='airflow_logs_cleanup',
default_args=args,
schedule_interval='#daily',
start_date=days_ago(0),
catchup=False,
max_active_runs=1,
tags=['airflow_maintenance'],
)
def truncate_table():
import mysql.connector
connection = mysql.connector.connect(host='localhost',
database='db_name',
user='username',
password='your password',
auth_plugin='mysql_native_password')
cursor = connection.cursor()
sql_select_query = """TRUNCATE TABLE xcom"""
cursor.execute(sql_select_query)
connection.commit()
connection.close()
print("XCOM Table truncated successfully")
def delete_binary_logs():
import mysql.connector
from datetime import datetime
date = datetime.today().strftime('%Y-%m-%d')
connection = mysql.connector.connect(host='localhost',
database='db_name',
user='username',
password='your_password',
auth_plugin='mysql_native_password')
cursor = connection.cursor()
query = 'PURGE BINARY LOGS BEFORE ' + "'" + str(date) + "'"
sql_select_query = query
cursor.execute(sql_select_query)
connection.commit()
connection.close()
print("Binary logs deleted successfully")
t1 = PythonOperator(
task_id='truncate_table',
python_callable=truncate_table, dag=dag
)
t2 = PythonOperator(
task_id='delete_binary_logs',
python_callable=delete_binary_logs, dag=dag
)
t2 << t1
'
I am surprized but it worked for me. Update your config as below:
base_log_folder=""
It is test in minio and in s3.
Our solution looks a lot like Franzi's:
Running on Airflow 2.0.1 (py3.8)
Override default logging configuration
Since we use a helm chart for airflow deployment it was easiest to push an env there, but it can also be done in the airflow.cfg or using ENV in dockerfile.
# Set custom logging configuration to enable log rotation for task logging
AIRFLOW__LOGGING__LOGGING_CONFIG_CLASS: "airflow_plugins.settings.airflow_local_settings.DEFAULT_LOGGING_CONFIG"
Then we added the logging configuration together with the custom log handler to a python module we build and install in the docker image. As described here: https://airflow.apache.org/docs/apache-airflow/stable/modules_management.html
Logging configuration snippet
This is only a copy on the default from the airflow codebase, but then the task logger gets a different handler.
DEFAULT_LOGGING_CONFIG: Dict[str, Any] = {
'version': 1,
'disable_existing_loggers': False,
'formatters': {
'airflow': {'format': LOG_FORMAT},
'airflow_coloured': {
'format': COLORED_LOG_FORMAT if COLORED_LOG else LOG_FORMAT,
'class': COLORED_FORMATTER_CLASS if COLORED_LOG else 'logging.Formatter',
},
},
'handlers': {
'console': {
'class': 'airflow.utils.log.logging_mixin.RedirectStdHandler',
'formatter': 'airflow_coloured',
'stream': 'sys.stdout',
},
'task': {
'class': 'airflow_plugins.log.rotating_file_task_handler.RotatingFileTaskHandler',
'formatter': 'airflow',
'base_log_folder': os.path.expanduser(BASE_LOG_FOLDER),
'filename_template': FILENAME_TEMPLATE,
'maxBytes': 10485760, # 10MB
'backupCount': 6,
},
...
RotatingFileTaskHandler
And finally the custom handler which is just a merge of the logging.handlers.RotatingFileHandler and the FileTaskHandler.
#
# 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.
"""File logging handler for tasks."""
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Optional
import requests
from airflow.configuration import AirflowConfigException, conf
from airflow.utils.helpers import parse_template_string
if TYPE_CHECKING:
from airflow.models import TaskInstance
class RotatingFileTaskHandler(logging.Handler):
"""
FileTaskHandler is a python log handler that handles and reads
task instance logs. It creates and delegates log handling
to `logging.FileHandler` after receiving task instance context.
It reads logs from task instance's host machine.
:param base_log_folder: Base log folder to place logs.
:param filename_template: template filename string
"""
def __init__(self, base_log_folder: str, filename_template: str, maxBytes=0, backupCount=0):
self.max_bytes = maxBytes
self.backup_count = backupCount
super().__init__()
self.handler = None # type: Optional[logging.FileHandler]
self.local_base = base_log_folder
self.filename_template, self.filename_jinja_template = parse_template_string(filename_template)
def set_context(self, ti: "TaskInstance"):
"""
Provide task_instance context to airflow task handler.
:param ti: task instance object
"""
local_loc = self._init_file(ti)
self.handler = logging.handlers.RotatingFileHandler(
filename=local_loc,
mode='a',
maxBytes=self.max_bytes,
backupCount=self.backup_count,
encoding='utf-8',
delay=False,
)
if self.formatter:
self.handler.setFormatter(self.formatter)
self.handler.setLevel(self.level)
def emit(self, record):
if self.handler:
self.handler.emit(record)
def flush(self):
if self.handler:
self.handler.flush()
def close(self):
if self.handler:
self.handler.close()
def _render_filename(self, ti, try_number):
if self.filename_jinja_template:
if hasattr(ti, 'task'):
jinja_context = ti.get_template_context()
jinja_context['try_number'] = try_number
else:
jinja_context = {
'ti': ti,
'ts': ti.execution_date.isoformat(),
'try_number': try_number,
}
return self.filename_jinja_template.render(**jinja_context)
return self.filename_template.format(
dag_id=ti.dag_id,
task_id=ti.task_id,
execution_date=ti.execution_date.isoformat(),
try_number=try_number,
)
def _read_grouped_logs(self):
return False
def _read(self, ti, try_number, metadata=None): # pylint: disable=unused-argument
"""
Template method that contains custom logic of reading
logs given the try_number.
:param ti: task instance record
:param try_number: current try_number to read log from
:param metadata: log metadata,
can be used for steaming log reading and auto-tailing.
:return: log message as a string and metadata.
"""
# Task instance here might be different from task instance when
# initializing the handler. Thus explicitly getting log location
# is needed to get correct log path.
log_relative_path = self._render_filename(ti, try_number)
location = os.path.join(self.local_base, log_relative_path)
log = ""
if os.path.exists(location):
try:
with open(location) as file:
log += f"*** Reading local file: {location}\n"
log += "".join(file.readlines())
except Exception as e: # pylint: disable=broad-except
log = f"*** Failed to load local log file: {location}\n"
log += "*** {}\n".format(str(e))
elif conf.get('core', 'executor') == 'KubernetesExecutor': # pylint: disable=too-many-nested-blocks
try:
from airflow.kubernetes.kube_client import get_kube_client
kube_client = get_kube_client()
if len(ti.hostname) >= 63:
# Kubernetes takes the pod name and truncates it for the hostname. This truncated hostname
# is returned for the fqdn to comply with the 63 character limit imposed by DNS standards
# on any label of a FQDN.
pod_list = kube_client.list_namespaced_pod(conf.get('kubernetes', 'namespace'))
matches = [
pod.metadata.name
for pod in pod_list.items
if pod.metadata.name.startswith(ti.hostname)
]
if len(matches) == 1:
if len(matches[0]) > len(ti.hostname):
ti.hostname = matches[0]
log += '*** Trying to get logs (last 100 lines) from worker pod {} ***\n\n'.format(
ti.hostname
)
res = kube_client.read_namespaced_pod_log(
name=ti.hostname,
namespace=conf.get('kubernetes', 'namespace'),
container='base',
follow=False,
tail_lines=100,
_preload_content=False,
)
for line in res:
log += line.decode()
except Exception as f: # pylint: disable=broad-except
log += '*** Unable to fetch logs from worker pod {} ***\n{}\n\n'.format(ti.hostname, str(f))
else:
url = os.path.join("http://{ti.hostname}:{worker_log_server_port}/log", log_relative_path).format(
ti=ti, worker_log_server_port=conf.get('celery', 'WORKER_LOG_SERVER_PORT')
)
log += f"*** Log file does not exist: {location}\n"
log += f"*** Fetching from: {url}\n"
try:
timeout = None # No timeout
try:
timeout = conf.getint('webserver', 'log_fetch_timeout_sec')
except (AirflowConfigException, ValueError):
pass
response = requests.get(url, timeout=timeout)
response.encoding = "utf-8"
# Check if the resource was properly fetched
response.raise_for_status()
log += '\n' + response.text
except Exception as e: # pylint: disable=broad-except
log += "*** Failed to fetch log file from worker. {}\n".format(str(e))
return log, {'end_of_log': True}
def read(self, task_instance, try_number=None, metadata=None):
"""
Read logs of given task instance from local machine.
:param task_instance: task instance object
:param try_number: task instance try_number to read logs from. If None
it returns all logs separated by try_number
:param metadata: log metadata,
can be used for steaming log reading and auto-tailing.
:return: a list of listed tuples which order log string by host
"""
# Task instance increments its try number when it starts to run.
# So the log for a particular task try will only show up when
# try number gets incremented in DB, i.e logs produced the time
# after cli run and before try_number + 1 in DB will not be displayed.
if try_number is None:
next_try = task_instance.next_try_number
try_numbers = list(range(1, next_try))
elif try_number < 1:
logs = [
[('default_host', f'Error fetching the logs. Try number {try_number} is invalid.')],
]
return logs, [{'end_of_log': True}]
else:
try_numbers = [try_number]
logs = [''] * len(try_numbers)
metadata_array = [{}] * len(try_numbers)
for i, try_number_element in enumerate(try_numbers):
log, metadata = self._read(task_instance, try_number_element, metadata)
# es_task_handler return logs grouped by host. wrap other handler returning log string
# with default/ empty host so that UI can render the response in the same way
logs[i] = log if self._read_grouped_logs() else [(task_instance.hostname, log)]
metadata_array[i] = metadata
return logs, metadata_array
def _init_file(self, ti):
"""
Create log directory and give it correct permissions.
:param ti: task instance object
:return: relative log path of the given task instance
"""
# To handle log writing when tasks are impersonated, the log files need to
# be writable by the user that runs the Airflow command and the user
# that is impersonated. This is mainly to handle corner cases with the
# SubDagOperator. When the SubDagOperator is run, all of the operators
# run under the impersonated user and create appropriate log files
# as the impersonated user. However, if the user manually runs tasks
# of the SubDagOperator through the UI, then the log files are created
# by the user that runs the Airflow command. For example, the Airflow
# run command may be run by the `airflow_sudoable` user, but the Airflow
# tasks may be run by the `airflow` user. If the log files are not
# writable by both users, then it's possible that re-running a task
# via the UI (or vice versa) results in a permission error as the task
# tries to write to a log file created by the other user.
relative_path = self._render_filename(ti, ti.try_number)
full_path = os.path.join(self.local_base, relative_path)
directory = os.path.dirname(full_path)
# Create the log file and give it group writable permissions
# TODO(aoen): Make log dirs and logs globally readable for now since the SubDag
# operator is not compatible with impersonation (e.g. if a Celery executor is used
# for a SubDag operator and the SubDag operator has a different owner than the
# parent DAG)
Path(directory).mkdir(mode=0o777, parents=True, exist_ok=True)
if not os.path.exists(full_path):
open(full_path, "a").close()
# TODO: Investigate using 444 instead of 666.
os.chmod(full_path, 0o666)
return full_path
Maybe a final note; the links in the airflow UI to the logging will now only open the latest logfile, not the older rotated files which are only accessible by means of SSH or any other interface to access the airflow logging path.
I don't think that there is a rotation mechanism but you can store them in S3 or google cloud storage as describe here : https://airflow.incubator.apache.org/configuration.html#logs

Resources