Coverage for ipam/querysets.py: 30%

1import heapq

3import netaddr

4from django.contrib.contenttypes.models import ContentType

5from django.db.models import Count, F, OuterRef, Q, Subquery, Value

6from django.db.models.expressions import RawSQL

7from django.db.models.functions import Cast, NullIf, Round

9from utilities.query import count_related

10from utilities.querysets import RestrictedQuerySet

12from .fields import IPAddressField

13from .lookups import Host

15__all__ = (

16 'ASNRangeQuerySet',

17 'IPAddressQuerySet',

18 'IPRangeQuerySet',

19 'PrefixQuerySet',

20 'VLANGroupQuerySet',

21 'VLANQuerySet',

22)

24# The host portion of an IP address (mask ignored), in the same form as the

25# ipam_ipaddress_host expression index.

26HOST_ADDRESS = Cast(Host('address'), output_field=IPAddressField())

29def _merge_intervals(intervals):

30 """

31 Return the union of (start, end) netaddr.IPAddress intervals, merged and sorted.

32 """

33 if not intervals: 33 ↛ 36line 33 didn't jump to line 36 because the condition on line 33 was always true

34 return []

36 intervals = sorted(intervals)

37 merged = [intervals[0]]

39 for start, end in intervals[1:]:

40 current_start, current_end = merged[-1]

41 # Adjacency math in int space; netaddr raises at the address-space maximum.

42 if start.version == current_end.version and int(start) <= int(current_end) + 1:

43 merged[-1] = (current_start, max(current_end, end))

44 else:

45 merged.append((start, end))

47 return merged

50class ASNRangeQuerySet(RestrictedQuerySet):

52 def annotate_asn_counts(self):

53 """

54 Annotate the number of ASNs which appear within each range.

55 """

56 from .models import ASN

58 # Because ASN does not have a foreign key to ASNRange, we create a fake column "_" with a consistent value

59 # that we can use to count ASNs and return a single value per ASNRange.

60 asns = ASN.objects.filter(

61 asn__gte=OuterRef('start'),

62 asn__lte=OuterRef('end')

63 ).order_by().annotate(_=Value(1)).values('_').annotate(c=Count('*')).values('c')

65 return self.annotate(asn_count=Subquery(asns))

68class IPAddressQuerySet(RestrictedQuerySet):

70 def count_distinct_hosts(self, exclude_intervals=()):

71 """

72 Count distinct host addresses, optionally excluding (start, end) netaddr.IPAddress intervals.

73 """

74 queryset = self

75 for start, end in exclude_intervals:

76 queryset = queryset.exclude(address__host_between=(start, end))

78 return queryset.aggregate(count=Count(HOST_ADDRESS, distinct=True))['count']

80 def count_distinct_hosts_pair(self, bounds, bounded_exclude=(), total_exclude=()):

81 """

82 Return two distinct host counts computed in a single scan, as a dict:

83 'bounded' counts hosts within the (first_ip, last_ip) bounds excluding the

84 bounded_exclude intervals; 'total' counts all hosts excluding the

85 total_exclude intervals. Interval arguments match the output of

86 IPRangeQuerySet.get_intervals(). Avoids a second scan of the host expression

87 index when both counts are needed. Use only when both counts are needed (e.g.

88 Prefix.get_ip_usage_summary()); single-purpose callers should prefer

89 count_distinct_hosts().

90 """

91 # The deduplicated column is already a bare host; plain comparisons beat

92 # the host_between lookup here, which would re-wrap it in HOST()::inet.

93 bounded_q = Q(host_address__range=(str(bounds[0]), str(bounds[1])))

94 for start, end in bounded_exclude:

95 bounded_q &= ~Q(host_address__range=(str(start), str(end)))

96 total_q = Q()

97 for start, end in total_exclude:

98 total_q &= ~Q(host_address__range=(str(start), str(end)))

100 hosts = self.order_by().annotate(host_address=HOST_ADDRESS).values('host_address').distinct()

101 return hosts.aggregate(

102 bounded=Count('host_address', filter=bounded_q),

103 # An empty Q is falsy; fall back to a plain count of all hosts.

104 total=Count('host_address', filter=total_q or None),

105 )

106

107 def _iter_distinct_hosts(self, first_ip, last_ip, batch_size):

108 """

109 Yield the distinct occupied hosts in [first_ip, last_ip] in ascending order,

110 fetched in LIMIT batches that resume just past the last seen host. (A

111 server-side cursor is unsuitable here: on autocommit connections Django

112 declares it WITH HOLD, which materializes the full result at DECLARE.)

113 """

