# sqlitestore.py - Storage backend that uses SQLite

#

# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

"""store repository data in SQLite (EXPERIMENTAL)

The sqlitestore extension enables the storage of repository data in SQLite.

This extension is HIGHLY EXPERIMENTAL. There are NO BACKWARDS COMPATIBILITY

GUARANTEES. This means that repositories created with this extension may

only be usable with the exact version of this extension/Mercurial that was

used. The extension attempts to enforce this in order to prevent repository

corruption.

In addition, several features are not yet supported or have known bugs:

* Only some data is stored in SQLite. Changeset, manifest, and other repository

  data is not yet stored in SQLite.

* Transactions are not robust. If the process is aborted at the right time

  during transaction close/rollback, the repository could be in an inconsistent

  state. This problem will diminish once all repository data is tracked by

  SQLite.

* Bundle repositories do not work (the ability to use e.g.

  `hg -R <bundle-file> log` to automatically overlay a bundle on top of the

  existing repository).

* Various other features don't work.

This extension should work for basic clone/pull, update, and commit workflows.

Some history rewriting operations may fail due to lack of support for bundle

repositories.

To use, activate the extension and set the ``storage.new-repo-backend`` config

option to ``sqlite`` to enable new repositories to use SQLite for storage.

"""

# To run the test suite with repos using SQLite by default, execute the

# following:

#

# HGREPOFEATURES="sqlitestore" run-tests.py \

#     --extra-config-opt extensions.sqlitestore= \

#     --extra-config-opt storage.new-repo-backend=sqlite

from __future__ import absolute_import

import hashlib

import sqlite3

import struct

import threading

import zlib

from mercurial.i18n import _

from mercurial.node import (

    nullid,

    nullrev,

    short,

)

from mercurial.thirdparty import attr

from mercurial import (

    ancestor,

    dagop,

    encoding,

    error,

    extensions,

    localrepo,

    mdiff,

    pycompat,

    registrar,

    util,

    verify,

)

from mercurial.interfaces import (

    repository,

    util as interfaceutil,

)

from mercurial.utils import storageutil

try:

    from mercurial import zstd

    zstd.__version__

except ImportError:

    zstd = None

configtable = {}

configitem = registrar.configitem(configtable)

# experimental config: storage.sqlite.compression

configitem(

    b'storage',

    b'sqlite.compression',

    default=b'zstd' if zstd else b'zlib',

    experimental=True,

)

# Note for extension authors: ONLY specify testedwith = 'ships-with-hg-core' for

# extensions which SHIP WITH MERCURIAL. Non-mainline extensions should

# be specifying the version(s) of Mercurial they are tested with, or

# leave the attribute unspecified.

testedwith = b'ships-with-hg-core'

REQUIREMENT = b'exp-sqlite-001'

REQUIREMENT_ZSTD = b'exp-sqlite-comp-001=zstd'

REQUIREMENT_ZLIB = b'exp-sqlite-comp-001=zlib'

REQUIREMENT_NONE = b'exp-sqlite-comp-001=none'

REQUIREMENT_SHALLOW_FILES = b'exp-sqlite-shallow-files'

CURRENT_SCHEMA_VERSION = 1

COMPRESSION_NONE = 1

COMPRESSION_ZSTD = 2

COMPRESSION_ZLIB = 3

FLAG_CENSORED = 1

FLAG_MISSING_P1 = 2

FLAG_MISSING_P2 = 4

