include more tests now that catalog.db is filled

Author: luabida

pysus/api/dadosgov/client.py

@@ -3,14 +3,15 @@
 import pathlib
 from collections.abc import Callable
 from datetime import datetime
-from typing import Annotated, Any, Optional
+from typing import TYPE_CHECKING, Annotated, Any, Optional
 
 import httpx
 from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, PrivateAttr
 from pysus import __version__
 from pysus.api.models import BaseRemoteClient, BaseRemoteFile
 
-from .models import Dataset
+if TYPE_CHECKING:
+    from .models import Dataset
 
 
 def to_datetime(value: Any) -> datetime | None:

pysus/api/dadosgov/models.py

@@ -180,7 +180,7 @@ async def _fetch_files(self) -> list[BaseRemoteFile]:
 
 class Dataset(BaseRemoteDataset):
     ids: list[str] = []
-    client: DadosGov
+    client: "DadosGov"
 
     def __repr__(self):
         return self.name
@@ -191,7 +191,7 @@ def formatter(self, filename: str) -> dict[str, Any]:
 
     async def _fetch_content(self) -> list[Group]:
         items: list[Group] = []
-        client: DadosGov = self.client
+        client: "DadosGov" = self.client
         if self.ids:
             for group_id in self.ids:
                 record = await client.get_dataset(group_id)

pysus/api/ducklake/catalog.py

