insights/products/notebooks/backend/models.py

from django.core.exceptions import ValidationError
from django.db import models
from django.db.models import JSONField, QuerySet
from django.utils import timezone

from insights.models.file_system.constants import DEFAULT_SURFACE
from insights.models.file_system.file_system_mixin import FileSystemSyncMixin
from insights.models.file_system.file_system_representation import FileSystemRepresentation
from insights.models.scoping.root_mixin import TeamScopedRootMixin
from insights.models.team import Team
from insights.models.utils import (
    RootTeamMixin,
    UUIDModel,
    UUIDTModel,
    build_partial_uniqueness_constraint,
    build_unique_relationship_check,
)
from insights.utils import generate_short_id


class Notebook(FileSystemSyncMixin, RootTeamMixin, UUIDTModel):
    class Visibility(models.TextChoices):
        INTERNAL = "internal", "internal"
        DEFAULT = "default", "default"

    short_id = models.CharField(max_length=12, blank=True, default=generate_short_id)
    team = models.ForeignKey("insights.Team", on_delete=models.CASCADE)
    title = models.CharField(max_length=256, blank=True, null=True)
    content: JSONField = JSONField(default=None, null=True, blank=True)
    text_content = models.TextField(blank=True, null=True)
    deleted = models.BooleanField(default=False)
    visibility = models.CharField(choices=Visibility, default=Visibility.DEFAULT, max_length=20)
    version = models.IntegerField(default=0)
    created_at = models.DateTimeField(auto_now_add=True, blank=True)
    created_by = models.ForeignKey("insights.User", on_delete=models.SET_NULL, null=True, blank=True)
    last_modified_at = models.DateTimeField(default=timezone.now)
    last_modified_by = models.ForeignKey(
        "insights.User",
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="modified_notebooks",
    )
    kernel_cpu_cores = models.FloatField(null=True, blank=True)
    kernel_memory_gb = models.FloatField(null=True, blank=True)
    kernel_idle_timeout_seconds = models.IntegerField(null=True, blank=True)

    class Meta:
        unique_together = ("team", "short_id")
        db_table = "insights_notebook"

    @classmethod
    def get_file_system_unfiled(cls, team: "Team", surface: str = DEFAULT_SURFACE) -> QuerySet["Notebook"]:
        base_qs = cls.objects.filter(team=team, deleted=False)
        return cls._filter_unfiled_queryset(base_qs, team, type="notebook", ref_field="short_id", surface=surface)

    def get_file_system_representation(self) -> FileSystemRepresentation:
        return FileSystemRepresentation(
            base_folder=self._get_assigned_folder("Unfiled/Notebooks"),
            type="notebook",  # sync with APIScopeObject in scopes.py
            ref=str(self.short_id),
            name=self.title or "Untitled",
            href=f"/notebooks/{self.short_id}",
            meta={"created_at": str(self.created_at), "created_by": self.created_by_id},
            should_delete=self.deleted or self.visibility == self.Visibility.INTERNAL,
        )


RELATED_OBJECTS = ("group", "account")


class ResourceNotebook(UUIDTModel):
    """
    Generic relationship table linking notebooks to various resources.
    This allows notebooks to be associated with multiple models.

    Uses explicit foreign keys rather than GenericForeignKey for better performance
    and type safety.
    """

    notebook = models.ForeignKey("notebooks.Notebook", on_delete=models.CASCADE, related_name="resources")

    # Relationships (exactly one must be set)
    # When adding a new foreign key, make sure to add the foreign key field and append field name
    # to the `RELATED_OBJECTS` tuple above.
    # Group relationship is not a foreign key because the table lives in another database instance, which would cause IntegrityErrors when inserting a new ResourceNotebook row.

    group = models.IntegerField(
        null=True,
        blank=True,
        db_column="group_id",
    )
    account = models.ForeignKey(
        "customer_analytics.Account",
        null=True,
        blank=True,
        on_delete=models.CASCADE,
        related_name="notebooks",
    )

    class Meta:
        unique_together = ("notebook", *RELATED_OBJECTS)
        constraints = [
            *[
                build_partial_uniqueness_constraint(
                    field="notebook",
                    related_field=related_field,
                    constraint_name=f"unique_notebook_{related_field}",
                )
                for related_field in RELATED_OBJECTS
            ],
            models.CheckConstraint(
                condition=build_unique_relationship_check(RELATED_OBJECTS), name="exactly_one_notebook_related_resource"
            ),
        ]
        db_table = "insights_resourcenotebook"

    def clean(self):
        super().clean()
        """Ensure that exactly one resource field is set."""
        if sum(map(bool, [getattr(self, field) for field in RELATED_OBJECTS])) != 1:
            raise ValidationError("Exactly one resource field must be set.")

    def save(self, *args, **kwargs):
        self.full_clean()
        return super().save(*args, **kwargs)

    @property
    def resource_type(self) -> str:
        """Return the type of the related resource."""
        for field in RELATED_OBJECTS:
            if getattr(self, field):
                return field
        return "unknown"

    @property
    def resource(self):
        """Return the actual related resource object."""
        resource_type = self.resource_type
        if resource_type != "unknown":
            return getattr(self, resource_type)
        return None

    def __str__(self) -> str:
        return f"Notebook {self.notebook.short_id} -> {self.resource_type} {self.resource}"


