Implement binary transfer of inet & cidr using ipaddr.

James William Pye · James William Pye · commit 8709585fbaf2 · 2009-05-17T17:08:24.000-07:00
closes 1010579
diff --git a/postgresql/documentation/driver.py b/postgresql/documentation/driver.py
@@ -1525,6 +1525,9 @@
  `postgresql.types.NUMERICOID`     `decimal.Decimal`                  numeric
  `postgresql.types.BYTEAOID`       `bytes`                            bytea
  `postgresql.types.TEXTOID`        `str`                              text
+
+ `postgresql.types.CIDROID`        `ipaddr.IPv4` or `ipaddr.IPv6`     cidr
+ `postgresql.types.INETOID`        `ipaddr.IPv4` or `ipaddr.IPv6`     inet
  ================================= ================================== ===========
 
 The mapping in the above table *normally* goes both ways. So when a parameter
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -1921,7 +1921,15 @@ def connect(self):
 
 		r = self.sys.activity_for(self.backend_id)
 		if r is not None:
-			self.client_address = r.get('client_addr')
+			# conditional initialization of client_address.
+			# pythons without ipaddr will likely give strings.
+			ca = r.get('client_addr')
+			if ca is not None:
+				if not isinstance(ca, str):
+					# it better be a ipaddr.BaseIP..
+					self.client_address = ca.ip_ext
+				else:
+					self.client_address = ca
 			self.client_port = r.get('client_port')
 			self.backend_start = r.get('backend_start')
 		try:
diff --git a/postgresql/protocol/typio.py b/postgresql/protocol/typio.py
@@ -421,6 +421,39 @@ def numeric_unpack(x):
 	pg_types.CIRCLEOID : (circle_pack, circle_unpack),
 }
 
+# Conditionally use ipaddr.
+# Condition goes away when 3.0 is deprecated.
+try:
+	import ipaddr
+	def net_pack(ip):
+		family = None
+		converted = False
+		while family is None:
+			if isinstance(ip, ipaddr.IPv4):
+				family = 2
+			elif isinstance(ip, ipaddr.IPv6):
+				family = 3
+			else:
+				if converted is True:
+					raise ValueError("unknown IP type: " + repr(ip))
+				ip = ipaddr.IP(ip)
+				converted = True
+		return ts.net_pack((family, ip.prefixlen, ip.packed))
+
+	def net_unpack(data):
+		family, mask, data = ts.net_unpack(data)
+		if family in (2, 3):
+			v = ipaddr.IP(data)
+		else:
+			raise ValueError("unknown net family: " + repr(family))
+		v.prefixlen = mask
+		return v
+	oid_to_io[pg_types.CIDROID] = (net_pack, net_unpack)
+	oid_to_io[pg_types.INETOID] = (net_pack, net_unpack)
+except ImportError:
+	oid_to_io[pg_types.CIDROID] = (None, None)
+	oid_to_io[pg_types.INETOID] = (None, None)
+
 def process_tuple(procs, tup, exception_handler):
 	"""
 	Call each item in `procs` with the corresponding
@@ -716,8 +749,6 @@ def __init__(self):
 			pg_types.VARCHAROID : (None, None),
 			pg_types.CSTRINGOID : (None, None),
 			pg_types.TEXTOID : (None, None),
-			pg_types.CIDROID : (None, None),
-			pg_types.INETOID : (None, None),
 			pg_types.REGPROCEDUREOID : (None, None),
 			pg_types.REGTYPEOID : (None, None),
 			pg_types.REGPROCOID : (None, None),
diff --git a/postgresql/protocol/typstruct.py b/postgresql/protocol/typstruct.py
@@ -313,11 +313,13 @@ def varbit_unpack(data):
 	"""
 	return long_unpack(data[0:4]), data[4:]
 
-def cidr_pack(family_mask_data):
+def net_pack(family_mask_data):
 	"""
 	Given a triple, yield the serialized form for transport.
 
 	Prepends the ``family``, ``mask`` and implicit ``is_cidr`` fields.
+
+	Supports cidr and inet types.
 	"""
 	(family, mask, data) = family_mask_data
 
@@ -328,9 +330,8 @@ def cidr_pack(family_mask_data):
 		byte_pack(len(data)),
 		data
 	))
-inet_pack = cidr_pack
 
-def cidr_unpack(data):
+def net_unpack(data):
 	"""
 	Given serialized cidr data, return a tuple:
 
@@ -343,7 +344,6 @@ def cidr_unpack(data):
 		raise ValueError("invalid size parameter")
 
 	return (family, mask, rd)
-inet_unpack = cidr_unpack
 
 def record_unpack(data):
 	"""
@@ -471,9 +471,9 @@ def return_arg(arg):
 	pg_types.BYTEAOID : (bytes, bytes),
 	pg_types.CHAROID : literal,
 
-#	pg_type.MACADDROID : literal,
-#	pg_type.INETOID : (cidr_pack, cidr_unpack),
-#	pg_type.CIDROID : (cidr_pack, cidr_unpack),
+#	pg_types.MACADDROID : literal,
+	pg_types.INETOID : (net_pack, net_unpack),
+	pg_types.CIDROID : (net_pack, net_unpack),
 
 	pg_types.DATEOID : (date_pack, date_unpack),
 	pg_types.ABSTIMEOID : (long_pack, long_unpack),
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -19,7 +19,7 @@
 from .. import unittest as pg_unittest
 from .. import lib as pg_lib
 
-type_samples = (
+type_samples = [
 	('smallint', (
 			((1 << 16) // 2) - 1, - ((1 << 16) // 2),
 			-1, 0, 1,
@@ -268,7 +268,40 @@
 			pg_types.varbit('010111101111'),
 		],
 	),
-)
+]
+
+try:
+	import ipaddr
+	type_samples.append((
+		'inet', [
+			ipaddr.IPv4('255.255.255.255'),
+			ipaddr.IPv4('127.0.0.1'),
+			ipaddr.IPv4('10.0.0.1'),
+			ipaddr.IPv4('0.0.0.0'),
+			ipaddr.IPv6('::1'),
+			ipaddr.IPv6('ffff' + ':ffff'*7),
+			ipaddr.IPv6('fe80::1'),
+			ipaddr.IPv6('fe80::1'),
+			ipaddr.IPv6('0::0'),
+		],
+	))
+	type_samples.append((
+		'cidr', [
+			ipaddr.IPv4('255.255.255.255/32'),
+			ipaddr.IPv4('127.0.0.0/8'),
+			ipaddr.IPv4('127.1.0.0/16'),
+			ipaddr.IPv4('10.0.0.0/32'),
+			ipaddr.IPv4('0.0.0.0/0'),
+			ipaddr.IPv6('ffff' + ':ffff'*7 + '/128'),
+			ipaddr.IPv6('::1/128'),
+			ipaddr.IPv6('fe80::1/128'),
+			ipaddr.IPv6('fe80::0/64'),
+			ipaddr.IPv6('fe80::0/16'),
+			ipaddr.IPv6('0::0/0'),
+		],
+	))
+except ImportError:
+	pass
 
 class test_driver(pg_unittest.TestCaseWithCluster):
 	"""
@@ -873,7 +906,7 @@ def testTypes(self):
 				"SELECT $1::" + typname
 			)
 			for sample in sample_data:
-				rsample = pb.first(sample)
+				rsample = list(pb.rows(sample))[0][0]
 				if isinstance(rsample, pg_types.Array):
 					rsample = rsample.nest()
 				self.failUnless(
