Direct video upload¶

# /// script
# requires-python = "==3.14.*"
# dependencies = [
#     "httpx2[http2]>=2.5",
#     "pydantic-settings==2.4.0",
#     "sqlalchemy>=2",
#     "platformdirs==4.3.6",
#     "stamina>=24",
#     "structlog==25.4.0",
# ]
# ///
import copy
import dataclasses
from collections.abc import Generator, Iterator
from datetime import datetime
from functools import cache, cached_property
from pathlib import Path
from typing import Annotated, Any, TypedDict

import httpx2
import platformdirs
import pytz
import sqlalchemy as sa
import stamina
import structlog
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, FilePath, PlainSerializer
from pydantic_settings import BaseSettings, SettingsConfigDict
from sqlalchemy import create_engine
from sqlalchemy.dialects.sqlite import JSON
from sqlalchemy.exc import NoResultFound
from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, Session, mapped_column


def create_custom_logger(name: str):
    """Create a logger with specific configuration."""
    return structlog.wrap_logger(
        structlog.PrintLogger(file=None),  # Prints to stdout
        processors=[
            structlog.stdlib.add_log_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.dev.ConsoleRenderer(),
        ],
        context_class=dict,
        logger_factory_args=(name,) if name else (),
    )


logger = structlog.get_logger("thestream")

APP_PATH = platformdirs.user_data_path("heapstream", appauthor="heapstream", ensure_exists=True)
SQLITE_path = APP_PATH / "upload-sdk.sqlite"
MB = 1024 * 1024


@cache
def get_engine() -> sa.engine.Engine:
    return create_engine(f"sqlite:///{SQLITE_path.as_posix()}", echo=False)


@cache
def create_all_tables() -> None:
    Base.metadata.create_all(get_engine())


class Base(MappedAsDataclass, DeclarativeBase):
    pass


class VideoUpload(Base):
    __tablename__ = "video_upload"
    project_id: Mapped[int] = mapped_column()
    id: Mapped[int] = mapped_column()
    path: Mapped[str] = mapped_column()
    created_on: Mapped[datetime] = mapped_column(default_factory=lambda: datetime.now(pytz.UTC), init=False)
    upload_id: Mapped[str] = mapped_column()
    key: Mapped[str] = mapped_column()
    data: Mapped[dict] = mapped_column(JSON(none_as_null=True), default_factory=dict)
    __table_args__ = (
        sa.PrimaryKeyConstraint("project_id", "id", name="video_upload_pk"),
        sa.Index("path_unique_idx", "project_id", "path", unique=True),
    )


def read_file_in_chunks(file_path: Path, chunk_size: int) -> Iterator[bytes]:
    """Lazy function (generator) to read a file piece by piece."""
    with open(file_path, "rb") as file_object:
        while True:
            data = file_object.read(chunk_size)
            if not data:
                break
            yield data


class PartDict(TypedDict):
    PartNumber: int
    ETag: str


PartsList = list[PartDict]


class Metadata(BaseModel):
    watermark_id: int | None = Field(default=None, description="Watermark to apply to this video.")
    normalize_audio: bool | None = Field(default=None, description="Normalize audio of the video.")
    auto_tt: list[str] | None = Field(
        default=None,
        description="Languages to auto transcribe the video. `None` will inherit project.auto_tt settings.",
    )
    encoding_tier: str | None = Field(
        default=None,
        description="Encoding tier to use. If set to `None` it will inherit project.encoding_tier.",
    )
    size: int = Field(default=0, description="Size of the video file in bytes.")


InIntOutStr = Annotated[
    int,
    BeforeValidator(int),
    PlainSerializer(func=str, return_type=str),
]


@dataclasses.dataclass(frozen=True)
class UploadPart:
    id: int
    data: bytes


def retry_exceptions(exc: Exception) -> bool:
    return True


def retry_skip_404(exc: Exception) -> bool:
    if isinstance(exc, httpx2.HTTPStatusError):
        if exc.response.status_code == 404:
            return False
    return True


