mbtiles2compactcache.py

# -------------------------------------------------------------------------------
# Name:        mbtilesRaster2compactcache
# Purpose:     Build compact cache V2 bundles from single MBTiles Raster dataset file.
#
# Author:      ltbam, luci6974
#
# Created:     18/07/2022
# Modified:    -
#
#  Copyright 2022 swisstopo. 2016 ESRI
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.?
#
# -------------------------------------------------------------------------------
#
# Converts .mbtile file to the esri Compact Cache V2 format
#
# Takes two arguments, the first one is the input .mbfile folder
# the second one being the output cache folder (_alllayers)
#
#
# This script is intended to transform big Files, so it will loop each record
# and resolve is it has to be exported based on the maxlvl parameter.
#
# It does not check the input tile format, and assumes that the file
# is a valid sqlite tile databases.
#
# -------------------------------------------------------------------------------
#
# Changeset
# Version 1.0.0 ltbam

import argparse
import sqlite3
import os
import struct
import shutil
import datetime
import re
import io
import time
import math

from queue import Queue
from threading import Lock, get_ident, Thread
import multiprocessing


class Bundle:
    # Bundle linear size in tiles
    BSZ = 128
    # Tiles per bundle
    BSZ2 = BSZ ** 2
    # Index size in bytes
    IDXSZ = BSZ2 * 8
    # Max size
    M = 2 ** 40

    def __init__(self, file_name):
        self.file_name = file_name
        self.curr_max = 0
        self.curr_offset = 0
        self.curr_index = []
        self.lock = Lock()
        self.fd = None
        regex = r"L(..)"
        self.level = re.findall(regex, self.file_name)[0]
        fname = self.file_name.split('.bundle')[0].split('\\')[-1]
        s_val = fname[1:].split('C')
        self.row_offset = int(s_val[0], 16)
        self.col_offset = int(s_val[1], 16)
        if not os.path.exists(self.file_name):
            self.init_content()

    def init_content(self):
        # print("t {0}: initializing: {1}".format(get_ident(), self.file_name))
        self.fd = open(self.file_name, "wb")
        # Empty bundle file header, lots of magic numbers
        header = struct.pack("<4I3Q6I",
                             3,  # Version
                             Bundle.BSZ2,  # numRecords
                             0,  # maxRecord Size
                             5,  # Offset Size
                             0,  # Slack Space
                             64 + Bundle.IDXSZ,  # File Size
                             40,  # User Header Offset.
                             20 + Bundle.IDXSZ,  # User Header Size
                             3,  # Legacy 1
                             16,  # Legacy 2
                             Bundle.BSZ2,  # Legacy 3
                             5,  # Legacy 4
                             Bundle.IDXSZ  # Index Size
                             )
        self.fd.write(header)
        # Write empty index.
        self.fd.write(struct.pack("<{}Q".format(Bundle.BSZ2), *((0,) * Bundle.BSZ2)))
        self.fd.close()
        self.fd = None

    def open(self):
        # Open the bundle
        # wait a bit if a thread closed a File before opening it again
        self.fd = open(self.file_name, "r+b")
        # Read the current max record size
        self.fd.seek(8)
        self.curr_max = int(struct.unpack("<I", self.fd.read(4))[0])
        # Read the index as longs in a list
        self.fd.seek(64)
        self.curr_index = list(struct.unpack("<{}Q".format(Bundle.BSZ2),
                                             self.fd.read(Bundle.IDXSZ)))
        # Go to end
        self.fd.seek(0, os.SEEK_END)
        self.curr_offset = self.fd.tell()

    def write_tile(self, tile, tile_size, row, col):
        self.fd.write(struct.pack("<I", tile_size))
        self.fd.write(tile)
        self.curr_offset += 4
        # Update the index, row major
        self.curr_index[(row % Bundle.BSZ) * Bundle.BSZ + col % Bundle.BSZ] = self.curr_offset + (tile_size << 40)
        self.curr_offset += tile_size
        # Update the current bundle max tile size
        self.curr_max = max(self.curr_max, tile_size)

    def cleanup(self):
        """
        Updates header and closes the current bundle
        """
        # Update the max rec size and file size, then close the file
        self.fd.seek(8)
        self.fd.write(struct.pack("<I", self.curr_max))
        self.fd.seek(24)
        self.fd.write(struct.pack("<Q", self.curr_offset))
        self.fd.seek(64)
        self.fd.write(struct.pack("<{}Q".format(Bundle.BSZ2), *self.curr_index))
        self.fd.close()
        self.fd = None
        # print("t {0}: cleaned up: {1}".format(get_ident(), self.file_name))
        # time.sleep(0.2)


