Merge pull request #33 from nkinkade/file-descriptor

Adds ability to pass a file descriptor when opening a new Reader() object

Author: oschwald

README.rst

@@ -40,17 +40,22 @@ provide `free GeoLite2 databases
 <http://dev.maxmind.com/geoip/geoip2/geolite2>`_. These files must be
 decompressed with ``gunzip``.
 
-After you have obtained a database and importing the module, call
-``open_database`` with a path to the database as the first argument.
-Optionally, you may pass a mode as the second argument. The modes are
-exported from ``maxminddb``. Valid modes are:
+After you have obtained a database and imported the module, call
+``open_database`` with a path, or file descriptor (in the case of MODE_FD),
+to the database as the first argument. Optionally, you may pass a mode as the
+second argument. The modes are exported from ``maxminddb``. Valid modes are:
 
 * MODE_MMAP_EXT - use the C extension with memory map.
 * MODE_MMAP - read from memory map. Pure Python.
 * MODE_FILE - read database as standard file. Pure Python.
 * MODE_MEMORY - load database into memory. Pure Python.
+* MODE_FD - load database into memory from a file descriptor. Pure Python.
 * MODE_AUTO - try MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order. Default.
 
+**NOTE**: When using ``MODE_FD``, it is the *caller's* responsibility to be
+sure that the file descriptor gets closed properly, even though this module
+*may* close it after the ``Reader`` object is created.
+
 The ``open_database`` function returns a ``Reader`` object. To look up an IP
 address, use the ``get`` method on this object. The method will return the
 corresponding values for the IP address from the database (e.g., a dictionary

maxminddb/__init__.py

@@ -9,21 +9,23 @@
     maxminddb.extension = None
 
 from maxminddb.const import (MODE_AUTO, MODE_MMAP, MODE_MMAP_EXT, MODE_FILE,
-                             MODE_MEMORY)
+                             MODE_MEMORY, MODE_FD)
 from maxminddb.decoder import InvalidDatabaseError
 
 
 def open_database(database, mode=MODE_AUTO):
     """Open a Maxmind DB database
 
     Arguments:
-        database -- A path to a valid MaxMind DB file such as a GeoIP2
-                    database file.
+        database -- A path to a valid MaxMind DB file such as a GeoIP2 database
+                    file, or a file descriptor in the case of MODE_FD.
         mode -- mode to open the database with. Valid mode are:
             * MODE_MMAP_EXT - use the C extension with memory map.
             * MODE_MMAP - read from memory map. Pure Python.
             * MODE_FILE - read database as standard file. Pure Python.
             * MODE_MEMORY - load database into memory. Pure Python.
+            * MODE_FD - the param passed via database is a file descriptor, not
+                        a path. This mode implies MODE_MEMORY.
             * MODE_AUTO - tries MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that
                           order. Default mode.
     """
@@ -35,7 +37,7 @@ def open_database(database, mode=MODE_AUTO):
                 "MODE_MMAP_EXT requires the maxminddb.extension module to be available"
             )
         return maxminddb.extension.Reader(database)
-    elif mode in (MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY):
+    elif mode in (MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY, MODE_FD):
         return maxminddb.reader.Reader(database, mode)
     raise ValueError('Unsupported open mode: {0}'.format(mode))
 

maxminddb/const.py

@@ -5,3 +5,4 @@
 MODE_MMAP = 2
 MODE_FILE = 4
 MODE_MEMORY = 8
+MODE_FD = 16

maxminddb/reader.py

@@ -17,7 +17,7 @@
 
 from maxminddb.compat import (byte_from_int, compat_ip_address, string_type,
                               string_type_name)
-from maxminddb.const import MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY
+from maxminddb.const import MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY, MODE_FD
 from maxminddb.decoder import Decoder
 from maxminddb.errors import InvalidDatabaseError
 from maxminddb.file import FileBuffer
