stratum-db 1.0.0__py3-none-any.whl

stratum/api.py

@@ -0,0 +1,41 @@
+import os
+import logging
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+def _get_config():
+    data_dir = Path(os.environ.get("STRATUM_DATA_DIR","./data"))
+
+    table_dir_raw = os.environ.get("STRATUM_TABLE_DIR")
+    table_dir = Path(table_dir_raw) if table_dir_raw else None
+
+    port_raw = os.environ.get("STRATUM_PORT","50051")
+
+    try:
+        port = int(port_raw)
+    except ValueError:
+        raise ValueError(
+            f"STRATUM_PORT must be an integer, got {port_raw!r}"
+        )
+
+    log_level_name = os.environ.get("STRATUM_LOG_LEVEL","INFO").upper()
+    log_level = getattr(logging, log_level_name)
+    if not isinstance(log_level, int):
+        raise ValueError(
+            f"STRATUM_LOG_LEVEL must be a valid level name, got {log_level_name!r}"
+        )
+
+    return data_dir, table_dir, port, log_level
+
+def main():
+    data_dir, table_dir, port, log_level = _get_config()
+    logging.basicConfig(level=log_level)
+
+    from stratum.grpc.server import serve
+    serve(data_dir = data_dir, table_dir=table_dir, port=port)
+
+if __name__ == "__main__":
+    main()

stratum/db.py

@@ -0,0 +1,90 @@
+import os
+
+
+_metadata_enabled = os.environ.get("STRATUM_METADATA_ENABLED", "false").strip().lower()
+if _metadata_enabled not in {"1", "true", "yes", "on", "0", "false", "no", "off"}:
+    raise ValueError(
+        "STRATUM_METADATA_ENABLED must be a boolean value (true/false, 1/0, yes/no, or on/off)"
+    )
+METADATA_ENABLED = _metadata_enabled in {"1", "true", "yes", "on"}
+
+
+if METADATA_ENABLED:
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import DeclarativeBase, sessionmaker
+
+    class Base(DeclarativeBase):
+        pass
+else:
+    Base = None
+
+
+def get_db_engine():
+    if METADATA_ENABLED:
+        database_url = os.environ.get("DATABASE_URL")
+        if not database_url:
+            raise RuntimeError(
+                "DATABASE_URL must be set when STRATUM_METADATA_ENABLED is enabled"
+            )
+        return create_engine(database_url, echo=True)
+    return None
+
+def get_session_factory():
+    db_engine = get_db_engine()
+    if db_engine is not None:
+        return sessionmaker(bind=db_engine)
+    return None
+
+def record_flush(filename, min_key, max_key, entry_count, file_size_bytes, created_at):
+    session_factory = get_session_factory()
+    if session_factory is None:
+        return None
+    session = session_factory()
+    from stratum.models import SSTableMeta
+    SSTableMetaEntry = SSTableMeta(
+        filename = filename,
+        min_key = min_key,
+        max_key = max_key,
+        entry_count = entry_count,
+        file_size_bytes = file_size_bytes,
+        created_at = created_at
+        )
+    try:
+        session.add(SSTableMetaEntry)
+        session.commit()
+    except Exception as e:
+        session.rollback()
+        print(f'[Metadata] record_flush failed: {e}')
+    finally:
+        session.close()
+
+def record_compaction(status, started_at, completed_at, tombstones_dropped, input_filenames, output_fields):
+    session_factory = get_session_factory()
+    if session_factory is None:
+        return None
+    session = session_factory()
+    from stratum.models import SSTableMeta, CompactionJob
+    try:
+        input_rows = (
+            session.query(SSTableMeta)
+            .filter(SSTableMeta.filename.in_(input_filenames))
+            .all()
+        )
+        for row in input_rows:
+            row.is_active = False
+
+        output_row = SSTableMeta(**output_fields)
+        compactionJobEntry = CompactionJob(
+                status = status,    
+                started_at = started_at,
+                completed_at = completed_at,
+                tombstones_dropped = tombstones_dropped,
+                output_sstable = output_row,
+                input_sstables = input_rows,)
+        session.add(compactionJobEntry)
+        session.commit()
+    except Exception as e:
+        session.rollback()
+        print(f'Metadata Write failed: {e}')
+    finally:
+        session.close()

stratum/engine.py

