plesna/plesna/storage/repository/fs_repository.py

from pathlib import Path

from pydantic import BaseModel, computed_field

from plesna.libs.string_tools import extract_values_from_pattern
from plesna.models.storage import Partition, Schema, Table
from plesna.storage.repository.repository import Repository


class FSTable(BaseModel):
    name: str
    repo_id: str
    schema_id: str
    id: str
    path: Path
    is_partitionned: bool
    partitions: list[str] = []

    @computed_field
    @property
    def ref(self) -> Table:
        if self.is_partitionned:
            datas = [str(self.path.absolute() / p) for p in self.partitions]
        else:
            datas = [str(self.path.absolute())]

        return Table(
            id=self.id,
            repo_id=self.repo_id,
            schema_id=self.schema_id,
            name=self.name,
            value=str(self.path.absolute()),
            partitions=self.partitions,
            datas=datas,
        )


class FSSchema(BaseModel):
    name: str
    repo_id: str
    id: str
    path: Path
    tables: list[str]

    @computed_field
    @property
    def ref(self) -> Schema:
        return Schema(
            id=self.id,
            repo_id=self.repo_id,
            name=self.name,
            value=str(self.path.absolute()),
            tables=self.tables,
        )


class FSRepositoryError(ValueError):
    pass


class FSRepository(Repository):
    """Repository based on files tree structure

    - first level: schemas
    - second level: tables
    - third level: partition (actual datas)

    """

    ID_FMT = {
        "schema": "{repo_id}-{schema_name}",
        "table": "{schema_id}-{table_name}",
    }

    def __init__(self, id: str, name: str, basepath: str):
        super().__init__(id, name)

        self._basepath = Path(basepath)
        assert self._basepath.exists()

    def ls(self, dir="", only_files=False, only_directories=False, recursive=False) -> list[str]:
        """List files in dir

        :param dir: relative path from self._basepath
        :param only_files: if true return only files
        :param only_directories: if true return only directories
        :param recursive: list content recursively (only for)
        :return: list of string describing path from self._basepath / dir
        """
        dirpath = self._basepath / dir

        if recursive:
            paths = dirpath.rglob("*")
        else:
            paths = dirpath.iterdir()

        if only_files:
            return [
                str(f.relative_to(dirpath))
                for f in paths
                if not f.is_dir() and not str(f).startswith(".")
            ]
        if only_directories:
            return [
                str(f.relative_to(dirpath))
                for f in paths
                if f.is_dir() and not str(f).startswith(".")
            ]

        return [str(f.relative_to(dirpath)) for f in paths if not str(f).startswith(".")]

    def parse_id(self, string: str, id_type: str) -> dict:
        if id_type not in self.ID_FMT:
            raise FSRepositoryError(
                "Wrong id_type. Gots {id_type} needs to be one of {self.ID_FMT.values}"
            )
        parsed = extract_values_from_pattern(self.ID_FMT[id_type], string)
        if not parsed:
            raise FSRepositoryError(
                f"Wrong format for {id_type}. Got {string} need {self.ID_FMT['id_type']}"
            )
        return parsed

    def schemas(self) -> list[str]:
        """List schemas (sub directories within basepath)"""
        subdirectories = self.ls("", only_directories=True)
        return [
            self.ID_FMT["schema"].format(repo_id=self.id, schema_name=d) for d in subdirectories
        ]

    def _schema(self, schema_id: str) -> FSSchema:
        """List schemas (sub directories within basepath)"""
        parsed = self.parse_id(schema_id, "schema")

        repo_id = parsed["repo_id"]
        schema_name = parsed["schema_name"]
        schema_path = self._basepath / schema_name

        if repo_id != self.id:
            raise FSRepositoryError("Trying to get schema that don't belong in this repository")

        tables = self.tables(schema_id)
        return FSSchema(
            name=schema_name,
            id=schema_id,
            repo_id=self.id,
            schema_id=schema_id,
            path=schema_path,
            tables=tables,
        )

    def schema(self, schema_id: str) -> Schema:
        return self._schema(schema_id).ref

    def _tables(self, schema_id: str) -> list[str]:
        parsed = self.parse_id(schema_id, "schema")
        tables = self.ls(parsed["schema_name"])
        return [self.ID_FMT["table"].format(table_name=t, schema_id=schema_id) for t in tables]

    def tables(self, schema_id: str = "") -> list[str]:
        if schema_id:
            return self._tables(schema_id)

        tables = []
        for schema in self.schemas():
            tables += self._tables(schema)
        return tables

    def _table(self, table_id: str) -> FSTable:
        """Get infos on the table"""
        parsed = self.parse_id(table_id, "table")
        schema = self._schema(parsed["schema_id"])

        if not schema.path.exists():
            raise FSRepositoryError(f"The schema {schema.id} does not exists.")

        table_subpath = f"{schema.name}/{parsed['table_name']}"
        table_path = self._basepath / table_subpath

        is_partitionned = table_path.is_dir()
        if is_partitionned:
            partitions = self.ls(table_subpath, only_files=True)
        else:
            partitions = []

        return FSTable(
            name=parsed["table_name"],
            id=table_id,
            repo_id=self.id,
            schema_id=schema.id,
            path=table_path,
            is_partitionned=is_partitionned,
            partitions=partitions,
        )

    def table(self, table_id: str) -> Table:
        return self._table(table_id).ref