start testing the programatically upload of the files on s3

Author: luabida

poetry.lock

@@ -42,6 +42,7 @@ bigtree = "^0.12.2"
 
 pyreaddbc = { version = ">=1.1.0", optional = true }
 pycparser = { version = "2.21", optional = true }
+dotenv = "^0.9.9"
 
 [tool.poetry.extras]
 dbc = ["pyreaddbc", "pycparser"]

pyproject.toml

@@ -11,6 +11,7 @@
 from .ducklake import DuckLakeClient
 from .ftp import FTPClient
 from .models import BaseLocalFile, BaseRemoteFile
+from .extensions import Parquet
 
 Base = declarative_base()
 
@@ -108,8 +109,7 @@ def _attach_client_catalog(self, name: str, path: str):
             existing = conn.exec_driver_sql(q, (abs_path,)).fetchone()
 
             if not existing:
-                conn.exec_driver_sql(f"ATTACH '{abs_path}' AS {
-                                     name} (READ_ONLY)")
+                conn.exec_driver_sql(f"ATTACH '{abs_path}' AS {name} (READ_ONLY)")
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self._ducklake:
@@ -122,7 +122,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     def _get_dest_path(self, file: BaseRemoteFile) -> Path:
         client_name = file.client.name.lower()
-        dataset_name = getattr(file.parent, "name", "unknown_dataset")
+        dataset_name = file.dataset.name.lower()
 
         group_name = ""
         if hasattr(file, "group") and file.group:
@@ -148,8 +148,7 @@ async def _update_state(
     ):
         with self.Session() as session:
             record = (
-                session.query(LocalFileState).filter_by(
-                    path=str(local_path)).first()
+                session.query(LocalFileState).filter_by(path=str(local_path)).first()
             )
             if not record:
                 record = LocalFileState(
@@ -231,36 +230,38 @@ async def download_to_parquet(
         file: BaseRemoteFile,
         token: str = None,
         callback: Callable[[int, int], None] = None,
-    ):
+    ) -> Parquet:
         local_file = await self.download(
             file=file,
             token=token,
             callback=callback,
         )
 
-        if hasattr(local_file, "to_parquet"):
-            original_path = local_file.path
+        if not hasattr(local_file, "to_parquet"):
+            raise NotImplementedError(
+                f"{local_file} can't be converted to Parquet",
+            )
 
-            parquet_file = await local_file.to_parquet(callback=callback)
+        original_path = local_file.path
 
-            await self._update_state(
-                local_path=parquet_file.path,
-                remote_path=file.path,
-                client_name=file.client.name.lower(),
-                status=DownloadStatus.COMPLETED,
-                year=file.year,
-                month=file.month,
-                state=file.state,
-                group=getattr(file.group, "name", None),
-            )
+        parquet_file = await local_file.to_parquet(callback=callback)
 
-            if original_path.exists() and original_path != parquet_file.path:
-                original_path.unlink()
-                await self._delete_record(str(original_path))
+        await self._update_state(
+            local_path=parquet_file.path,
+            remote_path=file.path,
+            client_name=file.client.name.lower(),
+            status=DownloadStatus.COMPLETED,
+            year=file.year,
+            month=file.month,
+            state=file.state,
+            group=getattr(file.group, "name", None),
+        )
 
-            return parquet_file
+        if original_path.exists() and original_path != parquet_file.path:
+            original_path.unlink()
+            await self._delete_record(str(original_path))
 
-        return local_file
+        return parquet_file
 
     def get_local_hierarchy(self):
         with self.Session() as session:

pysus/api/client.py

@@ -46,9 +46,7 @@ def path(self) -> str:
     def extension(self) -> str:
         if self.record.file_name:
             return pathlib.Path(self.record.file_name).suffix
-        return pathlib.Path(
-            self.record.url.split("/")[-1].split("?")[0]
-        ).suffix
+        return pathlib.Path(self.record.url.split("/")[-1].split("?")[0]).suffix
 
     @property
     def size(self) -> int:
@@ -75,9 +73,7 @@ async def _download(
         output: pathlib.Path | None = None,
         callback: Callable[[int], None] | None = None,
     ) -> pathlib.Path:
-        return await self.client._download_file(
-            self, output, callback=callback
-        )
+        return await self.client._download_file(self, output, callback=callback)
 
     async def fetch_size(self) -> int:
         try:
@@ -104,9 +100,13 @@ async def fetch_size(self) -> int:
 
 class Group(BaseRemoteGroup):
     record: ConjuntoDados