114 resume = first_ip

115 while True:

116 # order_by() first clears the default ordering, which would otherwise

117 # leak into SELECT and break distinct().

118 hosts = list(

119 self.filter(address__host_between=(resume, last_ip))

120 .order_by()

121 .annotate(host_address=HOST_ADDRESS)

122 .values_list('host_address', flat=True)

123 .distinct()

124 .order_by('host_address')[:batch_size]

125 )

126 for host in hosts:

127 yield host.ip

128 if len(hosts) < batch_size: 128 ↛ 130line 128 didn't jump to line 130 because the condition on line 128 was always true

129 return

130 last_host = hosts[-1].ip

131 if int(last_host) >= int(last_ip):

132 return

133 resume = netaddr.IPAddress(int(last_host) + 1, version=last_host.version)

134

135 def available_intervals(self, first_ip, last_ip, exclude_intervals=(), batch_size=5000):

136 """

137 Yield the unoccupied (start, end) netaddr.IPAddress intervals (inclusive)

138 within [first_ip, last_ip], in ascending order. exclude_intervals are

139 (start, end) netaddr.IPAddress pairs; they are merged and sorted internally,

140 intervals of a foreign address family are ignored, and addresses they cover

141 count as occupied. Consumption is lazy: a caller that stops early stops

142 fetching host batches.

143 """

144 if batch_size < 1: 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 raise ValueError('batch_size must be greater than zero')

146

147 first_int, last_int = int(first_ip), int(last_ip)

148 version = first_ip.version

149

150 if first_int > last_int: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true

151 return

152 # Normalize: the sweep below requires sorted, non-overlapping, same-family intervals.

153 exclude_intervals = _merge_intervals([

154 (start, end)

155 for start, end in exclude_intervals

156 if start.version == end.version == version

157 ])

158 intervals = [(int(start), int(end)) for start, end in exclude_intervals]

159

160 # Fast path: one merged excluded interval covers the entire span.

161 if intervals and intervals[0][0] <= first_int and intervals[0][1] >= last_int: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true

162 return

163

164 hosts = (

165 (int(host), int(host))

166 for host in self._iter_distinct_hosts(first_ip, last_ip, batch_size)

167 )

168

169 candidate = first_int

170 # Ties on `start` are harmless; the sweep handles overlapping intervals.

171 for start, end in heapq.merge(intervals, hosts):

172 if end < candidate: 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true

173 continue

174 if start > candidate: 174 ↛ 175line 174 didn't jump to line 175 because the condition on line 174 was never true

175 yield (

176 netaddr.IPAddress(candidate, version=version),

177 netaddr.IPAddress(min(start - 1, last_int), version=version),

178 )

179 candidate = max(candidate, end + 1)

180 if candidate > last_int: 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true

181 return

182

183 if candidate <= last_int: 183 ↛ exitline 183 didn't return from function 'available_intervals' because the condition on line 183 was always true

184 yield (

185 netaddr.IPAddress(candidate, version=version),

186 netaddr.IPAddress(last_int, version=version),

187 )

188

189 def first_available_host(self, first_ip, last_ip, exclude_intervals=()):

190 """

191 Return the first host in [first_ip, last_ip] neither present nor in an excluded interval (or None).

192 """

193 interval = next(self.available_intervals(first_ip, last_ip, exclude_intervals), None)

194 return interval[0] if interval else None

195

196

197class IPRangeQuerySet(RestrictedQuerySet):

198

199 def get_intervals(self, first_ip=None, last_ip=None):

200 """

201 Return ranges as merged (start, end) netaddr.IPAddress intervals, optionally clipped to the bounds.

202 """

203 intervals = []

204

205 # order_by() clears the default ordering; _merge_intervals() sorts anyway.

206 for start_address, end_address in self.order_by().values_list('start_address', 'end_address'): 206 ↛ 207line 206 didn't jump to line 207 because the loop on line 206 never started

207 start, end = start_address.ip, end_address.ip

208

209 if first_ip is not None:

210 if end < first_ip:

211 continue

212 start = max(start, first_ip)

213

214 if last_ip is not None:

215 if start > last_ip:

216 continue

217 end = min(end, last_ip)

218

219 intervals.append((start, end))

220

221 return _merge_intervals(intervals)

222

223

224class PrefixQuerySet(RestrictedQuerySet):

225

226 def annotate_hierarchy(self):