@@ -0,0 +1,220 @@
+import re
+from threading import RLock
+from pathlib import Path
+from sortedcontainers import SortedDict
+from datetime import datetime, timezone
+import heapq
+from itertools import groupby
+
+from stratum.memtable import MemTable
+from stratum.wal import WAL
+from stratum.sstable import SSTable
+from stratum import db
+from stratum.metrics import MetricsRegistry
+
+class Engine:
+    def __init__(self, data_dir, table_dir=None):
+        if not isinstance(data_dir, Path):
+            raise TypeError(
+                f"WAL requires a pathlib.Path, got {type(data_dir).__name__}. "
+                f"Wrap it with Path(...) at the call site."
+            )
+        if table_dir is None:
+            table_dir = data_dir
+        if not isinstance(table_dir, Path):
+            raise TypeError(
+                f"SSTable requires a pathlib.Path, got {type(table_dir).__name__}. "
+                f"Wrap it with Path(...) at the call site."
+            )
+
+        self.data_dir = data_dir
+        self.table_dir = table_dir
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        self.table_dir.mkdir(parents=True, exist_ok=True)
+
+        self.memtable = MemTable()
+        self.wal = WAL(self.data_dir)
+        self._seq_no = 0
+        self.lock = RLock()
+        self.sstables = []
+        self.tableCount = 0
+        self.metrics = MetricsRegistry()
+
+        self._load_sstables()
+        self._recover()
+        self._update_metrics()
+
+    def _update_metrics(self):
+        self.metrics.set_gauge("sstables", len(self.sstables))
+        self.metrics.set_gauge("memtable_entries", len(self.memtable.items()))
+
+    def _load_sstables(self):
+        for sst_path in sorted(self.table_dir.glob("*.sst")):
+            idx_path = sst_path.with_suffix('.idx')
+            if not idx_path.exists():
+                raise FileNotFoundError(f"Could not find {idx_path}")
+            self.sstables.append(SSTable.from_file(sst_path, idx_path))
+            match = re.match(r"^(\d+)", sst_path.name)
+            if match:
+                self.tableCount = max(self.tableCount, int(match.group(1)) + 1)
+
+    def _recover(self):
+        for seq_no, op_type, key, value in self.wal.replay():
+            self._seq_no = seq_no
+            if op_type == 1:
+                self.memtable.put(key, value, self._seq_no)
+            else:
+                self.memtable.delete(key, self._seq_no)
+
+    def put(self, key, value):
+        self.metrics.increment("puts_total")
+        with self.lock:
+            if self.memtable.is_full():
+                self.write_table()
+            self._seq_no += 1
+            self.wal.append(self._seq_no, 1, key, value)
+            self.memtable.put(key, value, self._seq_no)
+            self._update_metrics()
+
+    def write_table(self):
+        data = self.memtable.items()
+        self.tableCount += 1        
+        new_sstable = SSTable.flush(self.table_dir, self.tableCount, data)
+        if db.METADATA_ENABLED:
+            created_at = datetime.now(timezone.utc)
+            db.record_flush(
+                filename=new_sstable.sst_path.name,
+                min_key=new_sstable.min_key,
+                max_key=new_sstable.max_key,
+                entry_count=new_sstable.entry_count,
+                file_size_bytes=new_sstable.file_size_bytes,
+                created_at=created_at
+                )
+        self.sstables.append(new_sstable)
+        self.wal.truncate()
+        self.memtable = MemTable()
+        self.metrics.increment("flushes_total")
+        self._update_metrics()
+
+    def delete(self, key):
+        self.metrics.increment("deletes_total")
+        with self.lock:
+            self._seq_no += 1
+            self.wal.append(self._seq_no, 2, key, b"")
+            self.memtable.delete(key,self._seq_no)
+            self._update_metrics()
+
+    def get(self, key):
+        self.metrics.increment("gets_total")
+        with self.lock:
+            record = self.memtable.get_record(key)
+            if record is not None:
+                value, _, deleted = record
+                return None if deleted else value
+
+            for sstable in reversed(self.sstables):
+                if sstable.min_key is not None and not (sstable.min_key <= key <= sstable.max_key):
+                    continue
+                if not sstable.might_contain(key):
+                    continue
+                results = sstable.search(key)
+                if results:
+                    _, value, _, deleted = results[0]
+                    return None if deleted else value
+
+            return None
+
+    def scan(self, start_key, end_key):
+        self.metrics.increment("scans_total")
+        if start_key > end_key:
+            return
+        
+        with self.lock:
+            sstables = list(self.sstables)
+            memtable_items = self.memtable.items()
+
+        sources = []
+
+        for table in sstables:
+            if table.entry_count == 0:
+                continue
+
+            if table.max_key < start_key or table.min_key > end_key:
+                continue
+
+            sources.append(table.scan(start_key, end_key))
+
+        sources.append(
+            (key, value, seq_no, deleted) for key, (value, seq_no, deleted) in memtable_items if start_key <= key <= end_key
+        )
+
+        merged = heapq.merge(*sources, key = lambda entry: entry [0])
+
+        for key, entries in groupby(merged, key= lambda entry: entry[0]):
+            newest = max(entries, key = lambda entry: entry[2])
+            if not newest[3]:
+                yield key, newest[1]
+
+    def compact(self):
+        self.metrics.increment("compactions_total")
+        with self.lock:
+            started_at = datetime.now(timezone.utc)
+            entries = SortedDict()
+            input_filenames = []
+            for table in self.sstables:
+                input_filenames.append(table.sst_path.name)
+                table_entries = table.scan(table.min_key, table.max_key)
+                for key, val, seq_no, deleted in table_entries:
+                    old = entries.get(key)
+                    if old is None:
+                        entries[key] = (val, seq_no, deleted)
+                    else:
+                        if seq_no > entries[key][1]:
+                            entries[key] = (val, seq_no, deleted)
+            del_keys = []
+            for key, entry in entries.items():
+                if entry[2]:
+                    del_keys.append(key)
+            for key in del_keys:
+                del entries[key]
+
+            self.tableCount += 1
+            created_at = datetime.now(timezone.utc)
+            new_table = SSTable.flush(self.table_dir, self.tableCount, entries.items())
+
+            
+
+            if new_table.min_key is None:
+                new_table.min_key = b""
+                new_table.max_key = b""
+            
+            old_sst_paths = [t.sst_path for t in self.sstables]
+            old_idx_paths = [t.idx_path for t in self.sstables]
+            self.sstables = [new_table]
+            output_fields = {
+                'filename':new_table.sst_path.name,
+                'min_key':new_table.min_key,
+                'max_key':new_table.max_key,
+                'entry_count':new_table.entry_count,
+                'file_size_bytes':new_table.file_size_bytes,
+                'created_at':created_at
+            }
+            if db.METADATA_ENABLED:
+                completed_at = datetime.now(timezone.utc)
+                db.record_compaction(
+                status="Completed",
+                started_at=started_at,
+                completed_at=completed_at,
+                tombstones_dropped=len(del_keys),
+                input_filenames=input_filenames,
+                output_fields=output_fields
+                )
+            for sst_path, idx_path in zip(old_sst_paths,old_idx_paths):
+                try:
+                    Path(sst_path).unlink()
+                except OSError as e:
+                    print(f"Failed to delete {sst_path}: {e}")
+                try:
+                    Path(idx_path).unlink()
+                except OSError as e:
+                    print(f"Failed to delete {idx_path}: {e}")