class KernelRuntime(UUIDTModel):
    class Backend(models.TextChoices):
        MODAL = "modal", "modal"
        DOCKER = "docker", "docker"

    class Status(models.TextChoices):
        STARTING = "starting", "starting"
        RUNNING = "running", "running"
        STOPPED = "stopped", "stopped"
        TIMED_OUT = "timed_out", "timed_out"
        DISCARDED = "discarded", "discarded"
        ERROR = "error", "error"

    team = models.ForeignKey("insights.Team", on_delete=models.CASCADE)
    notebook = models.ForeignKey("notebooks.Notebook", on_delete=models.SET_NULL, null=True, blank=True)
    notebook_short_id = models.CharField(max_length=12)
    user = models.ForeignKey("insights.User", on_delete=models.SET_NULL, null=True, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    last_used_at = models.DateTimeField(default=timezone.now)
    status = models.CharField(choices=Status, default=Status.STARTING, max_length=20)
    backend = models.CharField(choices=Backend, default=Backend.DOCKER, max_length=20)
    kernel_id = models.CharField(max_length=64, null=True, blank=True)
    kernel_pid = models.IntegerField(null=True, blank=True)
    connection_file = models.TextField(null=True, blank=True)
    sandbox_id = models.CharField(max_length=128, null=True, blank=True)
    last_error = models.TextField(null=True, blank=True)
    server_url = models.TextField(null=True, blank=True)
    server_connect_token = models.TextField(null=True, blank=True)
    # The DuckDB objects a SQL node can currently SELECT from, as of this kernel's last run
    # (Journey 7). Kernel-scoped rather than notebook-scoped on purpose: a row is only ever
    # reused while its sandbox is verifiably alive, so the snapshot cannot outlive the kernel
    # that produced it and go stale — a restarted kernel gets a new row with this unset.
    frames: JSONField = JSONField(default=None, null=True, blank=True)
    # Which run's snapshot `frames` holds, as that run's created_at. The kernel executes runs
    # one at a time in arrival order, but their callbacks land on different web workers and can
    # arrive out of order, so a slow older callback must not overwrite a newer snapshot. Runs
    # are created before dispatch, so created_at orders them the same way the kernel ran them.
    frames_run_created_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        db_table = "insights_kernelruntime"
        indexes = [
            models.Index(fields=["team", "notebook_short_id", "user", "status"]),
        ]


class NotebookNodeRun(TeamScopedRootMixin, UUIDModel):
    """A single execution of a revamped-notebooks (SQLV2) node.
    The primary key is the run_id referenced by the run/callback/stream endpoints.
    """

    class Status(models.TextChoices):
        RUNNING = "running", "running"
        DONE = "done", "done"
        FAILED = "failed", "failed"
        # A user-requested stop (Journey 9); unlike FAILED, the envelope's captured
        # stdout/stderr still surface to the UI.
        INTERRUPTED = "interrupted", "interrupted"

    class NodeType(models.TextChoices):
        INSIGHTSQL = "insightsql", "insightsql"
        PYTHON = "python", "python"
        DUCKDB = "duckdb", "duckdb"

    # db_constraint=False: creating a real FK to the hot insights_team table locks it on deploy.
    # Tenant isolation is still enforced by the fail-closed TeamScopedRootMixin manager.
    team = models.ForeignKey("insights.Team", on_delete=models.CASCADE, db_constraint=False)
    notebook = models.ForeignKey("notebooks.Notebook", on_delete=models.CASCADE)
    # Who ran it. Kernels are per user, so this is the second half of a KernelRuntime's scope —
    # the callback needs it to file the frame snapshot without a user-blind lookup by id.
    # db_constraint=False: a real FK to the hot insights_user table locks it on deploy.
    # db_index=False: nothing queries runs by user — it is only ever read off a run we already
    # hold. DO_NOTHING keeps that true: SET_NULL would have Django's collector issue an
    # `UPDATE … WHERE user_id = …` against this unindexed column on every user delete, on the
    # table that grows fastest. Nothing enforces referential integrity here anyway
    # (db_constraint=False), and a dangling id already reads back as None.
    user = models.ForeignKey(
        "insights.User", on_delete=models.DO_NOTHING, null=True, blank=True, db_constraint=False, db_index=False
    )
    node_id = models.CharField(max_length=128)
    # How the run executed: insightsql pushed to Datastore (pages re-query by `code`); python and
    # duckdb ran in the sandbox kernel (pages slice the on-sandbox result frame by `result_id`).
    node_type = models.CharField(choices=NodeType, default=NodeType.INSIGHTSQL, max_length=20)
    # The node's code at run time — paging must re-query what produced the result,
    # not whatever the editor holds now.
    code = models.TextField(blank=True, default="")
    # The direct-query data source this run executed against (null = Insights's own Datastore).
    # Stored, not just passed through, because `code` is only meaningful on the engine that ran
    # it: a re-query or a cross-cell reference has to target the same one. A plain id rather
    # than an FK — a deleted source must not take run history with it.
    connection_id = models.UUIDField(null=True, blank=True)
    # Whether `code` went to the connection verbatim instead of being compiled from InsightsQL. Set,
    # it means `code` is the engine's own dialect and the InsightsQL parser can't read it.
    send_raw_query = models.BooleanField(default=False, db_default=False)
    status = models.CharField(choices=Status, default=Status.RUNNING, max_length=20)
    envelope: JSONField = JSONField(default=None, null=True, blank=True)
    result_id = models.UUIDField(null=True, blank=True)
    # Which kernel this run was dispatched to, so the callback can file the run's frame
    # snapshot against the right KernelRuntime. A plain id rather than an FK: runtime rows
    # are transient (starting/stopped/discarded/error) and run history must not be coupled
    # to their churn — see sql_v2_result_delivery.md, "Related model notes".
    kernel_runtime_id = models.UUIDField(null=True, blank=True)
    error = models.TextField(null=True, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "insights_notebooknoderun"
        indexes = [
            models.Index(fields=["team", "notebook", "node_id"]),
        ]