class BundleManager:
    b_list = {}

    def __init__(self):
        pass

    @staticmethod
    def fetch_worker(mb_tile_file, max_level, queue, num_workers):
        """
        Single thread that owns the DB cursor and feeds batches into the queue.
        Uses fetchmany() for efficient streaming without OFFSET scanning.
        """
        database = sqlite3.connect(mb_tile_file)
        row_cursor = database.cursor()

        if max_level != -1:
            row_cursor.execute(
                'SELECT zoom_level, tile_column, tile_row, tile_data '
                'FROM tiles WHERE zoom_level <= ?', (max_level,)
            )
        else:
            row_cursor.execute(
                'SELECT zoom_level, tile_column, tile_row, tile_data FROM tiles'
            )

        try:
            while True:
                rows = row_cursor.fetchmany(Application.rec_per_request)
                if not rows:
                    break
                queue.put(rows)  # blocks naturally if workers are slow (backpressure)
        except Exception as e:
            print(f"Fetcher error: {e}")
            raise
        finally:
            # Always send poison pills so workers are never left blocked on queue.get()
            for _ in range(num_workers):
                queue.put(None)
            row_cursor.close()
            database.close()

    @staticmethod
    def process(queue, arguments):
        """
        Worker thread: pulls batches from the queue and writes bundle files.
        Exits when it receives the None poison pill.
        """

        while True:
            batch = queue.get()
            if batch is None:  # poison pill — this worker is done
                break

            data = {}
            for rec in batch:
                level = 'L' + '{:02d}'.format(rec[0])
                output_path = os.path.join(arguments.destination, level)
                max_rows = 2 ** int(rec[0]) - 1
                tile = io.BytesIO(rec[3]).getvalue()
                row = max_rows - int(rec[2])
                col = int(rec[1])

                # Resolve the bundle
                start_row = (row // Bundle.BSZ) * Bundle.BSZ
                start_col = (col // Bundle.BSZ) * Bundle.BSZ
                bname = "R{:04x}C{:04x}".format(start_row, start_col)
                fname = os.path.join(output_path, bname + ".bundle")

                if fname not in data:
                    data[fname] = []
                data[fname].append([fname, tile, row, col])

            if data:
                BundleManager.add_tiles(data, arguments.lock)

    @staticmethod
    def add_tiles(data, lock):
        for bundle in data.keys():

            # init list one bundle/object at a time
            with lock:
                if bundle not in BundleManager.b_list:
                    BundleManager.b_list[bundle] = Bundle(bundle)

            # lock write operation on the bundle per thread
            # print("t {0}: enter lock: {1}".format(get_ident(), bundle))
            with BundleManager.b_list[bundle].lock:
                if not BundleManager.b_list[bundle].fd:
                    BundleManager.b_list[bundle].open()
                for t in data[bundle]:
                    tile = t[1]
                    tile_size = len(t[1])
                    row = t[2]
                    col = t[3]
                    # print("add tile row:{0} col:{1} buff:{2} path:{3}".format(row, col, tile_size, bundle))
                    # Open or create it, seek to end of bundle file
                    BundleManager.b_list[bundle].write_tile(tile, tile_size, row, col)

                BundleManager.b_list[bundle].cleanup()
                # free written bundle from mem.
                del BundleManager.b_list[bundle]

            # print("t {0}: left lock: {1}".format(get_ident(), bundle))

    @staticmethod
    def add_tile(output_path, byte_buffer, row, col=None):
        """
        Add this tile to the output cache

        :param output_path: path where the bundle is.
        :param byte_buffer: input tile as byte buffer
        :param row: row number
        :param col: column number
        """

        # Read the tile data
        tile = io.BytesIO(byte_buffer).getvalue()
        tile_size = len(tile)

        # resolve the bundle
        start_row = int((row / Bundle.BSZ)) * Bundle.BSZ
        start_col = int((col / Bundle.BSZ)) * Bundle.BSZ
        bname = "R{:04x}C{:04x}".format(start_row, start_col)
        fname = os.path.join(output_path, bname + ".bundle")

        print("add tile row:{0} col:{1} buff:{2} path:{3}".format(row, col, len(byte_buffer), output_path))

        # should be thread safe
        if not fname in BundleManager.b_list:
            BundleManager.b_list[fname] = Bundle(fname)

        # lock write operation on the bundle per thread
        with BundleManager.b_list[fname].lock:
            print("t {0}: enter lock: {1}".format(get_ident(), fname))
            # Open or create it, seek to end of bundle file
            if not BundleManager.b_list[fname].fd:
                BundleManager.b_list[fname].open()

            BundleManager.b_list[fname].write_tile(tile, tile_size, row, col)
            BundleManager.b_list[fname].cleanup()
            print("t {0}: left lock: {1}".format(get_ident(), fname))

        return fname


class Application:
    # Number of concurrent jobs for the export
    p_jobs = multiprocessing.cpu_count()
    # Records per fetchmany() call — used as the streaming chunk size
    rec_per_request = 500000
    # Queue depth: don't let the fetcher run too far ahead of workers
    queue_maxsize = p_jobs * 2

    def __init__(self):
        self.bm = BundleManager()

    @staticmethod
    def get_arguments():
        """
        Parses commandline arguments.

        :return: commandline arguments
        """
        parser = argparse.ArgumentParser()

        parser.add_argument('-s', '--source',
                            help='Input folder containing the mbtile files.', required=True)
        parser.add_argument('-d', '--destination',
                            help='Output for level folders.', required=True)
        parser.add_argument('-ml', '--max_level',
                            help='Do until this Level. If omitted, the maximum zoom level '
                                 'present in the MBTiles file is used automatically.',
                            default=-1, type=int, required=False)

        # Return the command line arguments.
        arguments = parser.parse_args()

        # validate folder parameters
        if not os.path.exists(arguments.source):
            parser.error("Input folder does not exist or is inaccessible.")

        return arguments


def get_max_level(mb_tile_file):
    """
    Resolve the maximum zoom level present in the MBTiles file.

    Strategy (fastest-first, safe for very large files):
      1. Check the 'metadata' table for the 'maxzoom' key — O(1) lookup,
         no tile rows touched at all.
      2. Fall back to SELECT MAX(zoom_level) FROM tiles — SQLite resolves
         this via the index on (zoom_level, tile_column, tile_row) without
         a full table scan.

    :param mb_tile_file: path to the .mbtiles SQLite file
    :return: integer max zoom level
    :raises RuntimeError: if the level cannot be determined
    """
    db = sqlite3.connect(mb_tile_file)
    try:
        # 1. Fast path: metadata table
        try:
            cur = db.execute(
                "SELECT value FROM metadata WHERE name = 'maxzoom' LIMIT 1"
            )
            row = cur.fetchone()
            if row is not None:
                return int(row[0])
        except sqlite3.DatabaseError:
            pass  # metadata table absent — fall through

        # 2. Index-assisted MAX() scan on the tiles table
        cur = db.execute("SELECT MAX(zoom_level) FROM tiles")
        row = cur.fetchone()
        if row is not None and row[0] is not None:
            return int(row[0])

        raise RuntimeError(
            "Cannot determine max zoom level: 'tiles' table appears to be empty."
        )
    finally:
        db.close()


#
# Entry point, start the Application
#
def main():
    app = Application()
    arguments = app.get_arguments()
    arguments.lock = Lock()

    # parse parameters
    mb_tile_file = arguments.source
    cache_output_folder = arguments.destination
    max_level_param = arguments.max_level
    print('Input file: {0}'.format(os.path.basename(mb_tile_file)))

    # Resolve max zoom level when not provided by the user
    if max_level_param == -1:
        print('--max_level not specified, querying MBTiles file…')
        max_level_param = get_max_level(mb_tile_file)
        print('Max zoom level detected: {0}'.format(max_level_param))

    # prepare output template
    if os.path.exists(cache_output_folder):
        shutil.rmtree(cache_output_folder)

    # creating lvl directories
    for lvl in range(max_level_param + 1):
        level = 'L' + '{:02d}'.format(lvl)
        dir = os.path.join(cache_output_folder, level)
        if not os.path.exists(dir):
            os.makedirs(dir)

    start_time = datetime.datetime.now()

    # One fetcher feeds all workers via a bounded queue (backpressure built-in)
    queue = Queue(maxsize=Application.queue_maxsize)

    fetcher = Thread(
        target=BundleManager.fetch_worker,
        args=(mb_tile_file, max_level_param, queue, app.p_jobs)
    )
    workers = [
        Thread(target=BundleManager.process, args=(queue, arguments))
        for _ in range(app.p_jobs)
    ]

    print('Exporting {0} rows at a time within {1} threads.'.format(
        app.rec_per_request, app.p_jobs))

    fetcher.start()
    for w in workers:
        w.start()

    fetcher.join()
    for w in workers:
        w.join()

    elapsed = (datetime.datetime.now() - start_time).total_seconds()
    print('Done in {:.2f}s'.format(elapsed))


if __name__ == '__main__':
    main()