@@ -38,30 +38,39 @@ def __init__(self, database, mode=MODE_AUTO):
         """Reader for the MaxMind DB file format
 
         Arguments:
-        database -- A path to a valid MaxMind DB file such as a GeoIP2
-                    database file.
+        database -- A path to a valid MaxMind DB file such as a GeoIP2 database
+                    file, or a file descriptor in the case of MODE_FD.
         mode -- mode to open the database with. Valid mode are:
             * MODE_MMAP - read from memory map.
             * MODE_FILE - read database as standard file.
             * MODE_MEMORY - load database into memory.
             * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
+            * MODE_FD - the param passed via database is a file descriptor, not
+                        a path. This mode implies MODE_MEMORY.
         """
         if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
             with open(database, 'rb') as db_file:
                 self._buffer = mmap.mmap(
                     db_file.fileno(), 0, access=mmap.ACCESS_READ)
                 self._buffer_size = self._buffer.size()
+            filename = database
         elif mode in (MODE_AUTO, MODE_FILE):
             self._buffer = FileBuffer(database)
             self._buffer_size = self._buffer.size()
+            filename = database
         elif mode == MODE_MEMORY:
             with open(database, 'rb') as db_file:
                 self._buffer = db_file.read()
                 self._buffer_size = len(self._buffer)
+            filename = database
+        elif mode == MODE_FD:
+            self._buffer = database.read()
+            self._buffer_size = len(self._buffer)
+            filename = database.name
         else:
             raise ValueError(
-                'Unsupported open mode ({0}). Only MODE_AUTO, '
-                ' MODE_FILE, and MODE_MEMORY are support by the pure Python '
+                'Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, '
+                'MODE_MEMORY and MODE_FD are supported by the pure Python '
                 'Reader'.format(mode))
 
         metadata_start = self._buffer.rfind(
@@ -72,7 +81,7 @@ def __init__(self, database, mode=MODE_AUTO):
             self.close()
             raise InvalidDatabaseError('Error opening database file ({0}). '
                                        'Is this a valid MaxMind DB file?'
-                                       ''.format(database))
+                                       ''.format(filename))
 
         metadata_start += len(self._METADATA_START_MARKER)
         metadata_decoder = Decoder(self._buffer, metadata_start)

tests/reader_test.py

@@ -4,6 +4,7 @@
 from __future__ import unicode_literals
 
 import logging
+import mock
 import os
 import sys
 import threading
@@ -22,7 +23,7 @@
 from maxminddb import open_database, InvalidDatabaseError
 from maxminddb.compat import FileNotFoundError
 from maxminddb.const import (MODE_AUTO, MODE_MMAP_EXT, MODE_MMAP, MODE_FILE,
-                             MODE_MEMORY)
+                             MODE_MEMORY, MODE_FD)
 
 if sys.version_info[:2] == (2, 6):
     import unittest2 as unittest
@@ -34,6 +35,18 @@
     unittest.TestCase.assertRegex = unittest.TestCase.assertRegexpMatches
 
 
+def get_reader_from_file_descriptor(filepath, mode):
+    """Patches open_database() for class TestFDReader()."""
+    if mode == MODE_FD:
+        with open(filepath, 'rb') as mmdb_fh:
+            return maxminddb.open_database(mmdb_fh, mode)
+    else:
+        # There are a few cases where mode is statically defined in
+        # BaseTestReader(). In those cases just call an unpatched
+        # open_database() with a string path.
+        return maxminddb.open_database(filepath, mode)
+
+
 class BaseTestReader(object):
     def test_reader(self):
         for record_size in [24, 28, 32]:
@@ -205,7 +218,7 @@ def test_double_close(self):
                           'Double close does not throw an exception')
 
     def test_closed_get(self):
-        if self.mode == MODE_MEMORY:
+        if self.mode in [MODE_MEMORY, MODE_FD]:
             return
         reader = open_database(
             'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
@@ -425,6 +438,17 @@ class TestMemoryReader(BaseTestReader, unittest.TestCase):
     readerClass = [maxminddb.reader.Reader]
 
 
+class TestFDReader(BaseTestReader, unittest.TestCase):
+    def setUp(self):
+        self.open_database_patcher = mock.patch('reader_test.open_database')
+        self.addCleanup(self.open_database_patcher.stop)
+        self.open_database = self.open_database_patcher.start()
+        self.open_database.side_effect = get_reader_from_file_descriptor
+
+    mode = MODE_FD
+    readerClass = [maxminddb.reader.Reader]
+
+
 class TestOldReader(unittest.TestCase):
     def test_old_reader(self):
         reader = maxminddb.Reader(