Customizing network importer with adapters

adapter.py

from my_custom_network_importer.models import CustomSite
from network_importer.adapters.network_importer.adapter import NetworkImporterAdapter
from network_importer.adapters.nautobot_api.adapter import NautobotAPIAdapter

class CustomNetworkImporterAdapter(NetworkImporterAdapter):
    
    site = CustomSite

    def load_batfish_device(self, device):
        super().load_batfish_device(device=device)
        self.load_batfish_site_asn(device=device)

    def load_batfish_site_asn(self, device):
        """Import Site ASN from the list of BGP Peers.

        Args:
            device (Device) Device object
        """
        bgp_sessions = (
            self.bfi.q.bgpPeerConfiguration(nodes=device.name).answer().frame()
        )

        if len(bgp_sessions) == 0:
            LOGGER.debug("No BGP session found for %s", device.name)
            return

        site = self.get(self.site, identifier=device.site_name)
        local_asn = int(bgp_sessions.iloc[0].Local_AS)

        if not site.asn:
            site.asn = local_asn
            LOGGER.info("Identified Site ASN as %s (from %s)", local_asn, device.name)
        elif site.asn != local_asn:
            LOGGER.warning(
                "Identified Site ASN as %s from %s but ASN is already configured as %s",
                local_asn,
                device.name,
                site.asn,
            )
        elif site.asn == local_asn:
            LOGGER.debug(
                "Identified Site ASN as %s (from %s), same as existing value",
                local_asn,
                device.name,
            )

class CustomNautobotAPIAdapter(NautobotAPIAdapter):

    site = CustomSite

    def load_nautobot_prefix(self, site):
        """Extend default method to also import Aggregate and ASN from Nautobot for a given site.

        Args:
            site (NautobotSite): Site to import prefix and aggregate from
        """

        super().load_nautobot_prefix(site=site)

        # Import Site ASN
        nb_site = self.nautobot.dcim.sites.get(site.remote_id)
        if nb_site.asn:
            site.asn = nb_site.asn
            LOGGER.info("Update %s ASN to %s", site.name, site.asn)

models.py

import pynautobot

from network_importer.adapters.nautobot_api.models import NautobotSite

class CustomSite(NautobotSite):

    asn: Optional[int]
    _attributes = tuple(list(NautobotSite._attributes) + ["asn"])


    def update(self, attrs: dict):
        """Update Site in Nautobot.

        Returns:
            CustomSite: DiffSync object
        """

        if not attrs["asn"] and self.asn:
            LOGGER.warning(
                "Unable to remove the ASN on an existing Site (Restricted). Please update %s manually",
                self.diffsync.name,
            )
            return super().update(attrs)

        nb_params = {"asn": attrs["asn"]}

        try:
            site = self.diffsync.nautobot.dcim.sites.get(self.remote_id)
            site.update(data=nb_params)
            LOGGER.info(
                "Updated Site %s (%s) in Nautobot", self.get_unique_id(), self.remote_id
            )
        except pynautobot.core.query.RequestError as exc:
            LOGGER.warning(
                "Unable to update Site %s in %s (%s)",
                self.get_unique_id(),
                self.diffsync.name,
                exc.error,
            )
            return None

        return super().update(attrs)

network_importer.toml

[adapters]
network_class = "my_custom_newtwork_importer.adapters.CustomNetworkImporterAdapter"
sot_class = "my_custom_network_importer.adapters.CustomNautobotAPIAdapter"