data.datasets.Dataset.__post_init__() - Code Metrics - Inspection of "Merge fixing dependencies ending with parallel tas..." - openego/eGon-data - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Push — dev ( 224a40...42ae28 )

by Stephan

created 2021-10-05 11:42 UTC

data.datasets.Dataset.__post_init__() C

↳ Parent: data.datasets

Complexity

Conditions

Size

Total Lines	42
Code Lines	31

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	9
eloc	31
nop	1
dl	0
loc	42
rs	6.6666
c	0
b	0
f	0

"""The API for configuring datasets."""

from __future__ import annotations

from collections import abc
from dataclasses import dataclass
from functools import reduce
from typing import Callable, Iterable, Set, Tuple, Union
import re

from airflow import DAG
from airflow.operators import BaseOperator as Operator
from airflow.operators.python_operator import PythonOperator
from sqlalchemy import Column, ForeignKey, Integer, String, Table, orm, tuple_
from sqlalchemy.ext.declarative import declarative_base

from egon.data import db, logger

Base = declarative_base()
SCHEMA = "metadata"


def setup():
    """Create the database structure for storing dataset information."""
    # TODO: Move this into a task generating the initial database structure.
    db.execute_sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA};")
    Model.__table__.create(bind=db.engine(), checkfirst=True)
    DependencyGraph.create(bind=db.engine(), checkfirst=True)


# TODO: Figure out how to use a mapped class as an association table.
#
# Trying it out, I ran into quite a few problems and didn't have time to do
# further research. The benefits are mostly just convenience, so it doesn't
# have a high priority. But I'd like to keep the code I started out with to
# have a starting point for me or anybody else trying again.
#
# class DependencyGraph(Base):
#     __tablename__ = "dependency_graph"
#     __table_args__ = {"schema": SCHEMA}
#     dependency_id = Column(Integer, ForeignKey(Model.id), primary_key=True,)
#     dependent_id = Column(Integer, ForeignKey(Model.id), primary_key=True,)

DependencyGraph = Table(
    "dependency_graph",
    Base.metadata,
    Column(
        "dependency_id",
        Integer,
        ForeignKey(f"{SCHEMA}.datasets.id"),
        primary_key=True,
    ),
    Column(
        "dependent_id",
        Integer,
        ForeignKey(f"{SCHEMA}.datasets.id"),
        primary_key=True,
    ),
    schema=SCHEMA,
)


class Model(Base):
    __tablename__ = "datasets"
    __table_args__ = {"schema": SCHEMA}
    id = Column(Integer, primary_key=True)
    name = Column(String, unique=True, nullable=False)
    version = Column(String, nullable=False)
    epoch = Column(Integer, default=0)
    dependencies = orm.relationship(
        "Model",
        secondary=DependencyGraph,
        primaryjoin=id == DependencyGraph.c.dependent_id,
        secondaryjoin=id == DependencyGraph.c.dependency_id,
        backref=orm.backref("dependents", cascade="all, delete"),
    )


#: A :class:`Task` is an Airflow :class:`Operator` or any
#: :class:`Callable <typing.Callable>` taking no arguments and returning
#: :obj:`None`. :class:`Callables <typing.Callable>` will be converted
#: to :class:`Operators <Operator>` by wrapping them in a
#: :class:`PythonOperator` and setting the :obj:`~PythonOperator.task_id`
#: to the :class:`Callable <typing.Callable>`'s
#: :obj:`~definition.__name__`, with underscores replaced with hyphens.
#: If the :class:`Callable <typing.Callable>`'s `__module__`__ attribute
#: contains the string :obj:`"egon.data.datasets."`, the
#: :obj:`~PythonOperator.task_id` is also prefixed with the module name,
#: followed by a dot and with :obj:`"egon.data.datasets."` removed.
#:
#: __ https://docs.python.org/3/reference/datamodel.html#index-34
Task = Union[Callable[[], None], Operator]
#: A graph of tasks is, in its simplest form, just a single node, i.e. a
#: single :class:`Task`. More complex graphs can be specified by nesting
#: :class:`sets <builtins.set>` and :class:`tuples <builtins.tuple>` of
#: :class:`TaskGraphs <TaskGraph>`. A set of :class:`TaskGraphs
#: <TaskGraph>` means that they are unordered and can be
#: executed in parallel. A :class:`tuple` specifies an implicit ordering so
#: a :class:`tuple` of :class:`TaskGraphs <TaskGraph>` will be executed
#: sequentially in the given order.
TaskGraph = Union[Task, Set["TaskGraph"], Tuple["TaskGraph", ...]]


