updated process dataset function

Author: cka-y

functions-python/batch_process_dataset/src/main.py

@@ -16,6 +16,7 @@
 
 import base64
 import json
+import logging
 import os
 import random
 import uuid
@@ -28,15 +29,13 @@
 from cloudevents.http import CloudEvent
 from google.cloud import storage
 from sqlalchemy import func
+from sqlalchemy.orm import Session
 
+from shared.database.database import with_db_session, refresh_materialized_view
 from shared.database_gen.sqlacodegen_models import Gtfsdataset, t_feedsearch, Gtfsfile
 from shared.dataset_service.main import DatasetTraceService, DatasetTrace, Status
-from shared.database.database import with_db_session, refresh_materialized_view
-import logging
-
 from shared.helpers.logger import init_logger, get_logger
-from shared.helpers.utils import download_and_get_hash
-from sqlalchemy.orm import Session
+from shared.helpers.utils import download_and_get_hash, get_hash_from_file
 
 init_logger()
 
@@ -177,6 +176,8 @@ def upload_file_to_storage(
                         id=str(uuid.uuid4()),
                         file_name=file_name,
                         file_size_bytes=os.path.getsize(file_path),
+                        hosted_url=file_blob.public_url if public else None,
+                        hash=get_hash_from_file(file_path),
                     )
                 )
         return blob, extracted_files

functions-python/helpers/utils.py

@@ -78,6 +78,17 @@ def download_url_content(url, with_retry=False):
         raise e
 
 
+def get_hash_from_file(file_path, hash_algorithm="sha256", chunk_size=8192):
+    """
+    Returns the hash of a file
+    """
+    hash_object = hashlib.new(hash_algorithm)
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            hash_object.update(chunk)
+    return hash_object.hexdigest()
+
+
 def download_and_get_hash(
     url,
     file_path,

functions-python/tasks_executor/src/tasks/dataset_files/rebuild_missing_dataset_files.py

@@ -1,16 +1,17 @@
 import logging
 import os
-import uuid
-import hashlib
 import shutil
+import ssl
 import tempfile
-import zipfile
 import urllib.request
-import ssl
+import uuid
+import zipfile
 
 from google.cloud import storage
+
 from shared.database.database import with_db_session
 from shared.database_gen.sqlacodegen_models import Gtfsdataset, Gtfsfile
+from shared.helpers.utils import get_hash_from_file
 
 # Disable SSL verification — trusted internal sources only
 ssl._create_default_https_context = ssl._create_unverified_context
@@ -115,15 +116,12 @@ def process_dataset(dataset: Gtfsdataset):
                     blob.upload_from_filename(file_path)
                     blob.make_public()
 
-                with open(file_path, "rb") as f:
-                    sha256_hash = hashlib.sha256(f.read()).hexdigest()
-
                 gtfs_files.append(
                     Gtfsfile(
                         id=str(uuid.uuid4()),
                         file_name=file_name,
                         file_size_bytes=os.path.getsize(file_path),
-                        hash=sha256_hash,
+                        hash=get_hash_from_file(file_path),
                         hosted_url=blob.public_url if dataset.latest else None,
                     )
                 )

functions-python/tasks_executor/tests/tasks/dataset_files/test_rebuild_missing_dataset_files.py

@@ -115,7 +115,7 @@ def test_process_dataset_latest(
         urlopen_mock.return_value.__enter__.return_value = mock_response
 
         mock_file = MagicMock()
-        mock_file.read.return_value = b"file content"
+        mock_file.read.side_effect = [b"chunk1", b"chunk2", b""]  # ends properly
         open_mock.return_value.__enter__.return_value = mock_file
 
         mock_blob = MagicMock()
@@ -159,7 +159,7 @@ def test_process_dataset_not_latest(
         urlopen_mock.return_value.__enter__.return_value = mock_response
 
         mock_file = MagicMock()
-        mock_file.read.return_value = b"file content"
+        mock_file.read.side_effect = [b"chunk1", b"chunk2", b""]  # ends properly
         open_mock.return_value.__enter__.return_value = mock_file
 
         mock_blob = MagicMock()