data_dag.operator_factory

class data_dag.operator_factory.OperatorFactory(*, task_id: Optional[str] = None)

An interface for writing operator factories.

_make_operators(*args, **kwargs) → None: Can be implemented instead of make_operator() to define an operator collection inside a default airflow.utils.task_group.TaskGroup

property default_task_id: str: If overridden, defines a default task ID when none is manually specified

get_task_id(): Provides the custom task ID if provided, else this factory’s default task ID

make_operator(*args, **kwargs) → Optional[Union[airflow.models.taskmixin.TaskMixin, Sequence[airflow.models.taskmixin.TaskMixin]]]

Converts this factory metadata into an operator.

If you need to create multiple operators, whether connected or not, implement _make_operators() instead, and they will be automatically wrapped in a airflow.utils.task_group.TaskGroup.

Returns:: Zero or more operator-like things. The code that calls this should know how to handle the possible return types for this particular factory.

class data_dag.operator_factory.OperatorComponent: A non-operator component for use in other operator factories. Just a proxy for pydantic.BaseModel.

class data_dag.operator_factory.SimpleOperatorFactory(*, task_id: Optional[str] = None)

Identical to OperatorFactory except that this represents predominantly a single field of metadata.

The model that inherits from SimpleOperatorFactory can only have a single non-required field (meaning no default value and not Optional). In return the constructor for this object, in addition to being callable with a dictionary of field values, can also be called with a simple literal to fill in the single required field.

Consider the following example:

class FilePath(SimpleOperatorFactory):
    path: str  # <-- single required field
    is_file: bool = True  # <-- optional field (because of default)
    mime_type: Optional[str]  # <-- optional field (because of Optional type)

    def make_operator(self):
        ...

Normally, this object could only be instantiated using a dictionary:

FilePath.parse_obj({'path': 'path/to/file.txt'})

# Or, in YAML:
# outer_object:
#   my_file:
#     path: 'path/to/file.txt'
# Or, in JSON:
# {"outer_object": {"my_file": {"path": "path/to/file.txt"}}}

However, because we inherit from SimpleOperatorFactory, we can instantiate a FilePath by specifying just the path literal:

FilePath.parse_obj('path/to/file.txt')

# Or, in YAML:
# outer_object:
#   my_file: 'path/to/file.txt'
# Or, in JSON:
# {"outer_object": {"my_file": "path/to/file.txt"}}

class data_dag.operator_factory.SimpleOperatorComponent: An extension of OperatorComponent to have the same single-field functionality as SimpleOperatorFactory.

class data_dag.operator_factory.DynamicOperatorFactory(*, task_id: Optional[str] = None)

An OperatorFactory that can automatically instantiate sub-classes based on the input data.

Consider the following example:

class InputFile(DynamicOperatorFactory, abc.ABC):
    pass

class LocalFile(InputFile):
    __type_name__ = 'local'

    path: str

class S3File(InputFile):
    __type_name__ = 's3'

    bucket: str
    key: str

InputFile.parse_obj({'type': 's3', 'bucket': 'my-bucket', 'key': 'my-key'})
# S3File(bucket='my-bucket', key='my-key')

Note how the type of object that gets instantiated is dynamically chosen from the data, rather than specified by the code. This allows a supertype to be used in code, and for the subtype to be chosen at runtime based on data.

To use a dynamic factory, define your base supertype to inherit directly from DynamicOperatorFactory and abc.ABC. The class can be totally empty, as in the example above. This top-level class will be populated with a dictionary that will automatically track subclasses as they get define.

Warning

It’s important to remember that, while subtypes are automatically tracked upon definition, they must still be imported somewhere. Make sure that when the supertype is imported, the subtypes also eventually get imported, or else they will be unavailable at DAG resolution time.

Subclasses must either define __type_name__ = "some_name" or else inherit from abc.ABC to indicate that they are abstract. Classes that are not abstract and not named will generate a warning.

A default subtype can be specified using __default_type_name__ in the top-level type. Note that this is the __type_name__ of the default subclass, not the class name itself.

By default, the subclass is chosen by the "type" key in the input data. This can be changed by setting __type_kwarg_name__ in the top-level type to some other string. This key will be stripped from the input data and all other keys will be passed along to the subtype’s constructor without further modification.

Attempting to construct a top-level object, either directly (with its constructor) or using parse_obj, without specifying a “type” (or whatever you renamed the key to be) will result in a TypeError.

Note

Pydantic already supports Union types, so why would we use a custom DynamicOperatorFactory instead?

Dynamic factories provide two key advantages:

The subtype selected is explicit rather than implicit. The subtypes don’t need to be distinguishable in any other way besides their __type_name__, nor is there any kind of ordering of the subtypes.
The list of options is automatically maintained, as long as the modules containing the subtypes are sure to be imported. That is, another component or factory can use the top-level type to annotate one of its fields, and the subtypes will automatically be implied.