implement K8sDatabaseManager

Author: andrii-i

src/jupyter_scheduler_k8s/__init__.py

@@ -1,9 +1,7 @@
 """Kubernetes backend for jupyter-scheduler."""
 
-# Import k8s_orm FIRST to auto-install K8s database backend before anything else
-from . import k8s_orm
-
 from .executors import K8sExecutionManager
+from .database_manager import K8sDatabaseManager
 
 __version__ = "0.1.0"
-__all__ = ["K8sExecutionManager"]
+__all__ = ["K8sExecutionManager", "K8sDatabaseManager"]

src/jupyter_scheduler_k8s/database_manager.py

@@ -0,0 +1,42 @@
+from kubernetes import client, config
+from jupyter_scheduler.managers import DatabaseManager
+
+from .k8s_orm import K8sSession
+
+
+class K8sDatabaseManager(DatabaseManager):
+    """Database manager that uses Kubernetes Jobs for storage."""
+    
+    def create_session(self, db_url: str):
+        """Create K8s session factory."""
+        if not db_url.startswith("k8s://"):
+            raise ValueError(f"K8sDatabaseManager only supports k8s:// URLs, got: {db_url}")
+            
+        namespace = db_url[6:] or "default"
+        
+        def session_factory():
+            return K8sSession(namespace=namespace)
+        return session_factory
+    
+    def create_tables(self, db_url: str, drop_tables: bool = False):
+        """Ensure K8s namespace exists."""
+        if not db_url.startswith("k8s://"):
+            return
+            
+        namespace = db_url[6:] or "default"
+        
+        try:
+            config.load_incluster_config()
+        except config.ConfigException:
+            config.load_kube_config()
+        
+        v1 = client.CoreV1Api()
+        
+        try:
+            v1.read_namespace(name=namespace)
+        except client.ApiException as e:
+            if e.status == 404:
+                namespace_body = client.V1Namespace(
+                    metadata=client.V1ObjectMeta(name=namespace)
+                )
+                v1.create_namespace(body=namespace_body)

src/jupyter_scheduler_k8s/k8s_orm.py

@@ -2,12 +2,10 @@
 
 import json
 import logging
-from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict
 from kubernetes import client, config
 from kubernetes.client.rest import ApiException
 
-from jupyter_scheduler.models import Status
 from jupyter_scheduler.utils import get_utc_timestamp
 
 logger = logging.getLogger(__name__)
@@ -42,7 +40,7 @@ def _init_k8s_client(self):
             self.k8s_batch = client.BatchV1Api()
             self.k8s_core = client.CoreV1Api()
 
-            # Test connection
+            # Validate connectivity before proceeding
             self.k8s_core.get_api_versions()
         except Exception as e:
             logger.error(f"Failed to initialize K8s clients: {e}")
@@ -58,15 +56,15 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self.rollback()
 
     def query(self, model_class):
-        """Return K8s query object."""
+        """Create query for model class."""
         return K8sQuery(self, model_class)
 
     def add(self, job):
-        """Buffer job creation for batch commit."""
+        """Buffer job for batch commit."""
         self._pending_operations.append(('create', job))
 
     def commit(self):
-        """Execute all buffered operations."""
+        """Execute buffered operations."""
         if not self._pending_operations:
             return
 
@@ -86,11 +84,12 @@ def commit(self):
             raise
 
     def rollback(self):
-        """Clear pending operations (K8s doesn't support true rollback)."""
+        """Clear pending operations."""
+        # K8s doesn't support transactions, only clear pending operations
         self._pending_operations.clear()
 
     def _job_to_dict(self, job) -> Dict[str, Any]:
-        """Convert SQLAlchemy Job model to dict."""
+        """Convert Job model to dict."""
         return {
             "job_id": job.job_id,
             "name": job.name,
@@ -104,11 +103,12 @@ def _job_to_dict(self, job) -> Dict[str, Any]:
         }
 
     def _create_k8s_job(self, job_data: Dict):
-        """Create placeholder K8s Job for database storage."""
+        """Create K8s Job for metadata storage."""
+        # Creates minimal busybox job that stores metadata in labels/annotations
         job_id = job_data['job_id']
         job_name = f"js-{job_id[:8]}-{job_id[-4:]}"
 
