bitcask.py

from inmemory_store import BitCaskKeyDirEntry, BitCaskKVPair
import os
from typing import Dict
import struct


class BitCaskDataStore:
    def __init__(self, datafile_name="database.db") -> None:
        self.datafile = datafile_name
        self.file_handle =  open(self.datafile, "a+b")
        self.keydir: Dict[str, BitCaskKeyDirEntry] = dict()
        self.build_keydir()
    
    def build_keydir(self):
        filesize = os.path.getsize(self.datafile)
        if filesize == 0:
            return 
        with open(self.datafile, 'rb') as f:
            while f.tell() < filesize:
                offset = f.tell()
                byte_data = f.read(16)
                timestamp, keysize, valuesize = struct.unpack("<QLL", byte_data)
                key = f.read(keysize).decode()
                value = f.read(valuesize).decode()
                self.keydir[key] = BitCaskKeyDirEntry(valuesize, offset)
            

    def get(self, key):
        """
        Get the offset information from the keydir
        Use the offset info to read the key and the value from the data file
        """

        if key not in self.keydir:
            return ""
        offset_info = self.keydir[key]
        file_offset = offset_info.offset
        with open(self.datafile, 'rb') as f:
            f.seek(file_offset)
            byte_data = f.read(16)
            timestamp, keysize, valuesize = struct.unpack("<QLL", byte_data)
            _ = f.read(keysize)
            value_bytes = f.read(valuesize)
        return value_bytes.decode()

    def put(self, key, value):
        """
        Insert a key-value pair into the bitcask database
        """
        
        # First, insert the key-value into the in-memory keydir
        valsize = len(value)

        offset = os.path.getsize(self.datafile)
        kdEntry = BitCaskKeyDirEntry(valsize, offset)
        self.keydir[key] = kdEntry

        # Append the key-value to the disk store now
        diskstore = BitCaskKVPair(key, value)
        byte_data = diskstore.encode()
        with open(self.datafile, "ab") as f:
            f.write(byte_data)

    def list_keys(self):
        """
        List all keys in the datastore
        """
        return list(self.keydir.keys())

    def sync(self):
        """
        Flush OS file buffers to disk to 
        persist pending writes to disk
        """
        self.file_handle.flush()
        os.fsync(self.file_handle)

    def close(self):
        """
        Flush pending writes to disk and close file handle
        """
        self.file_handle.flush()
        os.fsync(self.file_handle)
        self.file_handle.close()