227 """

228 Annotate the depth and number of child prefixes for each Prefix. Cast null VRF values to zero for

229 comparison. (NULL != NULL).

230 """

231 return self.annotate(

232 hierarchy_depth=RawSQL(

233 'SELECT COUNT(DISTINCT U0."prefix") AS "c" '

234 'FROM "ipam_prefix" U0 '

235 'WHERE (U0."prefix" >> "ipam_prefix"."prefix" '

236 'AND COALESCE(U0."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',

237 ()

238 ),

239 hierarchy_children=RawSQL(

240 'SELECT COUNT(U1."prefix") AS "c" '

241 'FROM "ipam_prefix" U1 '

242 'WHERE (U1."prefix" << "ipam_prefix"."prefix" '

243 'AND COALESCE(U1."vrf_id", 0) = COALESCE("ipam_prefix"."vrf_id", 0))',

244 ()

245 )

246 )

247

248

249class VLANGroupQuerySet(RestrictedQuerySet):

250

251 def annotate_utilization(self):

252 from .models import VLAN

253

254 # NullIf guards against legacy rows where total_vlan_ids was miscounted to

255 # 0 by the pre-fix VLANGroup.save(); without it, the annotation 500s.

256 return self.annotate(

257 vlan_count=count_related(VLAN, 'group'),

258 utilization=Round(F('vlan_count') * 100.0 / NullIf(F('total_vlan_ids'), Value(0)), 2),

259 )

260

261

262class VLANQuerySet(RestrictedQuerySet):

263

264 def get_for_site(self, site):

265 """

266 Return all VLANs in the specified site

267 """

268 from .models import VLANGroup

269 q = Q()

270 q |= Q(

271 scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),

272 scope_id=site.pk

273 )

274

275 if site.region:

276 q |= Q(

277 scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),

278 scope_id__in=site.region.get_ancestors(include_self=True)

279 )

280 if site.group:

281 q |= Q(

282 scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

283 scope_id__in=site.group.get_ancestors(include_self=True)

284 )

285

286 return self.filter(

287 Q(group__in=VLANGroup.objects.filter(q)) |

288 Q(site=site) |

289 Q(group__scope_id__isnull=True, site__isnull=True) | # Global group VLANs

290 Q(group__isnull=True, site__isnull=True) # Global VLANs

291 )

292

293 def get_for_site_group(self, site_group):

294 """

295 Return all VLANs available to the specified site group.

296 """

297 if site_group is None:

298 return self.none()

299 from .models import VLANGroup

300 q = Q(

301 scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

302 scope_id__in=site_group.get_ancestors(include_self=True)

303 )

304 return self.filter(

305 Q(group__in=VLANGroup.objects.filter(q)) |

306 Q(group__scope_id__isnull=True, site__isnull=True) | # Global group VLANs

307 Q(group__isnull=True, site__isnull=True) # Global VLANs

308 )

309

310 def get_for_device(self, device):

311 """

312 Return all VLANs available to the specified Device.

313 """

314 from .models import VLANGroup

315

316 # Find all relevant VLANGroups

317 q = Q()

318 if device.cluster_id:

319 # The Device's physical scope is evaluated below. For valid assignments,

320 # the Cluster's physical scope is already represented by that hierarchy.

321 q |= Q(

322 scope_type=ContentType.objects.get_by_natural_key('virtualization', 'cluster'),

323 scope_id=device.cluster_id

324 )

325 if device.cluster.group_id:

326 q |= Q(

327 scope_type=ContentType.objects.get_by_natural_key('virtualization', 'clustergroup'),

328 scope_id=device.cluster.group_id

329 )

330 if device.site.region:

331 q |= Q(

332 scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),

333 scope_id__in=device.site.region.get_ancestors(include_self=True)

334 )

335 if device.site.group:

336 q |= Q(

337 scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

338 scope_id__in=device.site.group.get_ancestors(include_self=True)

339 )

340 q |= Q(

341 scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),

342 scope_id=device.site_id

343 )

344 if device.location:

345 q |= Q(

346 scope_type=ContentType.objects.get_by_natural_key('dcim', 'location'),

347 scope_id__in=device.location.get_ancestors(include_self=True)

348 )

349 if device.rack:

350 q |= Q(

351 scope_type=ContentType.objects.get_by_natural_key('dcim', 'rack'),

352 scope_id=device.rack_id

353 )

354

355 # Return all applicable VLANs