-        # Create minimal job for metadata storage
+        # Busybox container runs once then exits, leaving metadata intact
         job_spec = client.V1JobSpec(
             template=client.V1PodTemplateSpec(
                 spec=client.V1PodSpec(
@@ -126,14 +126,14 @@ def _create_k8s_job(self, job_data: Dict):
             backoff_limit=0
         )
 
-        # Database labels for querying
+        # Labels enable fast K8s label selector queries
         labels = {
             "app.kubernetes.io/managed-by": "jupyter-scheduler-k8s",
             "jupyter-scheduler.io/job-id": self._sanitize(job_data["job_id"]),
             "jupyter-scheduler.io/status": self._sanitize(job_data["status"]),
         }
 
-        # Add schedule indicator for Job vs JobDefinition differentiation
+        # Differentiate Job from JobDefinition using schedule presence
         if job_data.get("schedule"):
             labels["jupyter-scheduler.io/has-schedule"] = "true"
         else:
@@ -142,7 +142,7 @@ def _create_k8s_job(self, job_data: Dict):
         if job_data.get("name"):
             labels["jupyter-scheduler.io/name"] = self._sanitize(job_data["name"])
 
-        # Full metadata in annotation
+        # Store complete job data in annotation for retrieval
         annotations = {
             "jupyter-scheduler.io/job-data": json.dumps(job_data)
         }
@@ -169,7 +169,8 @@ def _create_k8s_job(self, job_data: Dict):
                 raise
 
     def _sanitize(self, value: str) -> str:
-        """Sanitize for K8s labels."""
+        """Sanitize value for K8s labels."""
+        # K8s labels must be alphanumeric, max 63 chars
         value = str(value).lower()
         value = ''.join(c if c.isalnum() or c in '-_.' else '-' for c in value)
         return value.strip('-_.')[:63] or "none"
@@ -194,63 +195,63 @@ def __init__(self, session: K8sSession, model_class):
 
     def filter(self, condition):
         """Add filter condition."""
-        # Parse SQLAlchemy-style condition
+        # Convert SQLAlchemy conditions to K8s label selectors or annotation filters
         if hasattr(condition, 'left') and hasattr(condition.left, 'name'):
             field_name = condition.left.name
             value = getattr(condition.right, 'value', condition.right)
 
             if field_name in ['job_id', 'status', 'name']:
                 self._label_filters[f'jupyter-scheduler.io/{field_name.replace("_", "-")}'] = self.session._sanitize(str(value))
             else:
-                # Store for annotation-based filtering
+                # Complex fields stored in annotations, filtered post-query
                 self._filters[field_name] = value
         elif hasattr(condition, 'type') and condition.type.name == 'in_':
-            # Handle IN clauses like Job.status.in_(['COMPLETED', 'FAILED'])
+            # IN clauses require annotation filtering since K8s labels don't support OR
             field_name = condition.left.name
             if field_name == 'status':
-                # For IN clauses, we'll need to handle multiple label selectors
+                # Multiple values require post-query filtering
                 self._filters['status_in'] = [self.session._sanitize(str(v)) for v in condition.right.value]
 
         return self
 
     def update(self, values: Dict):
-        """Update job in K8s."""
-        # Build label selector from all label filters
+        """Update matching jobs."""
+        # Use labels for efficient K8s filtering
         label_selector = ",".join([f"{k}={v}" for k, v in self._label_filters.items()])
         if not label_selector:
             raise ValueError("Update requires filterable conditions")
 
-        # Find and update K8s Jobs
+        # Query matching jobs using label selector
         jobs = self.session.k8s_batch.list_namespaced_job(
             namespace=self.session.namespace,
             label_selector=label_selector
         )
 
         for job in jobs.items:
-            # Update annotation with new data
+            # Merge new values into existing job data
             if job.metadata.annotations and "jupyter-scheduler.io/job-data" in job.metadata.annotations:
                 job_data = json.loads(job.metadata.annotations["jupyter-scheduler.io/job-data"])
                 job_data.update(values)
                 job_data["update_time"] = get_utc_timestamp()
 
-                # Update annotation
+                # Store updated data back to annotation
                 job.metadata.annotations["jupyter-scheduler.io/job-data"] = json.dumps(job_data)
 
-                # Update corresponding labels if changed
+                # Sync searchable fields to labels for query performance
                 for field, value in values.items():
                     if field in ['status', 'name']:
                         label_key = f"jupyter-scheduler.io/{field.replace('_', '-')}"
                         job.metadata.labels[label_key] = self.session._sanitize(str(value))
 
-                # Patch the job
+                # Apply changes to K8s resource
                 self.session.k8s_batch.patch_namespaced_job(
                     name=job.metadata.name,
                     namespace=self.session.namespace, 
                     body=job
                 )
 
     def one(self):
-        """Get single job (throw if not found)."""
+        """Get single job or raise."""
         result = self.first()
         if result is None:
             raise ValueError("Job not found")
@@ -283,8 +284,8 @@ def delete(self):
             )
 
     def _get_matching_jobs(self):
-        """Get K8s jobs matching current filters."""
-        # Build label selector
+        """Query jobs matching filters."""
+        # Use labels for efficient server-side filtering
         label_selector = ",".join([f"{k}={v}" for k, v in self._label_filters.items()])
 
         jobs = self.session.k8s_batch.list_namespaced_job(
@@ -297,14 +298,14 @@ def _get_matching_jobs(self):
             if job.metadata.annotations and "jupyter-scheduler.io/job-data" in job.metadata.annotations:
                 job_data = json.loads(job.metadata.annotations["jupyter-scheduler.io/job-data"])
 
-                # Apply annotation-based filters
+                # Post-filter using annotation data for complex conditions
                 if self._matches_annotation_filters(job_data):
                     results.append(self._dict_to_job(job_data))
 
         return results
 
     def _matches_annotation_filters(self, job_data: Dict) -> bool:
-        """Check if job data matches annotation-based filters."""
+        """Check annotation-based filter matches."""
         for field, value in self._filters.items():
             if field == 'status_in':
                 if job_data.get('status') not in value:
@@ -316,7 +317,7 @@ def _matches_annotation_filters(self, job_data: Dict) -> bool:
                 if not job_data.get('start_time') or job_data['start_time'] < value:
                     return False
             elif field.endswith('_like'):
-                # Handle LIKE queries (e.g., name LIKE 'prefix%')
+                # SQL LIKE converted to string prefix matching
                 actual_field = field[:-5]
                 actual_value = job_data.get(actual_field, "")
                 if not actual_value.startswith(str(value).rstrip('%')):
@@ -327,101 +328,10 @@ def _matches_annotation_filters(self, job_data: Dict) -> bool:
         return True
 
     def _dict_to_job(self, job_data: Dict):
-        """Convert dict back to Job-like object."""
+        """Convert dict to Job-like object."""
         class JobRecord:
             def __init__(self, data):
                 for k, v in data.items():
                     setattr(self, k, v)
 
-        return JobRecord(job_data)
-
-
-# Store original functions for fallback
-_original_create_session = None
-_original_create_tables = None
-
-
-def k8s_create_session(db_url):
-    """K8s session factory that replaces SQLAlchemy."""
-    if db_url.startswith("k8s://"):
-        namespace = db_url[6:] or "default"
-        def session_factory():
-            return K8sSession(namespace=namespace)
-        return session_factory
-    else:
-        # Fallback to original SQLAlchemy implementation
-        if _original_create_session:
-            return _original_create_session(db_url)
-        else:
-            # Import here to avoid circular imports
-            from jupyter_scheduler.orm import create_session as original_create_session
-            return original_create_session(db_url)
-
-
-def k8s_create_tables(db_url, drop_tables=False, Base=None):
-    """K8s equivalent of create_tables - ensure namespace exists."""
-    if db_url.startswith("k8s://"):
-        namespace = db_url[6:] or "default"
-        
-        try:
-            config.load_incluster_config()
-        except config.ConfigException:
-            config.load_kube_config()
-        
-        k8s_core = client.CoreV1Api()
-        
-        # Ensure namespace exists
-        try:
-            k8s_core.read_namespace(name=namespace)
-        except ApiException as e:
-            if e.status == 404:
-                ns = client.V1Namespace(metadata=client.V1ObjectMeta(name=namespace))
-                k8s_core.create_namespace(body=ns)
-                logger.info(f"Created K8s namespace: {namespace}")
-        
-        logger.info(f"K8s database initialized in namespace: {namespace}")
-    else:
-        # Fallback to original SQLAlchemy implementation  
-        if _original_create_tables:
-            return _original_create_tables(db_url, drop_tables, Base)
-        else:
-            from jupyter_scheduler.orm import create_tables as original_create_tables
-            return original_create_tables(db_url, drop_tables, Base)
-
-
-def install_k8s_backend():
-    """Install K8s backend by monkey patching jupyter_scheduler.orm functions."""
-    try:
-        import jupyter_scheduler.orm as orm
-        
-        # Store originals for fallback
-        global _original_create_session, _original_create_tables
-        _original_create_session = orm.create_session
-        _original_create_tables = orm.create_tables
-        
-        # Replace with K8s-aware functions
-        orm.create_session = k8s_create_session
-        orm.create_tables = k8s_create_tables
-        
-        # Also monkey patch SQLAlchemy's create_engine to handle k8s:// URLs
-        import sqlalchemy
-        original_create_engine = sqlalchemy.create_engine
-        
-        def k8s_aware_create_engine(url, *args, **kwargs):
-            if str(url).startswith("k8s://"):
-                # Return a dummy engine object that won't be used
-                # since our k8s_create_tables handles k8s:// URLs
-                class DummyEngine:
-                    dialect = type('dialect', (), {'name': 'k8s'})()
-                return DummyEngine()
-            return original_create_engine(url, *args, **kwargs)
-        
-        sqlalchemy.create_engine = k8s_aware_create_engine
-        
-        logger.info("K8s database backend installed successfully")
-    except ImportError:
-        logger.warning("jupyter_scheduler not found, K8s backend not installed")
-
-
-# Auto-install K8s backend when this module is imported
-install_k8s_backend()
+        return JobRecord(job_data)