python-ox/ox/cache.py

# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4
# GPL 2008
import gzip
import zlib
import hashlib
import os
import StringIO
import time
import urlparse
import urllib2
import sqlite3

import chardet
from ox.utils import json

import net
from net import DEFAULT_HEADERS, getEncoding


cache_timeout = 30*24*60*60 # default is 30 days

COMPRESS_TYPES = (
    'text/html',
    'text/plain',
    'text/xml',
    'application/xhtml+xml',
    'application/x-javascript',
    'application/javascript',
    'application/ecmascript',
    'application/rss+xml'
)

def status(url, data=None, headers=DEFAULT_HEADERS, timeout=cache_timeout):
    '''
      >>> status('http://google.com')
      200
      >>> status('http://google.com/mysearch')
      404
    '''
    headers = getHeaders(url, data, headers)
    return int(headers['status'])

def exists(url, data=None, headers=DEFAULT_HEADERS, timeout=cache_timeout):
    '''
      >>> exists('http://google.com')
      True
      >>> exists('http://google.com/mysearch')
      False
    '''
    s = status(url, data, headers, timeout)
    if s >= 200 and s < 400:
        return True
    return False

def getHeaders(url, data=None, headers=DEFAULT_HEADERS, timeout=cache_timeout):
    url_headers = _readUrlCache(url, data, headers, timeout, "headers")
    if url_headers:
        url_headers = json.loads(url_headers)
    else:
        url_headers = net.getHeaders(url, data, headers)
        _saveUrlCache(url, data, -1, url_headers)
    return url_headers

class InvalidResult(Exception):
    """Base class for exceptions in this module."""
    def __init__(self, result, headers):
        self.result = result
        self.headers = headers

def readUrl(url, data=None, headers=DEFAULT_HEADERS, timeout=cache_timeout, valid=None):
    '''
        url     - url to load
        data    - possible post data
        headers - headers to send with request
        timeout - get from cache if cache not older than given seconds, -1 to get from cache
        valid   - function to check if result is ok, its passed result and headers
                  if this function fails, InvalidResult will be raised deal with it in your code 
    '''
    #FIXME: send last-modified / etag from cache and only update if needed
    if isinstance(url, unicode):
        url = url.encode('utf-8')
    result = _readUrlCache(url, data, headers, timeout)
    if not result:
        #print "get data", url
        try:
            url_headers, result = net.readUrl(url, data, headers, returnHeaders=True)
        except urllib2.HTTPError, e:
            e.headers['Status'] = "%s" % e.code
            url_headers = dict(e.headers)
            result = e.read()
            if url_headers.get('content-encoding', None) == 'gzip':
                result = gzip.GzipFile(fileobj=StringIO.StringIO(result)).read()
        if not valid or valid(result, url_headers):
            _saveUrlCache(url, data, result, url_headers)
        else:
            raise InvalidResult(result, url_headers)
    return result

def readUrlUnicode(url, data=None, headers=DEFAULT_HEADERS, timeout=cache_timeout, _readUrl=readUrl, valid=None):
    data = _readUrl(url, data, headers, timeout, valid)
    encoding = getEncoding(data)
    if not encoding:
        encoding = 'latin-1'
    return unicode(data, encoding)

def saveUrl(url, filename, overwrite=False):
    if not os.path.exists(filename) or overwrite:
        dirname = os.path.dirname(filename)
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        data = readUrl(url)
        f = open(filename, 'w')
        f.write(data)
        f.close()

def _getCacheBase():
    'cache base is eather ~/.ox/cache or can set via env variable oxCACHE'
    return os.environ.get('oxCACHE', os.path.expanduser('~/.ox/cache'))

def _getCacheDB():
    path = _getCacheBase()
    if not os.path.exists(path):
        os.makedirs(path)
    return os.path.join(path, "cache.sqlite")

def _connectDb():
    conn = sqlite3.connect(_getCacheDB(), timeout=10)
    conn.text_factory = str
    return conn

def _getSetting(c, key, default=None):
    c.execute('SELECT value FROM setting WHERE key = ?', (key, ))
    for row in c:
        return row[0]
    return default

def _setSetting(c, key, value):
    c.execute(u'INSERT OR REPLACE INTO setting values (?, ?)', (key, str(value)))