def prefix(o):
    module = o.__module__
    parent = f"{__name__}."
    return f"{module.replace(parent, '')}." if parent in module else ""


@dataclass
class Tasks(dict):
    first: Set[Task]
    last: Set[Task]
    graph: TaskGraph = ()

    def __init__(self, graph: TaskGraph):
        """Connect multiple tasks into a potentially complex graph.

        Parses a :class:`TaskGraph` into a :class:`Tasks` object.
        """
        if isinstance(graph, Callable):
            graph = PythonOperator(
                task_id=f"{prefix(graph)}{graph.__name__.replace('_', '-')}",
                python_callable=graph,
            )
        self.graph = graph
        if isinstance(graph, Operator):
            self.first = {graph}
            self.last = {graph}
            self[graph.task_id] = graph
        elif isinstance(graph, abc.Sized) and len(graph) == 0:
            self.first = {}
            self.last = {}
        elif isinstance(graph, abc.Set):
            results = [Tasks(subtasks) for subtasks in graph]
            self.first = {task for result in results for task in result.first}
            self.last = {task for result in results for task in result.last}
            self.update(reduce(lambda d1, d2: dict(d1, **d2), results, {}))
            self.graph = set(tasks.graph for tasks in results)
        elif isinstance(graph, tuple):
            results = [Tasks(subtasks) for subtasks in graph]
            for (left, right) in zip(results[:-1], results[1:]):
                for last in left.last:
                    for first in right.first:
                        last.set_downstream(first)
            self.first = results[0].first
            self.last = results[-1].last
            self.update(reduce(lambda d1, d2: dict(d1, **d2), results, {}))
            self.graph = tuple(tasks.graph for tasks in results)
        else:
            raise (
                TypeError(
                    "`egon.data.datasets.Tasks` got an argument of type:\n\n"
                    f"  {type(graph)}\n\n"
                    "where only `Task`s, `Set`s and `Tuple`s are allowed."
                )
            )


@dataclass
class Dataset:
    #: The name of the Dataset
    name: str
    #: The :class:`Dataset`'s version. Can be anything from a simple
    #: semantic versioning string like "2.1.3", to a more complex
    #: string, like for example "2021-01-01.schleswig-holstein.0" for
    #: OpenStreetMap data.
    #: Note that the latter encodes the :class:`Dataset`'s date, region
    #: and a sequential number in case the data changes without the date
    #: or region changing, for example due to implementation changes.
    version: str
    #: The first task(s) of this :class:`Dataset` will be marked as
    #: downstream of any of the listed dependencies. In case of bare
    #: :class:`Task`, a direct link will be created whereas for a
    #: :class:`Dataset` the link will be made to all of its last tasks.
    dependencies: Iterable[Union[Dataset, Task]] = ()
    #: The tasks of this :class:`Dataset`. A :class:`TaskGraph` will
    #: automatically be converted to :class:`Tasks`.
    tasks: Union[Tasks, TaskGraph] = ()

    def check_version(self, after_execution=()):
        def skip_task(task, *xs, **ks):
            with db.session_scope() as session:
                datasets = session.query(Model).filter_by(name=self.name).all()
                if self.version in [
                    ds.version for ds in datasets
                ] and not re.search(r"\.dev$", self.version):
                    logger.info(
                        f"Dataset '{self.name}' version '{self.version}'"
                        f" already executed. Skipping."
                    )
                else:
                    for ds in datasets:
                        session.delete(ds)
                    result = super(type(task), task).execute(*xs, **ks)
                    for function in after_execution:
                        function(session)
                    return result

        return skip_task

    def update(self, session):
        dataset = Model(name=self.name, version=self.version)
        dependencies = (
            session.query(Model)
            .filter(
                tuple_(Model.name, Model.version).in_(
                    [
                        (dataset.name, dataset.version)
                        for dependency in self.dependencies
                        if isinstance(dependency, Dataset)
                        or hasattr(dependency, "dataset")
                        for dataset in [
                            dependency.dataset
                            if isinstance(dependency, Operator)
                            else dependency
                        ]
                    ]
                )
            )
            .all()
        )
        dataset.dependencies = dependencies
        session.add(dataset)

    def __post_init__(self):
        self.dependencies = list(self.dependencies)
        if not isinstance(self.tasks, Tasks):
            self.tasks = Tasks(self.tasks)
        if len(self.tasks.last) > 1:
            # Explicitly create single final task, because we can't know
            # which of the multiple tasks finishes last.
            name = prefix(self)
            name = name if name else f"{self.name}."
            update_version = PythonOperator(
                task_id=f"{name}update-version",
                # Do nothing, because updating will be added later.
                python_callable=lambda *xs, **ks: None,
            )
            self.tasks = Tasks((self.tasks.graph, update_version))
        # Due to the `if`-block above, there'll now always be exactly
        # one task in `self.tasks.last` which the next line just
        # selects.
        last = list(self.tasks.last)[0]
        for task in self.tasks.values():
            task.dataset = self
            cls = task.__class__
            versioned = type(
                f"{self.name[0].upper()}{self.name[1:]} (versioned)",
                (cls,),
                {
                    "execute": self.check_version(
                        after_execution=[self.update] if task is last else []
                    )
                },
            )
            task.__class__ = versioned

        predecessors = [
            task
            for dataset in self.dependencies
            if isinstance(dataset, Dataset)
            for task in dataset.tasks.last
        ] + [task for task in self.dependencies if isinstance(task, Operator)]
        for p in predecessors:
            for first in self.tasks.first:
                p.set_downstream(first)

    def insert_into(self, dag: DAG):
        for task in self.tasks.values():
            for attribute in dag.default_args:
                if getattr(task, attribute) is None:
                    setattr(task, attribute, dag.default_args[attribute])
        dag.add_tasks(self.tasks.values())


