Flow exporting/importing (#187)

* stages/*: Add SerializerModel as base model, implement serializer property

* flows: add initial flow exporter and importer

* policies/*: implement .serializer for all policies

* root: fix missing dacite requirement

passbook/flows/transfer/common.py

@@ -0,0 +1,59 @@
+"""transfer common classes"""
+from dataclasses import asdict, dataclass, field, is_dataclass
+from json.encoder import JSONEncoder
+from typing import Any, Dict, List
+from uuid import UUID
+
+from passbook.lib.models import SerializerModel
+from passbook.lib.sentry import SentryIgnoredException
+
+
+def get_attrs(obj: SerializerModel) -> Dict[str, Any]:
+    """Get object's attributes via their serializer, and covert it to a normal dict"""
+    data = dict(obj.serializer(obj).data)
+    if "policies" in data:
+        data.pop("policies")
+    if "stages" in data:
+        data.pop("stages")
+    return data
+
+
+@dataclass
+class FlowBundleEntry:
+    """Single entry of a bundle"""
+
+    identifier: str
+    model: str
+    attrs: Dict[str, Any]
+
+    @staticmethod
+    def from_model(model: SerializerModel) -> "FlowBundleEntry":
+        """Convert a SerializerModel instance to a Bundle Entry"""
+        return FlowBundleEntry(
+            identifier=model.pk,
+            model=f"{model._meta.app_label}.{model._meta.model_name}",
+            attrs=get_attrs(model),
+        )
+
+
+@dataclass
+class FlowBundle:
+    """Dataclass used for a full export"""
+
+    version: int = field(default=1)
+    entries: List[FlowBundleEntry] = field(default_factory=list)
+
+
+class DataclassEncoder(JSONEncoder):
+    """Convert FlowBundleEntry to json"""
+
+    def default(self, o):
+        if is_dataclass(o):
+            return asdict(o)
+        if isinstance(o, UUID):
+            return str(o)
+        return super().default(o)
+
+
+class EntryInvalidError(SentryIgnoredException):
+    """Error raised when an entry is invalid"""

passbook/flows/transfer/exporter.py

@@ -0,0 +1,78 @@
+"""Flow exporter"""
+from json import dumps
+from typing import Iterator
+
+from passbook.flows.models import Flow, FlowStageBinding, Stage
+from passbook.flows.transfer.common import DataclassEncoder, FlowBundle, FlowBundleEntry
+from passbook.policies.models import Policy, PolicyBinding, PolicyBindingModel
+from passbook.stages.prompt.models import PromptStage
+
+
+class FlowExporter:
+    """Export flow with attached stages into json"""
+
+    flow: Flow
+    with_policies: bool
+    with_stage_prompts: bool
+
+    def __init__(self, flow: Flow):
+        self.flow = flow
+        self.with_policies = True
+        self.with_stage_prompts = True
+
+    def walk_stages(self) -> Iterator[FlowBundleEntry]:
+        """Convert all stages attached to self.flow into FlowBundleEntry objects"""
+        stages = (
+            Stage.objects.filter(flow=self.flow).select_related().select_subclasses()
+        )
+        for stage in stages:
+            if isinstance(stage, PromptStage):
+                pass
+            yield FlowBundleEntry.from_model(stage)
+
+    def walk_stage_bindings(self) -> Iterator[FlowBundleEntry]:
+        """Convert all bindings attached to self.flow into FlowBundleEntry objects"""
+        bindings = FlowStageBinding.objects.filter(target=self.flow).select_related()
+        for binding in bindings:
+            yield FlowBundleEntry.from_model(binding)
+
+    def walk_policies(self) -> Iterator[FlowBundleEntry]:
+        """Walk over all policies and their respective bindings"""
+        pbm_uuids = [self.flow.pbm_uuid]
+        for stage_subclass in Stage.__subclasses__():
+            if issubclass(stage_subclass, PolicyBindingModel):
+                pbm_uuids += stage_subclass.objects.filter(flow=self.flow).values_list(
+                    "pbm_uuid", flat=True
+                )
+        pbm_uuids += FlowStageBinding.objects.filter(target=self.flow).values_list(
+            "pbm_uuid", flat=True
+        )
+        policies = Policy.objects.filter(bindings__in=pbm_uuids).select_related()
+        for policy in policies:
+            yield FlowBundleEntry.from_model(policy)
+        bindings = PolicyBinding.objects.filter(target__in=pbm_uuids).select_related()
+        for binding in bindings:
+            yield FlowBundleEntry.from_model(binding)
+
+    def walk_stage_prompts(self) -> Iterator[FlowBundleEntry]:
+        """Walk over all prompts associated with any PromptStages"""
+        prompt_stages = PromptStage.objects.filter(flow=self.flow)
+        for stage in prompt_stages:
+            for prompt in stage.fields.all():
+                yield FlowBundleEntry.from_model(prompt)
+
+    def export(self) -> FlowBundle:
+        """Create a list of all objects including the flow"""
+        bundle = FlowBundle()
+        bundle.entries.append(FlowBundleEntry.from_model(self.flow))
+        if self.with_stage_prompts:
+            bundle.entries.extend(self.walk_stage_prompts())
+        bundle.entries.extend(self.walk_stages())
+        bundle.entries.extend(self.walk_stage_bindings())
+        if self.with_policies:
+            bundle.entries.extend(self.walk_policies())
+        return bundle
+
+    def export_to_string(self) -> str:
+        """Call export and convert it to json"""
+        return dumps(self.export(), cls=DataclassEncoder)

passbook/flows/transfer/importer.py

@@ -0,0 +1,134 @@
+"""Flow importer"""
+from json import loads
+from typing import Type
+
+from dacite import from_dict
+from dacite.exceptions import DaciteError
+from django.apps import apps
+from django.db import transaction
+from django.db.models import Model
+from rest_framework.serializers import BaseSerializer, Serializer
+from structlog import BoundLogger, get_logger
+
+from passbook.flows.models import Flow, FlowStageBinding, Stage
+from passbook.flows.transfer.common import (
+    EntryInvalidError,
+    FlowBundle,
+    FlowBundleEntry,
+)
+from passbook.lib.models import SerializerModel
+from passbook.policies.models import Policy, PolicyBinding, PolicyBindingModel
+from passbook.stages.prompt.models import Prompt
+
+ALLOWED_MODELS = (Flow, FlowStageBinding, Stage, Policy, PolicyBinding, Prompt)
+
+
+class FlowImporter:
+    """Import Flow from json"""
+
+    __import: FlowBundle
+
+    logger: BoundLogger
+
+    def __init__(self, json_input: str):
+        self.logger = get_logger()
+        import_dict = loads(json_input)
+        try:
+            self.__import = from_dict(FlowBundle, import_dict)
+        except DaciteError as exc:
+            raise EntryInvalidError from exc
+
+    def validate(self) -> bool:
+        """Validate loaded flow export, ensure all models are allowed
+        and serializers have no errors"""
+        if self.__import.version != 1:
+            self.logger.warning("Invalid bundle version")
+            return False
+        for entry in self.__import.entries:
+            try:
+                self._validate_single(entry)
+            except EntryInvalidError as exc:
+                self.logger.warning(exc)
+                return False
+        return True
+
+    def __get_pk_filed(self, model_class: Type[Model]) -> str:
+        fields = model_class._meta.get_fields()
+        pks = []
+        for field in fields:
+            # Ignore base PK from pbm as that isn't the same pk we exported
+            if field.model in [PolicyBindingModel]:
+                continue
+            # Ignore primary keys with _ptr suffix as those are surrogate and not what we exported
+            if field.name.endswith("_ptr"):
+                continue
+            if hasattr(field, "primary_key"):
+                if field.primary_key:
+                    pks.append(field.name)
+        if len(pks) > 1:
+            self.logger.debug(
+                "Found more than one fields with primary_key=True, using pk", pks=pks
+            )
+            return "pk"
+        return pks[0]
+
+    def _validate_single(self, entry: FlowBundleEntry) -> BaseSerializer:
+        """Validate a single entry"""
+        model_app_label, model_name = entry.model.split(".")
+        model: SerializerModel = apps.get_model(model_app_label, model_name)
+        if not isinstance(model(), ALLOWED_MODELS):
+            raise EntryInvalidError(f"Model {model} not allowed")
+
+        # If we try to validate without referencing a possible instance
+        # we'll get a duplicate error, hence we load the model here and return
+        # the full serializer for later usage
+        existing_models = model.objects.filter(pk=entry.identifier)
+        serializer_kwargs = {"data": entry.attrs}
+        if existing_models.exists():
+            self.logger.debug(
+                "initialise serializer with instance", instance=existing_models.first()
+            )
+            serializer_kwargs["instance"] = existing_models.first()
+        else:
+            self.logger.debug("initialise new instance", pk=entry.identifier)
+
+        serializer: Serializer = model().serializer(**serializer_kwargs)
+        is_valid = serializer.is_valid()
+        if not is_valid:
+            raise EntryInvalidError(f"Serializer errors {serializer.errors}")
+        if not existing_models.exists():
+            # only insert the PK if we're creating a new model, otherwise we get
+            # an integrity error
+            model_pk = self.__get_pk_filed(model)
+            serializer.validated_data[model_pk] = entry.identifier
+        return serializer
+
+    def apply(self) -> bool:
+        """Apply (create/update) flow json, in database transaction"""
+        transaction.set_autocommit(False)
+        successful = self._apply_models()
+        if not successful:
+            self.logger.debug("Reverting changes due to error")
+            transaction.rollback()
+            transaction.set_autocommit(True)
+            return False
+        self.logger.debug("Committing changes")
+        transaction.commit()
+        transaction.set_autocommit(True)
+        return True
+
+    def _apply_models(self) -> bool:
+        """Apply (create/update) flow json"""
+        for entry in self.__import.entries:
+            model_app_label, model_name = entry.model.split(".")
+            model: SerializerModel = apps.get_model(model_app_label, model_name)
+            # Validate each single entry
+            try:
+                serializer = self._validate_single(entry)
+            except EntryInvalidError as exc:
+                self.logger.error("entry not valid", entry=entry, error=exc)
+                return False
+
+            model = serializer.save()
+            self.logger.debug("updated model", model=model, pk=model.pk)
+        return True