def _createDb(conn, c):
    # Create table and indexes 
    c.execute('''CREATE TABLE IF NOT EXISTS cache (url_hash varchar(42) unique, domain text, url text,
                      post_data text, headers text, created int, data blob, only_headers int)''')
    c.execute('''CREATE INDEX IF NOT EXISTS cache_domain ON cache (domain)''')
    c.execute('''CREATE INDEX IF NOT EXISTS cache_url ON cache (url)''')
    c.execute('''CREATE INDEX IF NOT EXISTS cache_url_hash ON cache (url_hash)''')

    c.execute('''CREATE TABLE IF NOT EXISTS setting (key varchar(1024) unique, value text)''')
    if int(_getSetting(c, 'version', 0)) < 1:
        _setSetting(c, 'version', 1)
        c.execute('''ALTER TABLE cache ADD compressed INT DEFAULT 0''')
        conn.commit()

def _readUrlCache(url, data, headers=DEFAULT_HEADERS, timeout=-1, value="data"):
    r = None
    if timeout == 0:
        return r

    if data:
        url_hash = hashlib.sha1(url + '?' + data).hexdigest()
    else:
        url_hash = hashlib.sha1(url).hexdigest()

    conn = _connectDb()
    c = conn.cursor()
    _createDb(conn, c)

    sql = 'SELECT %s, compressed FROM cache WHERE url_hash=?' % value
    if timeout > 0:
        now = time.mktime(time.localtime())
        t = (url_hash, now-timeout)
        sql += ' AND created > ?'
    else:
        t = (url_hash, )
    if value != "headers":
        sql += ' AND only_headers != 1 '
    c.execute(sql, t)
    for row in c:
        r = row[0]
        if value == 'data':
            if row[1] == 1:
                r = zlib.decompress(r)
            else:
                r = str(r)
        break

    c.close()
    conn.close()
    return r

def _saveUrlCache(url, post_data, data, headers):
    if post_data:
        url_hash = hashlib.sha1(url + '?' + post_data).hexdigest()
    else:
        url_hash = hashlib.sha1(url).hexdigest()

    domain = ".".join(urlparse.urlparse(url)[1].split('.')[-2:])

    conn = _connectDb()
    c = conn.cursor()

    # Create table if not exists
    _createDb(conn, c)

    # Insert a row of data
    if not post_data: post_data=""
    only_headers = 0
    if data == -1:
        only_headers = 1
        data = ""
    created = time.mktime(time.localtime())
    content_type = headers.get('content-type', '').split(';')[0].strip()
    if content_type in COMPRESS_TYPES:
        compressed = 1
        data = zlib.compress(data)
    else:
        compressed = 0
    data = sqlite3.Binary(data)
    t = (url_hash, domain, url, post_data, json.dumps(headers), created, data, only_headers, compressed)
    c.execute(u"""INSERT OR REPLACE INTO cache values (?, ?, ?, ?, ?, ?, ?, ?, ?)""", t)

    # Save (commit) the changes and clean up
    conn.commit()
    c.close()
    conn.close()

def migrate_to_db():
    import re
    import os
    import sqlite3
    import glob

    conn = _connectDb()
    c = conn.cursor()
    _createDb(conn, c)

    files = glob.glob(_getCacheBase() + "/*/*/*/*/*")
    _files = filter(lambda x: not x.endswith(".headers"), files)

    for f in _files:
        info = re.compile("%s/(.*?)/../../../(.*)" % _getCacheBase()).findall(f)
        domain = url = info[0][0]
        url_hash = info[0][1]
        post_data = ""
        created = os.stat(f).st_ctime
        fd = open(f, "r")
        data = fd.read()
        fd.close()
        fd = open(f + ".headers", "r")
        headers = fd.read()
        fd.close()
        t = (url_hash, domain, url, post_data, headers, created, sqlite3.Binary(data), 0)
        c.execute(u"""INSERT OR REPLACE INTO cache values (?, ?, ?, ?, ?, ?, ?, ?)""", t)

    conn.commit()
    c.close()
    conn.close()

def compress_db():
    conn = _connectDb()
    c = conn.cursor()
    _createDb(conn, c)
    c.execute(u"""SELECT url_hash FROM cache WHERE compressed = 0""")
    ids = [row[0] for row in c]
    for url_hash in ids:
        c.execute(u"""SELECT headers, data FROM cache WHERE url_hash = ?""", (url_hash, ))
        headers = {}
        for row in c:
            headers = json.loads(row[0])
            data = row[1]
        
        content_type = headers.get('content-type', '').split(';')[0].strip()
        if content_type in COMPRESS_TYPES:
            data = zlib.compress(data)
            t = (sqlite3.Binary(data), url_hash)
            print url_hash, 'update'
            c.execute('UPDATE cache SET compressed = 1, data = ? WHERE url_hash = ?', t)

    conn.commit()
    print "optimizing database"
    c.execute('VACUUM')
    conn.commit()
    c.close()
    conn.close()