1			"""The API for configuring datasets."""
2
3			from __future__ import annotations
4
5			from collections import abc
6			from dataclasses import dataclass
7			from functools import reduce
8			from typing import Callable, Iterable, Set, Tuple, Union
9			import re
10
11			from airflow import DAG
12			from airflow.operators import BaseOperator as Operator
13			from airflow.operators.python_operator import PythonOperator
14			from sqlalchemy import Column, ForeignKey, Integer, String, Table, orm, tuple_
15			from sqlalchemy.ext.declarative import declarative_base
16
17			from egon.data import db, logger
18
19			Base = declarative_base()
20			SCHEMA = "metadata"
21
22
23			def setup():
24			"""Create the database structure for storing dataset information."""
25			# TODO: Move this into a task generating the initial database structure.
26			db.execute_sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA};")
27			Model.__table__.create(bind=db.engine(), checkfirst=True)
28			DependencyGraph.create(bind=db.engine(), checkfirst=True)
29
30
31			# TODO: Figure out how to use a mapped class as an association table.
32			#
33			# Trying it out, I ran into quite a few problems and didn't have time to do
34			# further research. The benefits are mostly just convenience, so it doesn't
35			# have a high priority. But I'd like to keep the code I started out with to
36			# have a starting point for me or anybody else trying again.
37			#
38			# class DependencyGraph(Base):
39			# __tablename__ = "dependency_graph"
40			# __table_args__ = {"schema": SCHEMA}
41			# dependency_id = Column(Integer, ForeignKey(Model.id), primary_key=True,)
42			# dependent_id = Column(Integer, ForeignKey(Model.id), primary_key=True,)
43
44			DependencyGraph = Table(
45			"dependency_graph",
46			Base.metadata,
47			Column(
48			"dependency_id",
49			Integer,
50			ForeignKey(f"{SCHEMA}.datasets.id"),
51			primary_key=True,
52			),
53			Column(
54			"dependent_id",
55			Integer,
56			ForeignKey(f"{SCHEMA}.datasets.id"),
57			primary_key=True,
58			),
59			schema=SCHEMA,
60			)
61
62
63			class Model(Base):
64			__tablename__ = "datasets"
65			__table_args__ = {"schema": SCHEMA}
66			id = Column(Integer, primary_key=True)
67			name = Column(String, unique=True, nullable=False)
68			version = Column(String, nullable=False)
69			epoch = Column(Integer, default=0)
70			dependencies = orm.relationship(
71			"Model",
72			secondary=DependencyGraph,
73			primaryjoin=id == DependencyGraph.c.dependent_id,
74			secondaryjoin=id == DependencyGraph.c.dependency_id,
75			backref=orm.backref("dependents", cascade="all, delete"),
76			)
77
78
79			#: A :class:`Task` is an Airflow :class:`Operator` or any
80			#: :class:`Callable <typing.Callable>` taking no arguments and returning
81			#: :obj:`None`. :class:`Callables <typing.Callable>` will be converted
82			#: to :class:`Operators <Operator>` by wrapping them in a
83			#: :class:`PythonOperator` and setting the :obj:`~PythonOperator.task_id`
84			#: to the :class:`Callable <typing.Callable>`'s
85			#: :obj:`~definition.__name__`, with underscores replaced with hyphens.
86			#: If the :class:`Callable <typing.Callable>`'s `__module__`__ attribute
87			#: contains the string :obj:`"egon.data.datasets."`, the
88			#: :obj:`~PythonOperator.task_id` is also prefixed with the module name,
89			#: followed by a dot and with :obj:`"egon.data.datasets."` removed.
90			#:
91			#: __ https://docs.python.org/3/reference/datamodel.html#index-34
92			Task = Union[Callable[[], None], Operator]
93			#: A graph of tasks is, in its simplest form, just a single node, i.e. a
94			#: single :class:`Task`. More complex graphs can be specified by nesting
95			#: :class:`sets <builtins.set>` and :class:`tuples <builtins.tuple>` of
96			#: :class:`TaskGraphs <TaskGraph>`. A set of :class:`TaskGraphs
97			#: <TaskGraph>` means that they are unordered and can be
98			#: executed in parallel. A :class:`tuple` specifies an implicit ordering so
99			#: a :class:`tuple` of :class:`TaskGraphs <TaskGraph>` will be executed
100			#: sequentially in the given order.
101			TaskGraph = Union[Task, Set["TaskGraph"], Tuple["TaskGraph", ...]]
102
103
104			def prefix(o):
105			module = o.__module__
106			parent = f"{__name__}."
107			return f"{module.replace(parent, '')}." if parent in module else ""
108
109
110			@dataclass
111			class Tasks(dict):
112			first: Set[Task]
113			last: Set[Task]
114			graph: TaskGraph = ()
115
116			def __init__(self, graph: TaskGraph):
117			"""Connect multiple tasks into a potentially complex graph.
118
119			Parses a :class:`TaskGraph` into a :class:`Tasks` object.
120			"""
121			if isinstance(graph, Callable):
122			graph = PythonOperator(
123			task_id=f"{prefix(graph)}{graph.__name__.replace('_', '-')}",
124			python_callable=graph,
125			)
126			self.graph = graph
127			if isinstance(graph, Operator):
128			self.first = {graph}
129			self.last = {graph}
130			self[graph.task_id] = graph
131			elif isinstance(graph, abc.Sized) and len(graph) == 0:
132			self.first = {}
133			self.last = {}
134			elif isinstance(graph, abc.Set):
135			results = [Tasks(subtasks) for subtasks in graph]
136			self.first = {task for result in results for task in result.first}
137			self.last = {task for result in results for task in result.last}
138			self.update(reduce(lambda d1, d2: dict(d1, **d2), results, {}))
139			self.graph = set(tasks.graph for tasks in results)
140			elif isinstance(graph, tuple):
141			results = [Tasks(subtasks) for subtasks in graph]
142			for (left, right) in zip(results[:-1], results[1:]):
143			for last in left.last:
144			for first in right.first:
145			last.set_downstream(first)
146			self.first = results[0].first
147			self.last = results[-1].last
148			self.update(reduce(lambda d1, d2: dict(d1, **d2), results, {}))
149			self.graph = tuple(tasks.graph for tasks in results)
150			else:
151			raise (
152			TypeError(
153			"`egon.data.datasets.Tasks` got an argument of type:\n\n"
154			f" {type(graph)}\n\n"
155			"where only `Task`s, `Set`s and `Tuple`s are allowed."
156			)
157			)
158
159
160			@dataclass
161			class Dataset:
162			#: The name of the Dataset
163			name: str
164			#: The :class:`Dataset`'s version. Can be anything from a simple
165			#: semantic versioning string like "2.1.3", to a more complex
166			#: string, like for example "2021-01-01.schleswig-holstein.0" for
167			#: OpenStreetMap data.
168			#: Note that the latter encodes the :class:`Dataset`'s date, region
169			#: and a sequential number in case the data changes without the date
170			#: or region changing, for example due to implementation changes.
171			version: str
172			#: The first task(s) of this :class:`Dataset` will be marked as
173			#: downstream of any of the listed dependencies. In case of bare
174			#: :class:`Task`, a direct link will be created whereas for a
175			#: :class:`Dataset` the link will be made to all of its last tasks.
176			dependencies: Iterable[Union[Dataset, Task]] = ()
177			#: The tasks of this :class:`Dataset`. A :class:`TaskGraph` will
178			#: automatically be converted to :class:`Tasks`.
179			tasks: Union[Tasks, TaskGraph] = ()
180
181			def check_version(self, after_execution=()):
182			def skip_task(task, xs, *ks):
183			with db.session_scope() as session:
184			datasets = session.query(Model).filter_by(name=self.name).all()
185			if self.version in [
186			ds.version for ds in datasets
187			] and not re.search(r"\.dev$", self.version):
188			logger.info(
189			f"Dataset '{self.name}' version '{self.version}'"
190			f" already executed. Skipping."
191			)
192			else:
193			for ds in datasets:
194			session.delete(ds)
195			result = super(type(task), task).execute(xs, *ks)
196			for function in after_execution:
197			function(session)
198			return result
199
200			return skip_task
201
202			def update(self, session):
203			dataset = Model(name=self.name, version=self.version)
204			dependencies = (
205			session.query(Model)
206			.filter(
207			tuple_(Model.name, Model.version).in_(
208			[
209			(dataset.name, dataset.version)
210			for dependency in self.dependencies
211			if isinstance(dependency, Dataset)
212			or hasattr(dependency, "dataset")
213			for dataset in [
214			dependency.dataset
215			if isinstance(dependency, Operator)
216			else dependency
217			]
218			]
219			)
220			)
221			.all()
222			)
223			dataset.dependencies = dependencies
224			session.add(dataset)
225
226			def __post_init__(self):
227			self.dependencies = list(self.dependencies)
228			if not isinstance(self.tasks, Tasks):
229			self.tasks = Tasks(self.tasks)
230			if len(self.tasks.last) > 1:
231			# Explicitly create single final task, because we can't know
232			# which of the multiple tasks finishes last.
233			name = prefix(self)
234			name = name if name else f"{self.name}."
235			update_version = PythonOperator(
236			task_id=f"{name}update-version",
237			# Do nothing, because updating will be added later.
238			python_callable=lambda xs, *ks: None,
239			)
240			self.tasks = Tasks((self.tasks.graph, update_version))
241			# Due to the `if`-block above, there'll now always be exactly
242			# one task in `self.tasks.last` which the next line just
243			# selects.
244			last = list(self.tasks.last)[0]
245			for task in self.tasks.values():
246			task.dataset = self
247			cls = task.__class__
248			versioned = type(
249			f"{self.name[0].upper()}{self.name[1:]} (versioned)",
250			(cls,),
251			{
252			"execute": self.check_version(
253			after_execution=[self.update] if task is last else []
254			)
255			},
256			)
257			task.__class__ = versioned
258
259			predecessors = [
260			task
261			for dataset in self.dependencies
262			if isinstance(dataset, Dataset)
263			for task in dataset.tasks.last
264			] + [task for task in self.dependencies if isinstance(task, Operator)]
265			for p in predecessors:
266			for first in self.tasks.first:
267			p.set_downstream(first)
268
269			def insert_into(self, dag: DAG):
270			for task in self.tasks.values():
271			for attribute in dag.default_args:
272			if getattr(task, attribute) is None:
273			setattr(task, attribute, dag.default_args[attribute])
274			dag.add_tasks(self.tasks.values())
275

openego / eGon-data

Push — dev ( 224a40...42ae28 )

data.datasets.Dataset.__post_init__() C

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like