CREATE_SCHEMA = [

    # Deltas are stored as content-indexed blobs.

    # compression column holds COMPRESSION_* constant for how the

    # delta is encoded.

    'CREATE TABLE delta ('

    '    id INTEGER PRIMARY KEY, '

    '    compression INTEGER NOT NULL, '

    '    hash BLOB UNIQUE ON CONFLICT ABORT, '

    '    delta BLOB NOT NULL '

    ')',

    # Tracked paths are denormalized to integers to avoid redundant

    # storage of the path name.

    'CREATE TABLE filepath ('

    '    id INTEGER PRIMARY KEY, '

    '    path BLOB NOT NULL '

    ')',

    'CREATE UNIQUE INDEX filepath_path ON filepath (path)',

    # We have a single table for all file revision data.

    # Each file revision is uniquely described by a (path, rev) and

    # (path, node).

    #

    # Revision data is stored as a pointer to the delta producing this

    # revision and the file revision whose delta should be applied before

    # that one. One can reconstruct the delta chain by recursively following

    # the delta base revision pointers until one encounters NULL.

    #

    # flags column holds bitwise integer flags controlling storage options.

    # These flags are defined by the FLAG_* constants.

    'CREATE TABLE fileindex ('

    '    id INTEGER PRIMARY KEY, '

    '    pathid INTEGER REFERENCES filepath(id), '

    '    revnum INTEGER NOT NULL, '

    '    p1rev INTEGER NOT NULL, '

    '    p2rev INTEGER NOT NULL, '

    '    linkrev INTEGER NOT NULL, '

    '    flags INTEGER NOT NULL, '

    '    deltaid INTEGER REFERENCES delta(id), '

    '    deltabaseid INTEGER REFERENCES fileindex(id), '

    '    node BLOB NOT NULL '

    ')',

    'CREATE UNIQUE INDEX fileindex_pathrevnum '

    '    ON fileindex (pathid, revnum)',

    'CREATE UNIQUE INDEX fileindex_pathnode ON fileindex (pathid, node)',

    # Provide a view over all file data for convenience.

    'CREATE VIEW filedata AS '

    'SELECT '

    '    fileindex.id AS id, '

    '    filepath.id AS pathid, '

    '    filepath.path AS path, '

    '    fileindex.revnum AS revnum, '

    '    fileindex.node AS node, '

    '    fileindex.p1rev AS p1rev, '

    '    fileindex.p2rev AS p2rev, '

    '    fileindex.linkrev AS linkrev, '

    '    fileindex.flags AS flags, '

    '    fileindex.deltaid AS deltaid, '

    '    fileindex.deltabaseid AS deltabaseid '

    'FROM filepath, fileindex '

    'WHERE fileindex.pathid=filepath.id',

    'PRAGMA user_version=%d' % CURRENT_SCHEMA_VERSION,

]

def resolvedeltachain(db, pathid, node, revisioncache, stoprids, zstddctx=None):

    """Resolve a delta chain for a file node."""

    # TODO the "not in ({stops})" here is possibly slowing down the query

    # because it needs to perform the lookup on every recursive invocation.

    # This could possibly be faster if we created a temporary query with

    # baseid "poisoned" to null and limited the recursive filter to

    # "is not null".

    res = db.execute(

        'WITH RECURSIVE '

        '    deltachain(deltaid, baseid) AS ('

        '        SELECT deltaid, deltabaseid FROM fileindex '

        '            WHERE pathid=? AND node=? '

        '        UNION ALL '

        '        SELECT fileindex.deltaid, deltabaseid '

        '            FROM fileindex, deltachain '

        '            WHERE '

        '                fileindex.id=deltachain.baseid '

        '                AND deltachain.baseid IS NOT NULL '

        '                AND fileindex.id NOT IN ({stops}) '

        '    ) '

        'SELECT deltachain.baseid, compression, delta '

        'FROM deltachain, delta '

        'WHERE delta.id=deltachain.deltaid'.format(

            stops=','.join(['?'] * len(stoprids))

        ),

        tuple([pathid, node] + list(stoprids.keys())),

    )

    deltas = []

    lastdeltabaseid = None

    for deltabaseid, compression, delta in res:

        lastdeltabaseid = deltabaseid

        if compression == COMPRESSION_ZSTD:

            delta = zstddctx.decompress(delta)

        elif compression == COMPRESSION_NONE:

            delta = delta

        elif compression == COMPRESSION_ZLIB:

            delta = zlib.decompress(delta)

        else:

            raise SQLiteStoreError(

                b'unhandled compression type: %d' % compression

            )

        deltas.append(delta)

    if lastdeltabaseid in stoprids:

        basetext = revisioncache[stoprids[lastdeltabaseid]]

    else:

        basetext = deltas.pop()

    deltas.reverse()

    fulltext = mdiff.patches(basetext, deltas)

    # SQLite returns buffer instances for blob columns on Python 2. This

    # type can propagate through the delta application layer. Because

    # downstream callers assume revisions are bytes, cast as needed.

    if not isinstance(fulltext, bytes):

        fulltext = bytes(delta)

    return fulltext