-    _formatter: Callable[[Recurso, "Group"], dict[str, Any]] | None = (
-        PrivateAttr(default=None)
-    )
+    _formatter: (
+        Callable[
+            [Recurso, "Group"],
+            dict[str, Any],
+        ]
+        | None
+    ) = PrivateAttr(default=None)
 
     def __init__(
         self,
@@ -136,10 +136,13 @@ def description(self) -> str:
     async def _fetch_files(self) -> list[File]:
         files = []
         for recurso in self.record.resources:
-            metadata = (
-                self._formatter(recurso, self) if self._formatter else {}
+            metadata = self._formatter(recurso, self) if self._formatter else {}
+            file = File(
+                record=recurso,
+                dataset=self.dataset,
+                group=self,
+                _metadata=metadata,
             )
-            file = File(record=recurso, parent=self, _metadata=metadata)
             files.append(file)
         return files
 
@@ -162,8 +165,6 @@ async def _fetch_content(self) -> list[Group]:
             for group_id in self.ids:
                 record = await client.get_dataset(group_id)
                 items.append(
-                    Group(
-                        record=record, dataset=self, formatter=self.formatter
-                    )
+                    Group(record=record, dataset=self, formatter=self.formatter)
                 )
         return items

pysus/api/dadosgov/models.py

@@ -10,6 +10,7 @@
     Integer,
     String,
     Table,
+    Sequence,
 )
 from sqlalchemy.orm import declarative_base, relationship
 
@@ -42,7 +43,11 @@ class Origin(enum.Enum):
 class CatalogDataset(CatalogTable):
     __tablename__ = "datasets"
 
-    id = Column(Integer, primary_key=True)
+    id = Column(
+        Integer,
+        Sequence("datasets_id_seq", schema="pysus"),
+        primary_key=True,
+    )
     name = Column(String, nullable=False, unique=True, index=True)
     long_name = Column(String, nullable=False)
     description = Column(String, nullable=True)
@@ -68,7 +73,11 @@ class CatalogDataset(CatalogTable):
 class ColumnDefinition(CatalogTable):
     __tablename__ = "dataset_columns"
 
-    id = Column(Integer, primary_key=True)
+    id = Column(
+        Integer,
+        Sequence("columns_id_seq", schema="pysus"),
+        primary_key=True,
+    )
     dataset_id = Column(
         Integer,
         ForeignKey("pysus.datasets.id"),
@@ -96,7 +105,11 @@ class ColumnDefinition(CatalogTable):
 class DatasetGroup(CatalogTable):
     __tablename__ = "dataset_groups"
 
-    id = Column(Integer, primary_key=True)
+    id = Column(
+        Integer,
+        Sequence("groups_id_seq", schema="pysus"),
+        primary_key=True,
+    )
     name = Column(String, nullable=False)
     dataset_id = Column(
         Integer,
@@ -123,7 +136,11 @@ class DatasetGroup(CatalogTable):
 class CatalogFile(CatalogTable):
     __tablename__ = "files"
 
-    id = Column(Integer, primary_key=True)
+    id = Column(
+        Integer,
+        Sequence("files_id_seq", schema="pysus"),
+        primary_key=True,
+    )
     dataset_id = Column(
         Integer, ForeignKey("pysus.datasets.id"), nullable=False, index=True
     )
@@ -137,6 +154,7 @@ class CatalogFile(CatalogTable):
     size = Column(Integer, nullable=False)
     rows = Column(Integer, nullable=False)
     modified = Column(DateTime, nullable=False)
+    origin_modified = Column(DateTime, nullable=True)
     sha256 = Column(String(64), nullable=True, index=True)
     year = Column(Integer, nullable=True, index=True)
     month = Column(Integer, nullable=True, index=True)
@@ -151,5 +169,6 @@ class CatalogFile(CatalogTable):
     __table_args__ = (
         Index("ix_files_dataset_group", "dataset_id", "group_id"),
         Index("ix_files_temporal", "year", "month"),
+        Index("ix_files_lookup", "dataset_id", "group_id", "year", "month", "state"),
         {"schema": "pysus"},
     )

pysus/api/ducklake/catalog.py

@@ -155,12 +155,14 @@ async def connect(self, force: bool = False):
 
     async def close(self):
         if self._engine:
-            if self._is_authenticated:
-                await self._upload_catalog()
-
             await anyio.to_thread.run_sync(self._engine.dispose)
+
             self._engine = None
             self._Session = None
+
+            if self._is_authenticated:
+                await self._upload_catalog()
+
             self._s3_client = None
 
     async def _download_file(

pysus/api/ducklake/client.py

@@ -18,18 +18,29 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 from pydantic import Field, PrivateAttr
+from dbfread import DBF as DBFReader
+
 from pysus import CACHEPATH
 from pysus.api.models import BaseCompressedFile, BaseLocalFile, BaseTabularFile
 
 from .types import FileType
 
+import sys
+import ctypes.util
+
 try:
-    from dbfread import DBF as DBFReader
-    from pyreaddbc import dbc2dbf
+    LIBFFI = True
+    if sys.platform.startswith("linux"):
+        LIBFFI = ctypes.util.find_library("ffi") is not None
+
+    if LIBFFI:
+        from pyreaddbc import dbc2dbf
 
-    FTP_IMPORT = True
+        DBC_IMPORT = True
+    else:
+        DBC_IMPORT = False
 except ImportError:
-    FTP_IMPORT = False
+    DBC_IMPORT = False
 
 
 class File(BaseLocalFile):
@@ -160,6 +171,10 @@ def _get_reader_sync():
 class Parquet(BaseTabularFile):
     type: FileType = Field("Parquet")
 
+    @property
+    def schema(self) -> pa.Schema:
+        return pq.read_schema(self.path)
+
     @property
     def columns(self) -> list[str]:
         return pq.read_schema(self.path).names
@@ -607,7 +622,9 @@ def _extract():
 class FTPNotImported(BaseTabularFile):
     type: FileType = Field(None)
     import_err: ClassVar[str] = """
-        run "pip install pysus[dbc]" to handle DBC files
+        run "pip install pysus[dbc]" to handle DBC files.
+        Make sure you also have libffi installed on the system. It may not work
+        on Windows
     """
 
     @property
@@ -647,7 +664,7 @@ class ExtensionFactory:
         ".csv": CSV,
         ".parquet": Parquet,
         ".dbf": DBF,
-        ".dbc": DBC if FTP_IMPORT else FTPNotImported,
+        ".dbc": DBC if DBC_IMPORT else FTPNotImported,
         ".pdf": PDF,
         ".json": JSON,
     }