Added Post-import Hook (#205)

* add: Post Import record hook, rename related for clarity

- Renamed `SourceAdapter.post_import` => `post_load`
- Renamed `SourceModelWrapper.post_import` => `post_process_references`
- Renamed `PreImport` => `PreImportRecord`

* fix: Linters

* fix: Tests

Author: snaselj

changes/205.added

@@ -0,0 +1 @@
+Added post-import hook.

nautobot_netbox_importer/diffsync/adapters/netbox.py

@@ -92,7 +92,8 @@ def load(self) -> None:
             fix_power_feed_locations(self)
         if self.options.unrack_zero_uheight_devices:
             unrack_zero_uheight_devices(self)
-        self.post_import()
+
+        self.post_load()
 
     def import_to_nautobot(self) -> None:
         """Import a NetBox export file into Nautobot."""

nautobot_netbox_importer/diffsync/models/custom_fields.py

@@ -8,7 +8,7 @@
     EMPTY_VALUES,
     DiffSyncBaseModel,
     ImporterPass,
-    PreImportResult,
+    PreImportRecordResult,
     SourceAdapter,
     SourceField,
     fields,
@@ -95,14 +95,14 @@ def setup(adapter: SourceAdapter) -> None:
     """Map NetBox custom fields to Nautobot."""
     choice_sets = {}
 
-    def create_choice_set(source: RecordData, importer_pass: ImporterPass) -> PreImportResult:
+    def create_choice_set(source: RecordData, importer_pass: ImporterPass) -> PreImportRecordResult:
         if importer_pass == ImporterPass.DEFINE_STRUCTURE:
             choice_sets[source.get("id")] = [
                 *_convert_choices(source.get("base_choices")),
                 *_convert_choices(source.get("extra_choices")),
             ]
 
-        return PreImportResult.USE_RECORD
+        return PreImportRecordResult.USE_RECORD
 
     def define_choice_set(field: SourceField) -> None:
         def choices_importer(source: RecordData, target: DiffSyncBaseModel) -> None:
@@ -144,7 +144,7 @@ def create_choices(choices: list, custom_field_uid: Uid) -> None:
     # Defined in NetBox but not in Nautobot
     adapter.configure_model(
         "extras.CustomFieldChoiceSet",
-        pre_import=create_choice_set,
+        pre_import_record=create_choice_set,
     )
 
     adapter.configure_model(

nautobot_netbox_importer/diffsync/models/dcim.py

@@ -4,7 +4,13 @@
 from uuid import UUID
 
 from nautobot_netbox_importer.base import RecordData
-from nautobot_netbox_importer.generator import DiffSyncBaseModel, PreImportResult, SourceAdapter, SourceField, fields
+from nautobot_netbox_importer.generator import (
+    DiffSyncBaseModel,
+    PreImportRecordResult,
+    SourceAdapter,
+    SourceField,
+    fields,
+)
 
 from .locations import define_location
 
@@ -39,13 +45,13 @@ def units_importer(source: RecordData, target: DiffSyncBaseModel) -> None:
     field.set_importer(units_importer)
 
 
-def _pre_import_cable_termination(source: RecordData, _) -> PreImportResult:
+def _pre_import_cable_termination(source: RecordData, _) -> PreImportRecordResult:
     cable_end = source.pop("cable_end").lower()
     source["id"] = source.pop("cable")
     source[f"termination_{cable_end}_type"] = source.pop("termination_type")
     source[f"termination_{cable_end}_id"] = source.pop("termination_id")
 
-    return PreImportResult.USE_RECORD
+    return PreImportRecordResult.USE_RECORD
 
 
 def setup(adapter: SourceAdapter) -> None:
@@ -69,7 +75,7 @@ def setup(adapter: SourceAdapter) -> None:
     adapter.configure_model(
         "dcim.cabletermination",
         extend_content_type="dcim.cable",
-        pre_import=_pre_import_cable_termination,
+        pre_import_record=_pre_import_cable_termination,
     )
     adapter.configure_model(
         "dcim.interface",

nautobot_netbox_importer/diffsync/models/object_change.py

@@ -1,23 +1,23 @@
 """NetBox to Nautobot Object Change Model Mapping."""
 
 from nautobot_netbox_importer.base import RecordData
-from nautobot_netbox_importer.generator import ImporterPass, PreImportResult, SourceAdapter, fields
+from nautobot_netbox_importer.generator import ImporterPass, PreImportRecordResult, SourceAdapter, fields
 
 
 def setup(adapter: SourceAdapter) -> None:
     """Map NetBox object change to Nautobot."""
 
-    def skip_disabled_object_types(source: RecordData, importer_pass: ImporterPass) -> PreImportResult:
+    def skip_disabled_object_types(source: RecordData, importer_pass: ImporterPass) -> PreImportRecordResult:
         """Disabled object types are not in Nautobot and should be skipped."""
         if importer_pass != ImporterPass.IMPORT_DATA:
-            return PreImportResult.USE_RECORD
+            return PreImportRecordResult.USE_RECORD
         object_type = source.get("changed_object_type", None)
         wrapper = adapter.get_or_create_wrapper(object_type)
-        return PreImportResult.SKIP_RECORD if wrapper.disable_reason else PreImportResult.USE_RECORD
+        return PreImportRecordResult.SKIP_RECORD if wrapper.disable_reason else PreImportRecordResult.USE_RECORD
 
     adapter.configure_model(
         "extras.ObjectChange",
-        pre_import=skip_disabled_object_types,
+        pre_import_record=skip_disabled_object_types,
         disable_related_reference=True,
         fields={
             "postchange_data": "object_data",

nautobot_netbox_importer/generator/__init__.py

@@ -7,7 +7,7 @@
     DiffSyncBaseModel,
     ImporterPass,
     InvalidChoiceValueIssue,
-    PreImportResult,
+    PreImportRecordResult,
     SourceAdapter,
     SourceContentType,
     SourceDataGenerator,
@@ -26,7 +26,7 @@
     "InternalFieldType",
     "InvalidChoiceValueIssue",
     "NautobotAdapter",
-    "PreImportResult",
+    "PreImportRecordResult",
     "SourceAdapter",
     "SourceContentType",
     "SourceDataGenerator",

nautobot_netbox_importer/generator/fields.py

@@ -18,7 +18,7 @@
     ImporterPass,
     InternalFieldType,
     InvalidChoiceValueIssue,
-    PreImportResult,
+    PreImportRecordResult,
     RecordData,
     SourceAdapter,
     SourceContentType,
@@ -116,7 +116,7 @@ def role(
     e.g., RackRole with `name = "Network"` and DeviceRole with `name = "Network"` to avoid duplicates.
     """
 
-    def cache_roles(source: RecordData, importer_pass: ImporterPass) -> PreImportResult:
+    def cache_roles(source: RecordData, importer_pass: ImporterPass) -> PreImportRecordResult:
         if importer_pass == ImporterPass.DEFINE_STRUCTURE:
             name = source.get("name", "").capitalize()
             if not name:
@@ -126,12 +126,12 @@ def cache_roles(source: RecordData, importer_pass: ImporterPass) -> PreImportRes
             if not uid:
                 _ROLE_NAME_TO_UID_CACHE[name] = nautobot_uid
 
-        return PreImportResult.USE_RECORD
+        return PreImportRecordResult.USE_RECORD
 
     role_wrapper = adapter.configure_model(
         source_content_type,
         nautobot_content_type="extras.role",
-        pre_import=cache_roles,
+        pre_import_record=cache_roles,
         identifiers=("name",),
         fields={
             # Include color to allow setting the default Nautobot value, import fails without it.

nautobot_netbox_importer/generator/source.py

@@ -131,7 +131,7 @@ class ImporterPass(Enum):
     IMPORT_DATA = 2
 
 
-class PreImportResult(Enum):
+class PreImportRecordResult(Enum):
     """Pre Import Response."""
 
     SKIP_RECORD = False
@@ -149,7 +149,8 @@ class SourceFieldSource(Enum):
     IDENTIFIER = auto()  # Fields used as identifiers
 
 
-PreImport = Callable[[RecordData, ImporterPass], PreImportResult]
+PreImportRecord = Callable[[RecordData, ImporterPass], PreImportRecordResult]
+PostImportRecord = Callable[[RecordData, DiffSyncBaseModel], None]
 SourceDataGenerator = Callable[[], Iterable[SourceRecord]]
 SourceFieldImporter = Callable[[RecordData, DiffSyncBaseModel], None]
 GetPkFromData = Callable[[RecordData], Uid]
@@ -205,7 +206,8 @@ def configure_model(
         default_reference: Optional[RecordData] = None,
         flags: Optional[DiffSyncModelFlags] = None,
         nautobot_flags: Optional[DiffSyncModelFlags] = None,
-        pre_import: Optional[PreImport] = None,
+        pre_import_record: Optional[PreImportRecord] = None,
+        post_import_record: Optional[PostImportRecord] = None,
         disable_related_reference: Optional[bool] = None,
         forward_references: Optional[ForwardReferences] = None,
         fill_dummy_data: Optional[FillDummyData] = None,
@@ -257,8 +259,10 @@ def configure_model(
             wrapper.flags = flags
         if nautobot_flags is not None:
             wrapper.nautobot.flags = nautobot_flags
-        if pre_import:
-            wrapper.pre_import = pre_import
+        if pre_import_record:
+            wrapper.pre_import_record = pre_import_record
+        if post_import_record:
+            wrapper.post_import_record = post_import_record
         if disable_related_reference is not None:
             wrapper.disable_related_reference = disable_related_reference
         if forward_references:
@@ -349,7 +353,7 @@ def get_nautobot_content_type_uid(self, content_type: ContentTypeValue) -> int:
     def load(self) -> None:
         """Load data from the source."""
         self.import_data()
-        self.post_import()
+        self.post_load()
 
     def import_data(self) -> None:
         """Import data from the source."""
@@ -379,9 +383,9 @@ def import_data(self) -> None:
         for content_type, data in get_source_data():
             self.wrappers[content_type].second_pass(data)
 
-    def post_import(self) -> None:
+    def post_load(self) -> None:
         """Post import processing."""
-        while any(wrapper.post_import() for wrapper in self.wrappers.values()):
+        while any(wrapper.post_process_references() for wrapper in self.wrappers.values()):
             pass
 
         for nautobot_wrapper in self.get_imported_nautobot_wrappers():
@@ -458,7 +462,8 @@ def __init__(self, adapter: SourceAdapter, content_type: ContentTypeStr, nautobo
 
         # Source fields defintions
         self.fields: OrderedDict[FieldName, SourceField] = OrderedDict()
-        self.pre_import: Optional[PreImport] = None
+        self.pre_import_record: Optional[PreImportRecord] = None
+        self.post_import_record: Optional[PostImportRecord] = None
 
         if self.disable_reason:
             self.adapter.logger.debug("Created disabled %s", self)
@@ -504,8 +509,8 @@ def cache_record_uids(self, source: RecordData, nautobot_uid: Optional[Uid] = No
 
     def first_pass(self, data: RecordData) -> None:
         """Firts pass of data import."""
-        if self.pre_import:
-            if self.pre_import(data, ImporterPass.DEFINE_STRUCTURE) != PreImportResult.USE_RECORD:
+        if self.pre_import_record:
+            if self.pre_import_record(data, ImporterPass.DEFINE_STRUCTURE) != PreImportRecordResult.USE_RECORD:
                 self.stats.first_pass_skipped += 1
                 return
 
@@ -522,14 +527,17 @@ def second_pass(self, data: RecordData) -> None:
         if self.disable_reason:
             return
 
-        if self.pre_import:
-            if self.pre_import(data, ImporterPass.IMPORT_DATA) != PreImportResult.USE_RECORD:
+        if self.pre_import_record:
+            if self.pre_import_record(data, ImporterPass.IMPORT_DATA) != PreImportRecordResult.USE_RECORD:
                 self.stats.second_pass_skipped += 1
                 return
 
         self.stats.second_pass_used += 1
 
-        self.import_record(data)
+        target = self.import_record(data)
+
+        if self.post_import_record:
+            self.post_import_record(data, target)
 
     def get_summary(self, content_type_id) -> SourceModelSummary:
         """Get a summary of the model."""
@@ -544,7 +552,8 @@ def get_summary(self, content_type_id) -> SourceModelSummary:
             identifiers=self.identifiers,
             disable_related_reference=self.disable_related_reference,
             forward_references=self.forward_references and self.forward_references.__name__ or None,
-            pre_import=self.pre_import and self.pre_import.__name__ or None,
+            pre_import=self.pre_import_record and self.pre_import_record.__name__ or None,
+            post_import=self.post_import_record and self.post_import_record.__name__ or None,
             fields=sorted(fields, key=lambda field: field.name),
             flags=str(self.flags),
             default_reference_uid=serialize_to_summary(self.default_reference_uid),
@@ -817,7 +826,7 @@ def set_default_reference(self, data: RecordData) -> None:
         """Set the default reference to this model."""
         self.default_reference_uid = self.cache_record(data)
 
-    def post_import(self) -> bool:
+    def post_process_references(self) -> bool:
         """Post import processing.
 
         Assigns referenced content_types to referencing instances.

nautobot_netbox_importer/summary.py

@@ -110,6 +110,7 @@ class SourceModelSummary(NamedTuple):
         disable_related_reference: Whether related references are disabled
         forward_references: Configuration for forward references handling
         pre_import: Pre-import processing function/method name
+        post_import: Post-import processing function/method name
         fields: List of field summaries for this model
         flags: Feature flags applied to this model
         default_reference_uid: Default UID used for reference when actual reference is missing
@@ -124,11 +125,12 @@ class SourceModelSummary(NamedTuple):
     identifiers: Optional[List[FieldName]]
     disable_related_reference: bool
     forward_references: Optional[str]
-    pre_import: Optional[str]
     fields: List[FieldSummary]
     flags: str
     default_reference_uid: Optional[Uid]
     stats: SourceModelStats
+    pre_import: Optional[str] = None
+    post_import: Optional[str] = None
 
 
 class NautobotModelSummary(NamedTuple):