def insertdelta(db, compression, hash, delta):

    try:

        return db.execute(

            'INSERT INTO delta (compression, hash, delta) VALUES (?, ?, ?)',

            (compression, hash, delta),

        ).lastrowid

    except sqlite3.IntegrityError:

        return db.execute(

            'SELECT id FROM delta WHERE hash=?', (hash,)

        ).fetchone()[0]

class SQLiteStoreError(error.StorageError):

    pass

@attr.s

class revisionentry(object):

    rid = attr.ib()

    rev = attr.ib()

    node = attr.ib()

    p1rev = attr.ib()

    p2rev = attr.ib()

    p1node = attr.ib()

    p2node = attr.ib()

    linkrev = attr.ib()

    flags = attr.ib()

@interfaceutil.implementer(repository.irevisiondelta)

@attr.s(slots=True)

class sqliterevisiondelta(object):

    node = attr.ib()

    p1node = attr.ib()

    p2node = attr.ib()

    basenode = attr.ib()

    flags = attr.ib()

    baserevisionsize = attr.ib()

    revision = attr.ib()

    delta = attr.ib()

    linknode = attr.ib(default=None)

@interfaceutil.implementer(repository.iverifyproblem)

@attr.s(frozen=True)

class sqliteproblem(object):

    warning = attr.ib(default=None)

    error = attr.ib(default=None)

    node = attr.ib(default=None)

@interfaceutil.implementer(repository.ifilestorage)

class sqlitefilestore(object):

    """Implements storage for an individual tracked path."""

    def __init__(self, db, path, compression):

        self._db = db

        self._path = path

        self._pathid = None

        # revnum -> node

        self._revtonode = {}

        # node -> revnum

        self._nodetorev = {}

        # node -> data structure

        self._revisions = {}

        self._revisioncache = util.lrucachedict(10)

        self._compengine = compression

        if compression == b'zstd':

            self._cctx = zstd.ZstdCompressor(level=3)

            self._dctx = zstd.ZstdDecompressor()

        else:

            self._cctx = None

            self._dctx = None

        self._refreshindex()

    def _refreshindex(self):

        self._revtonode = {}

        self._nodetorev = {}

        self._revisions = {}

        res = list(

            self._db.execute(

                'SELECT id FROM filepath WHERE path=?', (self._path,)

            )

        )

        if not res:

            self._pathid = None

            return

        self._pathid = res[0][0]

        res = self._db.execute(

            'SELECT id, revnum, node, p1rev, p2rev, linkrev, flags '

            'FROM fileindex '

            'WHERE pathid=? '

            'ORDER BY revnum ASC',

            (self._pathid,),

        )

        for i, row in enumerate(res):

            rid, rev, node, p1rev, p2rev, linkrev, flags = row

            if i != rev:

                raise SQLiteStoreError(

                    _(b'sqlite database has inconsistent revision numbers')

                )

            if p1rev == nullrev:

                p1node = nullid

            else:

                p1node = self._revtonode[p1rev]

            if p2rev == nullrev:

                p2node = nullid

            else:

                p2node = self._revtonode[p2rev]

            entry = revisionentry(

                rid=rid,

                rev=rev,

                node=node,

                p1rev=p1rev,

                p2rev=p2rev,

                p1node=p1node,

                p2node=p2node,

                linkrev=linkrev,

                flags=flags,

            )

            self._revtonode[rev] = node

            self._nodetorev[node] = rev

            self._revisions[node] = entry

    # Start of ifileindex interface.

    def __len__(self):

        return len(self._revisions)

    def __iter__(self):

        return iter(pycompat.xrange(len(self._revisions)))

    def hasnode(self, node):

        if node == nullid:

            return False

        return node in self._nodetorev

    def revs(self, start=0, stop=None):

        return storageutil.iterrevs(

            len(self._revisions), start=start, stop=stop

        )

    def parents(self, node):

        if node == nullid:

            return nullid, nullid

        if node not in self._revisions:

            raise error.LookupError(node, self._path, _(b'no node'))

        entry = self._revisions[node]

        return entry.p1node, entry.p2node

    def parentrevs(self, rev):

        if rev == nullrev:

            return nullrev, nullrev

        if rev not in self._revtonode:

            raise IndexError(rev)

        entry = self._revisions[self._revtonode[rev]]

        return entry.p1rev, entry.p2rev

    def rev(self, node):

        if node == nullid:

            return nullrev

        if node not in self._nodetorev: