Check that the database is of the right type before doing a lookup

Author: oschwald

HISTORY.rst

@@ -3,6 +3,17 @@
 History
 -------
 
+0.7.0 (2014-09-XX)
+++++++++++++++++++
+
+- The ``geoip2.database.Reader`` lookup methods (e.g., ``city()``, ``isp()``)
+  now raise a ``TypeError`` if they are used with a database that does not
+  match the method. In particular, doing a ``city()`` lookup on a GeoIP2
+  Country database will result in an error and vice versa.
+- A ``metadata()`` method has been added to the ``geoip2.database.Reader``
+  class. This returns a ``maxminddb.reader.Metadata`` object with information
+  about the database.
+
 0.6.0 (2014-07-22)
 ++++++++++++++++++
 

geoip2/database.py

@@ -4,6 +4,8 @@
 ======================
 
 """
+import inspect
+
 import geoip2
 import geoip2.models
 import geoip2.errors
@@ -53,7 +55,9 @@ def country(self, ip_address):
 
         """
 
-        return self._model_for(geoip2.models.Country, ip_address)
+        return self._model_for(geoip2.models.Country,
+                               ['GeoIP2-Country', 'GeoLite2-Country'],
+                               ip_address)
 
     def city(self, ip_address):
         """Get the City object for the IP address
@@ -63,33 +67,9 @@ def city(self, ip_address):
         :returns: :py:class:`geoip2.models.City` object
 
         """
-        return self._model_for(geoip2.models.City, ip_address)
-
-    def city_isp_org(self, ip_address):
-        """Get the City object for the IP address
-
-        :param ip_address: IPv4 or IPv6 address as a string. If no address
-          is provided.
-
-        :returns: :py:class:`geoip2.models.City` object
-
-        .. deprecated:: 0.6.0
-           Use :py:method:`city` instead.
-        """
-
-        return self.city(ip_address)
-
-    def omni(self, ip_address):
-        """Get the Insights object for the IP address
-
-        :param ip_address: IPv4 or IPv6 address as a string.
-
-        :returns: :py:class:`geoip2.models.Insights` object
-
-        .. deprecated:: 0.6.0
-           Use :py:method:`city` instead.
-        """
-        return self._model_for(geoip2.models.Insights, ip_address)
+        return self._model_for(geoip2.models.City,
+                               ['GeoIP2-City', 'GeoLite2-City'],
+                               ip_address)
 
     def connection_type(self, ip_address):
         """Get the ConnectionType object for the IP address
@@ -99,7 +79,9 @@ def connection_type(self, ip_address):
         :returns: :py:class:`geoip2.models.ConnectionType` object
 
         """
-        return self._flat_model_for(geoip2.models.ConnectionType, ip_address)
+        return self._flat_model_for(geoip2.models.ConnectionType,
+                                    ['GeoIP2-Connection-Type'],
+                                    ip_address)
 
     def domain(self, ip_address):
         """Get the Domain object for the IP address
@@ -109,7 +91,9 @@ def domain(self, ip_address):
         :returns: :py:class:`geoip2.models.Domain` object
 
         """
-        return self._flat_model_for(geoip2.models.Domain, ip_address)
+        return self._flat_model_for(geoip2.models.Domain,
+                                    ['GeoIP2-Domain'],
+                                    ip_address)
 
     def isp(self, ip_address):
         """Get the ISP object for the IP address
@@ -119,25 +103,35 @@ def isp(self, ip_address):
         :returns: :py:class:`geoip2.models.ISP` object
 
         """
-        return self._flat_model_for(geoip2.models.ISP, ip_address)
-
-    def _get(self, ip_address):
+        return self._flat_model_for(geoip2.models.ISP,
+                                    ['GeoIP2-ISP'],
+                                    ip_address)
+
+    def _get(self, types, ip_address):
+        if not self.metadata().database_type in types:
+            caller = inspect.stack()[2][3]
+            raise TypeError("The %s method cannot be used with the "
+                            "%s database" %
+                            (caller, self.metadata().database_type))
         record = self._db_reader.get(ip_address)
         if record is None:
             raise geoip2.errors.AddressNotFoundError(
                 "The address %s is not in the database." % ip_address)
         return record
 
-    def _model_for(self, model_class, ip_address):
-        record = self._get(ip_address)
+    def _model_for(self, model_class, types, ip_address):
+        record = self._get(types, ip_address)
         record.setdefault('traits', {})['ip_address'] = ip_address
         return model_class(record, locales=self._locales)
 
-    def _flat_model_for(self, model_class, ip_address):
-        record = self._get(ip_address)
+    def _flat_model_for(self, model_class, types, ip_address):
+        record = self._get(types, ip_address)
         record['ip_address'] = ip_address
         return model_class(record)
 
+    def metadata(self):
+        return self._db_reader.metadata()
+
     def close(self):
         """Closes the GeoIP2 database"""
 

tests/data

@@ -29,7 +29,7 @@ def test_default_locale(self):
 
     def test_language_list(self):
         reader = geoip2.database.Reader(
-            'tests/data/test-data/GeoIP2-City-Test.mmdb',
+            'tests/data/test-data/GeoIP2-Country-Test.mmdb',
             ['xx', 'ru', 'pt-BR', 'es', 'en'])
         record = reader.country('81.2.69.160')
 
@@ -38,7 +38,7 @@ def test_language_list(self):
 
     def test_has_ip_address(self):
         reader = geoip2.database.Reader(
-            'tests/data/test-data/GeoIP2-City-Test.mmdb')
+            'tests/data/test-data/GeoIP2-Country-Test.mmdb')
         record = reader.country('81.2.69.160')
         self.assertEqual(record.traits.ip_address, '81.2.69.160')
         reader.close()
@@ -49,7 +49,16 @@ def test_unknown_address(self):
         with self.assertRaisesRegex(geoip2.errors.AddressNotFoundError,
                                     'The address 10.10.10.10 is not in the '
                                     'database.'):
-            reader.city_isp_org('10.10.10.10')
+            reader.city('10.10.10.10')
+        reader.close()
+
+    def test_wrong_database(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoIP2-City-Test.mmdb')
+        with self.assertRaisesRegex(TypeError,
+                                    'The country method cannot be used with '
+                                    'the GeoIP2-City database'):
+            reader.country('1.1.1.1')
         reader.close()
 
     def test_invalid_address(self):