356 return self.filter(

357 Q(group__in=VLANGroup.objects.filter(q)) |

358 Q(site=device.site) |

359 Q(group__scope_id__isnull=True, site__isnull=True) | # Global group VLANs

360 Q(group__isnull=True, site__isnull=True) # Global VLANs

361 )

362

363 def get_for_virtualmachine(self, vm):

364 """

365 Return all VLANs available to the specified VirtualMachine.

366 """

367 from .models import VLANGroup

368

369 # Find all relevant VLANGroups

370 q = Q()

371 site = vm.site

372 if vm.cluster:

373 # Add VLANGroups scoped to the assigned cluster (or its group)

374 q |= Q(

375 scope_type=ContentType.objects.get_by_natural_key('virtualization', 'cluster'),

376 scope_id=vm.cluster_id

377 )

378 if vm.cluster.group:

379 q |= Q(

380 scope_type=ContentType.objects.get_by_natural_key('virtualization', 'clustergroup'),

381 scope_id=vm.cluster.group_id

382 )

383 # Looking all possible cluster scopes

384 if vm.cluster.scope_type == ContentType.objects.get_by_natural_key('dcim', 'location'):

385 site = site or vm.cluster.scope.site

386 q |= Q(

387 scope_type=ContentType.objects.get_by_natural_key('dcim', 'location'),

388 scope_id__in=vm.cluster.scope.get_ancestors(include_self=True)

389 )

390 elif vm.cluster.scope_type == ContentType.objects.get_by_natural_key('dcim', 'site'):

391 site = site or vm.cluster.scope

392 q |= Q(

393 scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),

394 scope_id=vm.cluster.scope.pk

395 )

396 elif vm.cluster.scope_type == ContentType.objects.get_by_natural_key('dcim', 'sitegroup'):

397 q |= Q(

398 scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

399 scope_id__in=vm.cluster.scope.get_ancestors(include_self=True)

400 )

401 elif vm.cluster.scope_type == ContentType.objects.get_by_natural_key('dcim', 'region'):

402 q |= Q(

403 scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),

404 scope_id__in=vm.cluster.scope.get_ancestors(include_self=True)

405 )

406 # VM can be assigned to a site without a cluster so checking assigned site independently

407 if site:

408 # Add VLANGroups scoped to the assigned site (or its group or region)

409 q |= Q(

410 scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),

411 scope_id=site.pk

412 )

413 if site.region:

414 q |= Q(

415 scope_type=ContentType.objects.get_by_natural_key('dcim', 'region'),

416 scope_id__in=site.region.get_ancestors(include_self=True)

417 )

418 if site.group:

419 q |= Q(

420 scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

421 scope_id__in=site.group.get_ancestors(include_self=True)

422 )

423 vlan_groups = VLANGroup.objects.filter(q)

424

425 # Return all applicable VLANs

426 q = (

427 Q(group__in=vlan_groups) |

428 Q(group__scope_id__isnull=True, site__isnull=True) | # Global group VLANs

429 Q(group__isnull=True, site__isnull=True) # Global VLANs

430 )

431 if site:

432 q |= Q(site=site)

433

434 return self.filter(q)

435

436 def get_related_to_sites(self, sites, *, negate=False):

437 """

438 Return VLANs related to any of the given sites, directly or through a scoped group.

439 Unlike get_for_site(), region-scoped and globally available VLANs are excluded.

440 Pass negate=True to exclude the related VLANs instead.

441 """

442 from dcim.models import SiteGroup

443

444 site_ids = set()

445 site_group_ids = set()

446 for site in sites:

447 site_ids.add(site.pk)

448 if site.group_id:

449 site_group_ids.add(site.group_id)

450

451 if not site_ids:

452 return self if negate else self.none()

453

454 q = Q(site_id__in=site_ids) | Q(

455 group__scope_type=ContentType.objects.get_by_natural_key('dcim', 'site'),

456 group__scope_id__in=site_ids

457 )

458

459 if site_group_ids:

460 # A site group scope reaches its descendants, so match the selected groups' ancestors.

461 ancestor_groups = SiteGroup.objects.none()

462 for site_group in SiteGroup.objects.filter(pk__in=site_group_ids):

463 ancestor_groups |= site_group.get_ancestors(include_self=True)

464

465 q |= Q(

466 group__scope_type=ContentType.objects.get_by_natural_key('dcim', 'sitegroup'),

467 group__scope_id__in=ancestor_groups.values('pk')

468 )

469

470 return self.exclude(q) if negate else self.filter(q)