stratum/grpc/server.py

@@ -0,0 +1,45 @@
+from concurrent import futures
+from pathlib import Path
+import grpc
+
+from stratum.grpc import stratum_pb2,stratum_pb2_grpc
+from stratum.engine import Engine
+
+class StratumServicer(stratum_pb2_grpc.StratumServicer):
+    def __init__(self, engine):
+        self.engine = engine
+
+    @staticmethod
+    def _require_key(key, context, field_name):
+        if not key:
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT,
+                f"{field_name} must not be empty",
+            )
+
+    def Get(self, request, context):
+        self._require_key(request.key, context, "key")
+        value = self.engine.get(request.key)
+        if value is None:
+            return stratum_pb2.GetResponse(found = False)
+        return stratum_pb2.GetResponse(value=value, found=True)
+
+    def Put(self, request, context):
+        self._require_key(request.key, context, "key")
+        self.engine.put(request.key, request.value)
+        return stratum_pb2.PutResponse(success=True)
+
+    def Scan(self, request, context):
+        self._require_key(request.start_key, context, "start_key")
+        self._require_key(request.end_key, context, "end_key")
+        for key, value in self.engine.scan(request.start_key, request.end_key):
+            yield stratum_pb2.KVPair(key=key, value=value)
+
+def serve(data_dir, table_dir=None, port=50051):
+    engine = Engine(data_dir = data_dir, table_dir= table_dir)
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
+    stratum_pb2_grpc.add_StratumServicer_to_server(StratumServicer(engine), server)
+    server.add_insecure_port(f'[::]:{port}')
+    server.start()
+    print(f"gRPC server running on port {port}")
+    server.wait_for_termination()

stratum/grpc/stratum_pb2.py

