Allow defining a custom geolocation function (fixes #70)

docs/index.rst

@@ -14,6 +14,7 @@ Basic metrics tracking for your `Flask`_ application. The core of library is ver
 The following is optional:
 
 * `freegeoip.net <http://freegeoip.net/>`_ integration for storing geography of the visitor.
+* Alternatively you can define a custom geolocation function, for use with the geoip2 package for example.
 * Unique visitor tracking if you are wanting to use Flask's cookie storage.
 * Summation hooks for live count of common web analysis statistics such as hit counts.
 
@@ -55,6 +56,8 @@ Usage
     app.config['TRACK_USAGE_USE_FREEGEOIP'] = False
     # You can use a different instance of freegeoip like so
     # app.config['TRACK_USAGE_FREEGEOIP_ENDPOINT'] = 'https://example.org/api/'
+    # (You can also define a custom geolocation function when initializing the extension)
+
     app.config['TRACK_USAGE_INCLUDE_OR_EXCLUDE_VIEWS'] = 'include'
 
     # We will just print out the data for the example
@@ -68,7 +71,9 @@ Usage
     t = TrackUsage(app, [
         PrintWriter(),
         OutputWriter(transform=lambda s: "OUTPUT: " + str(s))
-    ])
+        ],
+        custom_geolocation_function # (Optional)
+    )
 
     # Include the view in the metrics
     @t.include
@@ -156,7 +161,8 @@ TRACK_USAGE_USE_FREEGEOIP
 
 **Default**: False
 
-Turn FreeGeoIP integration on or off. If set to true, then geography information is also stored in the usage logs.
+Turn FreeGeoIP integration on or off. If set to true, then the geography information is also stored in the usage logs.
+Alternatively you can define a custom geolocation lookup function when initializing the extension.
 
 .. versionchanged:: 1.1.
    The default server for using geoip integration changed to extreme-ip-lookup.com

src/flask_track_usage/__init__.py

@@ -56,7 +56,7 @@ class TrackUsage(object):
     Tracks basic usage of Flask applications.
     """
 
-    def __init__(self, app=None, storage=None, _fake_time=None):
+    def __init__(self, app=None, storage=None, ip_lookup_func=None, _fake_time=None):
         """
         Create the instance.
 
@@ -78,9 +78,9 @@ def __init__(self, app=None, storage=None, _fake_time=None):
         self._fake_time = _fake_time
 
         if app is not None and storage is not None:
-            self.init_app(app, storage)
+            self.init_app(app, storage, ip_lookup_func)
 
-    def init_app(self, app, storage):
+    def init_app(self, app, storage, ip_lookup_func):
         """
         Initialize the instance with the app.
 
@@ -90,6 +90,7 @@ def init_app(self, app, storage):
         """
         self.app = app
         self._storages = storage
+        self.ip_lookup_func = ip_lookup_func
         self._use_freegeoip = app.config.get(
             'TRACK_USAGE_USE_FREEGEOIP', False)
         self._freegeoip_endpoint = app.config.get(
@@ -190,7 +191,10 @@ def after_request(self, response):
             data['username'] = str(ctx.request.authorization.username)
         elif getattr(self.app, 'login_manager', None) and current_user and not current_user.is_anonymous:
             data['username'] = str(current_user)
-        if self._use_freegeoip:
+        if self.ip_lookup_func:
+            clean_ip = quote_plus(str(ctx.request.remote_addr))
+            data['ip_info'] = self.ip_lookup_func(clean_ip)
+        elif self._use_freegeoip:
             clean_ip = quote_plus(str(ctx.request.remote_addr))
             if '{ip}' in self._freegeoip_endpoint:
                 url = self._freegeoip_endpoint.format(ip=clean_ip)

src/flask_track_usage/storage/sql.py

@@ -79,6 +79,7 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage",
         """
 
         import sqlalchemy as sql
+        from sqlalchemy.dialects.postgresql import JSONB
         if db:
             self._eng = db.engine
             self._metadata = db.metadata
@@ -106,7 +107,7 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage",
                     sql.Column('remote_addr', sql.String(24)),
                     sql.Column('xforwardedfor', sql.String(24)),
                     sql.Column('authorization', sql.Boolean),
-                    sql.Column('ip_info', sql.String(1024)),
+                    sql.Column('ip_info', JSONB()),
                     sql.Column('path', sql.String(128)),
                     sql.Column('speed', sql.Float),
                     sql.Column('datetime', sql.DateTime),
@@ -118,7 +119,10 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage",
             else:
                 self._metadata.reflect(bind=self._eng)
                 self.track_table = self._metadata.tables[table_name]
-
+                 # Patch the ip_info column type after reflection
+                if 'ip_info' in self.track_table.c:
+                    self.track_table.c.ip_info.type = JSONB()
+
     def store(self, data):
         """
         Executed on "function call".
@@ -131,16 +135,7 @@ def store(self, data):
         """
         user_agent = data["user_agent"]
         utcdatetime = datetime.datetime.fromtimestamp(data['date'])
-        if data["ip_info"]:
-            t = {}
-            for key in data["ip_info"]:
-                t[key] = data["ip_info"][key]
-                if not len(json.dumps(t)) < 1024:
-                    del t[key]
-                    break
-            ip_info_str = json.dumps(t)
-        else:
-            ip_info_str = None
+
         with self._eng.begin() as con:
             stmt = self.track_table.insert().values(
                 url=data['url'],
@@ -156,7 +151,7 @@ def store(self, data):
                 remote_addr=data["remote_addr"],
                 xforwardedfor=data["xforwardedfor"],
                 authorization=data["authorization"],
-                ip_info=ip_info_str,
+                ip_info=data["ip_info"],
                 path=data["path"],
                 speed=data["speed"],
                 datetime=utcdatetime,