class VideoUploadManager(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    session: httpx2.Client
    project_id: int
    video_path: Path
    host: str
    video_id: InIntOutStr = Field(default=0)
    key: str = Field(default="")
    upload_id: str = Field(default="")
    # chunk size must be exact on all parts (only the last part can be different)
    chunk_size: int = 500 * MB
    parts: PartsList = Field(default_factory=lambda: [])
    metadata: Metadata = Field(default_factory=Metadata)
    existing_part_ids: set[int] = Field(default_factory=set)

    @property
    def base_url(self) -> str:
        return f"{self.host}/api/v1/project/{self.project_id}"

    @property
    def upload_url(self) -> str:
        return f"{self.base_url}/direct_upload/s3/multipart"

    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> "VideoUploadManager":
        # httpx2.Client contains internal locks; share the same instance rather than copying
        return type(self)(
            session=self.session,
            project_id=self.project_id,
            video_path=self.video_path,
            host=self.host,
            chunk_size=self.chunk_size,
            metadata=copy.deepcopy(self.metadata, memo or {}),
        )

    def model_post_init(self, __context: Any) -> None:
        create_all_tables()
        self.set_video_path(self.video_path)

    def set_video_path(self, video_path: Path) -> None:
        """
        Change the video_path before starting to upload.
        :param video_path: Path to video file.
        :return: None
        """
        self.video_path = video_path
        self.metadata.size = video_path.stat().st_size

    @stamina.retry(on=retry_skip_404)
    def get_remote_video(self, video_id: int) -> httpx2.Response:
        """Get remote video."""
        r0 = self.session.get(f"{self.base_url}/video/{video_id}")
        r0.raise_for_status()
        return r0

    def start_video_upload(self) -> httpx2.Response | None:
        """
        Start the direct upload.
        :return: Response of the "start" request. Includes the video_id, S3 key, S3 `uploadId` to start
        uploading parts and complete the upload.
        """
        # check for existing videoUpload in local db!
        query = (
            sa.select(VideoUpload)
            .filter(VideoUpload.project_id == self.project_id)
            .filter(VideoUpload.path == self.video_path.as_posix())
        )
        with Session(get_engine()) as session:
            try:
                existing_video = session.scalars(query).one()
            except NoResultFound:
                return self.create_new_video()
            else:
                # check if remote video exists!!!
                try:
                    self.get_remote_video(existing_video.id)
                except httpx2.HTTPStatusError as e:
                    # if 404 error, remove video and continue uploading!
                    if e.response.status_code == 404:
                        logger.info("remote_video_not_found", video_id=existing_video.id)
                        session.delete(existing_video)
                        session.commit()
                        return self.create_new_video()
                    raise
                # resume by listing parts and getting correct part!
                self.video_id = existing_video.id
                self.key = existing_video.key
                self.upload_id = existing_video.upload_id
                # add existing parts!
                existing_parts = self.list_parts()
                for part in existing_parts:
                    self.existing_part_ids.add(part["PartNumber"])
                    self.track_part(part)

    def track_part(self, part: PartDict):
        self.parts.append(part)

    @stamina.retry(on=retry_exceptions)
    def create_new_video(self) -> httpx2.Response:
        filename = self.video_path.name
        metadata_dict = self.metadata.model_dump(exclude_unset=True)
        data = {"filename": filename, "metadata": metadata_dict}
        r0 = self.session.post(self.upload_url, json=data)
        # create video and save it!
        r0.raise_for_status()
        self.video_id = int(r0.json()["video"]["id"])
        self.key = r0.json()["key"]
        self.upload_id = r0.json()["uploadId"]
        new_video = VideoUpload(
            path=self.video_path.as_posix(),
            project_id=self.project_id,
            id=self.video_id,
            upload_id=self.upload_id,
            key=self.key,
        )
        with Session(get_engine()) as session:
            session.add(new_video)
            session.commit()
        return r0

    def do_video_upload(self) -> int:
        """
        Start & complete a video upload.
        :return: the `id` of the video that was uploaded.
        """
        self.start_video_upload()
        self.upload_parts()
        self.complete_upload()
        return self.video_id

    @stamina.retry(on=retry_exceptions)
    def complete_upload(self) -> httpx2.Response:
        """
        Complete the direct upload after uploading all parts.
        :return: Response of the "complete" request.
        """
        url3 = f"{self.upload_url}/{self.upload_id}/complete"
        r3 = self.session.post(url3, params={"key": self.key}, json={"parts": self.parts})
        r3.raise_for_status()
        # delete the video in sql
        query = (
            sa.delete(VideoUpload)
            .filter(VideoUpload.project_id == self.project_id)
            .filter(VideoUpload.id == self.video_id)
        )
        with Session(get_engine()) as session:
            session.execute(query)
            session.commit()
        return r3

    def upload_parts(self) -> None:
        """Upload all parts of the file."""
        for part in self.parts_iterator():
            if part.id in self.existing_part_ids:
                logger.info("skipping_part", part_id=part.id)
                continue
            self.upload_part(part)

    def parts_iterator(self) -> Generator[UploadPart, None, None]:
        for part_id, chunk in enumerate(read_file_in_chunks(self.video_path, self.chunk_size), start=1):
            yield UploadPart(part_id, chunk)

    def upload_part(self, part: UploadPart) -> None:
        """
        Upload a single part.
        :param part: The part object.
        :return:None
        """
        upload_url = self.get_upload_url(part)
        self.upload_part_data(upload_url, part)

    @stamina.retry(on=retry_exceptions)
    def upload_part_data(self, upload_url: str, part: UploadPart) -> None:
        # upload part to s3
        r2 = self.session.put(upload_url, content=part.data)
        r2.raise_for_status()
        # track ETags
        e_tag = r2.headers["ETag"]
        self.track_part({"PartNumber": part.id, "ETag": e_tag})
        logger.info("uploaded_part", part_id=part.id)

    @stamina.retry(on=retry_exceptions)
    def get_upload_url(self, part: UploadPart) -> str:
        url1 = f"{self.upload_url}/{self.upload_id}/{part.id}"
        # get part url
        r1 = self.session.get(url1, params={"key": self.key})
        r1.raise_for_status()
        upload_url = r1.json()["url"]
        return upload_url

    @stamina.retry(on=retry_exceptions)
    def abort_upload(self) -> httpx2.Response:
        """
        Abort a direct upload.
        """
        url2 = f"{self.upload_url}/{self.upload_id}"
        r0 = self.session.delete(url2, params={"key": self.key})
        r0.raise_for_status()
        return r0

    @stamina.retry(on=retry_exceptions)
    def list_parts(self) -> list[PartDict]:
        url2 = f"{self.upload_url}/{self.upload_id}"
        r0 = self.session.get(url2, params={"key": self.key})
        r0.raise_for_status()
        return r0.json()


class Settings(
    BaseSettings,
    cli_parse_args=True,
    cli_enforce_required=True,
    cli_hide_none_type=True,
    cli_avoid_json=True,
):
    model_config = SettingsConfigDict(nested_model_default_partial_update=True)

    apikey_id: str = Field(description="Your ApiKey.id")
    apikey_key: str = Field(description="Your ApiKey.key")
    project_id: int = Field(description="Your Project.id")
    video_path: FilePath = Field(description="Absolute path to a video file in your computer.")
    host: str = Field(default="https://app.heapstream.com", description="Host of the api.")
    metadata: Metadata = Field(default=Metadata(), description="Metadata to set for the video.")

    @cached_property
    def session(self) -> httpx2.Client:
        return httpx2.Client(
            auth=httpx2.BasicAuth(username=self.apikey_id, password=self.apikey_key),
            follow_redirects=True,
        )


def main() -> None:
    """Direct upload a video to HeapStream."""
    # https://github.com/pydantic/pydantic/pull/3972
    settings = Settings()
    upload_manager = VideoUploadManager(
        session=settings.session,
        project_id=settings.project_id,
        video_path=settings.video_path,
        host=settings.host,
        metadata=settings.metadata,
    )
    upload_manager.do_video_upload()
    logger.info("video_uploaded", video_id=upload_manager.video_id)


if __name__ == "__main__":
    main()