@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: stratum.proto
+# Protobuf Python Version: 7.35.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    7,
+    35,
+    1,
+    '',
+    'stratum.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rstratum.proto\x12\x07stratum\"\x19\n\nGetRequest\x12\x0b\n\x03key\x18\x01 \x01(\x0c\"+\n\x0bGetResponse\x12\r\n\x05value\x18\x01 \x01(\x0c\x12\r\n\x05\x66ound\x18\x02 \x01(\x08\"(\n\nPutRequest\x12\x0b\n\x03key\x18\x01 \x01(\x0c\x12\r\n\x05value\x18\x02 \x01(\x0c\"\x1e\n\x0bPutResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\"1\n\x0bScanRequest\x12\x11\n\tstart_key\x18\x01 \x01(\x0c\x12\x0f\n\x07\x65nd_key\x18\x02 \x01(\x0c\"$\n\x06KVPair\x12\x0b\n\x03key\x18\x01 \x01(\x0c\x12\r\n\x05value\x18\x02 \x01(\x0c\x32\x9e\x01\n\x07Stratum\x12\x30\n\x03Get\x12\x13.stratum.GetRequest\x1a\x14.stratum.GetResponse\x12\x30\n\x03Put\x12\x13.stratum.PutRequest\x1a\x14.stratum.PutResponse\x12/\n\x04Scan\x12\x14.stratum.ScanRequest\x1a\x0f.stratum.KVPair0\x01\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'stratum_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_GETREQUEST']._serialized_start=26
+  _globals['_GETREQUEST']._serialized_end=51
+  _globals['_GETRESPONSE']._serialized_start=53
+  _globals['_GETRESPONSE']._serialized_end=96
+  _globals['_PUTREQUEST']._serialized_start=98
+  _globals['_PUTREQUEST']._serialized_end=138
+  _globals['_PUTRESPONSE']._serialized_start=140
+  _globals['_PUTRESPONSE']._serialized_end=170
+  _globals['_SCANREQUEST']._serialized_start=172
+  _globals['_SCANREQUEST']._serialized_end=221
+  _globals['_KVPAIR']._serialized_start=223
+  _globals['_KVPAIR']._serialized_end=259
+  _globals['_STRATUM']._serialized_start=262
+  _globals['_STRATUM']._serialized_end=420
+# @@protoc_insertion_point(module_scope)

stratum/grpc/stratum_pb2_grpc.py

@@ -0,0 +1,183 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from . import stratum_pb2 as stratum__pb2
+
+GRPC_GENERATED_VERSION = '1.83.1'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + ' but the generated code in stratum_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class StratumStub:
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Get = channel.unary_unary(
+                '/stratum.Stratum/Get',
+                request_serializer=stratum__pb2.GetRequest.SerializeToString,
+                response_deserializer=stratum__pb2.GetResponse.FromString,
+                _registered_method=True)
+        self.Put = channel.unary_unary(
+                '/stratum.Stratum/Put',
+                request_serializer=stratum__pb2.PutRequest.SerializeToString,
+                response_deserializer=stratum__pb2.PutResponse.FromString,
+                _registered_method=True)
+        self.Scan = channel.unary_stream(
+                '/stratum.Stratum/Scan',
+                request_serializer=stratum__pb2.ScanRequest.SerializeToString,
+                response_deserializer=stratum__pb2.KVPair.FromString,
+                _registered_method=True)
+
+
+class StratumServicer:
+    """Missing associated documentation comment in .proto file."""
+
+    def Get(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Put(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Scan(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_StratumServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'Get': grpc.unary_unary_rpc_method_handler(
+                    servicer.Get,
+                    request_deserializer=stratum__pb2.GetRequest.FromString,
+                    response_serializer=stratum__pb2.GetResponse.SerializeToString,
+            ),
+            'Put': grpc.unary_unary_rpc_method_handler(
+                    servicer.Put,
+                    request_deserializer=stratum__pb2.PutRequest.FromString,
+                    response_serializer=stratum__pb2.PutResponse.SerializeToString,
+            ),
+            'Scan': grpc.unary_stream_rpc_method_handler(
+                    servicer.Scan,
+                    request_deserializer=stratum__pb2.ScanRequest.FromString,
+                    response_serializer=stratum__pb2.KVPair.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'stratum.Stratum', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('stratum.Stratum', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class Stratum:
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Get(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/stratum.Stratum/Get',
+            stratum__pb2.GetRequest.SerializeToString,
+            stratum__pb2.GetResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Put(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/stratum.Stratum/Put',
+            stratum__pb2.PutRequest.SerializeToString,
+            stratum__pb2.PutResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Scan(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_stream(
+            request,
+            target,
+            '/stratum.Stratum/Scan',
+            stratum__pb2.ScanRequest.SerializeToString,
+            stratum__pb2.KVPair.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

stratum/init_db.py

@@ -0,0 +1,6 @@
+from stratum.db import Base, get_session_factory, get_db_engine
+from stratum.models import SSTableMeta, CompactionJob
+
+database_engine = get_db_engine()
+Base.metadata.create_all(database_engine)
+print("Database initialized")