@@ -24,7 +24,12 @@ class Base(DeclarativeBase):
 file_columns = Table(
     "file_columns",
     Base.metadata,
-    Column("file_id", Integer, ForeignKey("pysus.files.id"), primary_key=True),
+    Column(
+        "file_id",
+        Integer,
+        ForeignKey("pysus.files.id"),
+        primary_key=True,
+    ),
     Column(
         "column_id",
         Integer,
@@ -196,7 +201,10 @@ class CatalogFile(CatalogTable):
         back_populates="files",
     )
     columns: Mapped[list["ColumnDefinition"]] = relationship(
-        "ColumnDefinition", secondary=file_columns, back_populates="files"
+        "ColumnDefinition",
+        secondary=file_columns,
+        back_populates="files",
+        cascade="all, delete",
     )
 
     __table_args__ = (

pysus/api/extensions.py

@@ -40,7 +40,7 @@
 
 
 class File(BaseLocalFile):
-    type: FileType = Field(None)
+    type: FileType = Field("FILE")
 
     async def load(self) -> bytes:
         return await to_thread.run_sync(self.path.read_bytes)
@@ -181,24 +181,28 @@ def rows(self) -> int:
 
     async def load(self, parse: bool = True) -> pd.DataFrame:
         def _load():
-            df = pd.read_parquet(self.path)
+            df = pd.read_parquet(self.path, engine="pyarrow")
             return self.parse_dftypes(df) if parse else df
 
         return await to_thread.run_sync(_load)
 
     async def stream(
-        self, chunk_size: int = 10000
+        self, chunk_size: int = 10000, parse: bool = False
     ) -> AsyncGenerator[pd.DataFrame, None]:
         parquet_file = await to_thread.run_sync(pq.ParquetFile, self.path)
 
         for batch in parquet_file.iter_batches(batch_size=chunk_size):
             df = batch.to_pandas()
+            if parse:
+                df = self.parse_dftypes(df)
             yield df
             await asyncio.sleep(0)
 
     @staticmethod
     def parse_dftypes(df: pd.DataFrame) -> pd.DataFrame:
         def str_to_int(string):
+            if pd.isna(string):
+                return string
             clean = str(string).replace(" ", "")
             return int(clean) if clean.isnumeric() else string
 
@@ -211,15 +215,16 @@ def str_to_date(string):
             return string
 
         cols_to_date = ["DT_NOTIFIC", "DT_SIN_PRI", "DT_NASC", "DT_INTER"]
-        cols_to_int = ["CODMUNRES", "SEXO", "IDADE"]
+        cols_to_int = ["CODMUNRES", "IDADE"]
 
         for col in df.columns:
             if col in cols_to_date:
                 df[col] = df[col].map(str_to_date)
             elif col in cols_to_int:
                 df[col] = df[col].map(str_to_int)
 
-        df = df.map(lambda x: "" if str(x).isspace() else x)
+        df = df.replace(r"^\s+$", "", regex=True)
+
         return df.convert_dtypes()
 
 
@@ -733,5 +738,11 @@ async def instantiate(cls, path: str | Path) -> BaseLocalFile:
         path = Path(path).expanduser().resolve()
         if await to_thread.run_sync(path.is_dir):
             return Directory(path=path, type="DIR")
+
         FileClass = await cls.get_file_class(path)
-        return FileClass(path=path, type=FileClass.type)
+        file_type = getattr(FileClass, "type", "FILE")
+
+        if not isinstance(file_type, str):
+            file_type = "FILE"
+
+        return FileClass(path=path, type=file_type)

pysus/api/ftp/client.py

@@ -57,9 +57,7 @@ def description(self) -> str:
         """
 
     @property
-    def ftp(self) -> FTPLib:
-        if not self._ftp:
-            raise ConnectionError("FTP Not properly connected")
+    def ftp(self) -> FTPLib | None:
         return self._ftp
 
     async def connect(self) -> None:

pysus/api/ftp/models.py

@@ -17,23 +17,28 @@
 )
 from pysus.api.types import State
 
-from .client import FTP, FTPFileInfo, FTPGroupInfo
+from .client import FTP, FTPFileInfo
 
 
 class File(BaseRemoteFile):
     _info: FTPFileInfo = PrivateAttr()
 
     def __init__(self, **data):
         info = data.pop("_info", None)
-        path = data.pop("path", None)
+        if "path" not in data and info and "path" in info:
+            data["path"] = info["path"]
 
         super().__init__(**data)
 
-        if info is not None:
-            self._info = info
-
-        if path is not None:
-            self._path = path
+        self._info = info
+        group_data = self._info.get("group")
+        if group_data:
+            self.group = Group(
+                path=str(self.path.parent),
+                dataset=self.dataset,
+                long_name=group_data.get("long_name", ""),
+                description=group_data.get("description", ""),
+            )
 
     def __repr__(self) -> str:
         return self.name
@@ -53,10 +58,6 @@ def modify(self) -> datetime:
             raise ValueError("File requires a modify date")
         return m
 
-    @property
-    def group_info(self) -> FTPGroupInfo | None:
-        return self._info.get("group")
-
     @property
     def year(self) -> int | None:
         return self._info.get("year")
@@ -109,7 +110,10 @@ async def content(self) -> list[Directory | File]:
     async def load(self) -> None:
         if not isinstance(self.client, FTP):
             raise ValueError("no ftp client found")
-        raw_infos = await self.client._list_directory(self.path, self.formatter)
+        raw_infos = await self.client._list_directory(
+            self.path,
+            self.formatter,
+        )
         self._content = []
 
         current_group = (
@@ -157,8 +161,11 @@ def __init__(
         dataset: Dataset,
         long_name: str,
         description: str = "",
+        **data: Any,
     ):
-        super().__init__(dataset=dataset)
+        data.update({"dataset": dataset, "path": path})
+        super().__init__(**data)
+
         self._long_name = long_name
         self._description = description
         self._dir = Directory(

pysus/api/models.py

@@ -130,7 +130,7 @@ async def to_parquet(
         chunk_size: int = 10000,
         callback: Callable[[int, int], None] | None = None,
     ) -> Parquet:
-        from pysus.api.extensions import ExtensionFactory
+        from pysus.api.extensions import ExtensionFactory, Parquet
 
         if output_path is None:
             output_path = self.path.with_suffix(".parquet")

pysus/api/types.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
 FileType = Literal[
-    None,
+    "FILE",
     "DIR",
     "PARQUET",
     "CSV",

pysus/management/client.py

@@ -13,10 +13,12 @@
     ColumnDefinition,
     DatasetGroup,
     Origin,
+    file_columns,
 )
 from pysus.api.extensions import Parquet
 from pysus.api.ftp.models import File as FTPFile
 from pysus.api.models import BaseRemoteFile
+from sqlalchemy import delete
 
 
 class CatalogManager:
@@ -59,6 +61,7 @@ async def upload(
     ) -> None:
         if not self.pysus._ducklake:
             raise ConnectionError("DuckLake is not connected")
+
         with self.pysus._ducklake._Session() as session:
             dataset = self._get_or_create_dataset(session, file)
             group = self._get_or_create_group(session, file, dataset)
@@ -67,7 +70,7 @@ async def upload(
             if not self._should_upload(file, cat_file):
                 return
 
-            session.commit()
+            session.flush()
 
         parquet_ext = await self.pysus.download_to_parquet(
             file=file, token=self.dadosgov_token, callback=callback
@@ -78,25 +81,37 @@ async def upload(
             f"/{file.dataset.name.lower()}/{parquet_ext.path.name}"
         )
 
-        await self._upload_to_s3(parquet_ext.path, s3_key)
-
         with self.pysus._ducklake._Session() as session:
-            current_dataset = self._get_or_create_dataset(session, file)
-            current_group = self._get_or_create_group(
-                session, file, current_dataset
+            existing_conflict = (
+                session.query(CatalogFile)
+                .filter(
+                    CatalogFile.path == s3_key,
+                    CatalogFile.dataset_id == dataset.id,
+                )
+                .first()
             )
 
-            cat_file = self._get_or_create_file(
-                session, file, current_dataset, current_group
-            )
+            if existing_conflict and existing_conflict.id != cat_file.id:
+                session.execute(
+                    delete(file_columns).where(
+                        file_columns.c.file_id == existing_conflict.id
+                    )
+                )
+                session.flush()
+                session.delete(existing_conflict)
+                session.flush()
+
+            cat_file = session.merge(cat_file)
+
+            await self._upload_to_s3(parquet_ext.path, s3_key)
 
             cat_file.path = s3_key
             cat_file.size = parquet_ext.size
             cat_file.rows = parquet_ext.rows
             cat_file.modified = datetime.utcnow()
             cat_file.origin_modified = file.modify
             cat_file.columns = self._get_or_create_columns(
-                session, current_dataset, parquet_ext
+                session, dataset, parquet_ext
             )
 
             session.commit()
@@ -140,9 +155,8 @@ def _get_or_create_dataset(
         ds_name = file.dataset.name.lower()
         ds = session.query(CatalogDataset).filter_by(name=ds_name).first()
         if not ds:
-            origin = (
-                Origin.FTP if file.client.name.lower() == "ftp" else Origin.API
-            )
+            is_ftp = file.client.name.lower() == "ftp"
+            origin = Origin.FTP if is_ftp else Origin.API
             ds = CatalogDataset(
                 name=ds_name, long_name=file.dataset.long_name, origin=origin
             )
@@ -201,7 +215,7 @@ def _get_or_create_file(
                 size=0,
                 rows=0,
                 modified=datetime.min,
-                origin_path=file.path,
+                origin_path=str(file.path),
                 year=file.year,
                 month=file.month,
                 state=file.state,