Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -77,4 +77,7 @@ coverage_reports
tf.plan

# CSV generation output files
functions-python/**/*.csv
functions-python/**/*.csv

# Local emulators
.cloudstorage
1 change: 1 addition & 0 deletions functions-python/batch_process_dataset/.coveragerc
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ omit =
*/database_gen/*
*/dataset_service/*
*/shared/*
*/scripts/*

[report]
exclude_lines =
Expand Down
3 changes: 2 additions & 1 deletion functions-python/batch_process_dataset/requirements_dev.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@ Faker
pytest~=7.4.3
urllib3-mock
requests-mock
python-dotenv~=1.0.0
python-dotenv~=1.0.0
gcp-storage-emulator
47 changes: 30 additions & 17 deletions functions-python/batch_process_dataset/src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,18 +128,15 @@ def download_content(self, temporary_file_path):
logger=self.logger,
)
is_zip = zipfile.is_zipfile(temporary_file_path)
if is_zip:
extracted_file_path = os.path.join(
temporary_file_path.split(".")[0], "extracted"
)
with zipfile.ZipFile(temporary_file_path, "r") as zip_ref:
zip_ref.extractall(os.path.dirname(extracted_file_path))
# List all files in the extracted directory
extracted_files = os.listdir(os.path.dirname(extracted_file_path))
self.logger.info(f"Extracted files: {extracted_files}")
return file_hash, is_zip

def upload_file_to_storage(self, source_file_path, dataset_stable_id):
def upload_file_to_storage(
self,
source_file_path,
dataset_stable_id,
extracted_files_path,
public=True,
):
"""
Uploads a file to the GCP bucket
"""
Expand All @@ -153,12 +150,12 @@ def upload_file_to_storage(self, source_file_path, dataset_stable_id):
blob = bucket.blob(target_path)
with open(source_file_path, "rb") as file:
blob.upload_from_file(file)
blob.make_public()
if public:
blob.make_public()

base_path, _ = os.path.splitext(source_file_path)
extracted_files_path = os.path.join(base_path, "extracted")
extracted_files: List[Gtfsfile] = []
Comment thread
davidgamez marked this conversation as resolved.
if not os.path.exists(extracted_files_path):
if not extracted_files_path or not os.path.exists(extracted_files_path):
Comment thread
davidgamez marked this conversation as resolved.
self.logger.warning(
f"Extracted files path {extracted_files_path} does not exist."
)
Expand All @@ -170,7 +167,8 @@ def upload_file_to_storage(self, source_file_path, dataset_stable_id):
f"{self.feed_stable_id}/{dataset_stable_id}/extracted/{file_name}"
)
file_blob.upload_from_filename(file_path)
file_blob.make_public()
if public:
file_blob.make_public()
self.logger.info(
f"Uploaded extracted file {file_name} to {file_blob.public_url}"
)
Expand All @@ -183,7 +181,7 @@ def upload_file_to_storage(self, source_file_path, dataset_stable_id):
)
return blob, extracted_files

def upload_dataset(self) -> DatasetFile or None:
def upload_dataset(self, public=True) -> DatasetFile or None:
"""
Uploads a dataset to a GCP bucket as <feed_stable_id>/latest.zip and
<feed_stable_id>/<feed_stable_id>-<upload_datetime>.zip
Expand All @@ -203,12 +201,12 @@ def upload_dataset(self) -> DatasetFile or None:
self.logger.info(
f"[{self.feed_stable_id}] File hash is {file_sha256_hash}."
)

if self.latest_hash != file_sha256_hash:
self.logger.info(
f"[{self.feed_stable_id}] Dataset has changed (hash {self.latest_hash}"
f"-> {file_sha256_hash}). Uploading new version."
)
extracted_files_path = self.unzip_files(temp_file_path)
self.logger.info(
f"Creating file {self.feed_stable_id}/latest.zip in bucket {self.bucket_name}"
)
Expand All @@ -224,7 +222,10 @@ def upload_dataset(self) -> DatasetFile or None:
f" in bucket {self.bucket_name}"
)
_, extracted_files = self.upload_file_to_storage(
temp_file_path, dataset_stable_id
temp_file_path,
dataset_stable_id,
extracted_files_path,
public=public,
)

return DatasetFile(
Expand All @@ -246,6 +247,18 @@ def upload_dataset(self) -> DatasetFile or None:
os.remove(temp_file_path)
return None

def unzip_files(self, temp_file_path):
extracted_files_path = os.path.join(temp_file_path.split(".")[0], "extracted")
self.logger.info(f"Unzipping files to {extracted_files_path}")
# Create the directory for extracted files if it does not exist
os.makedirs(extracted_files_path, exist_ok=True)
with zipfile.ZipFile(temp_file_path, "r") as zip_ref:
zip_ref.extractall(path=extracted_files_path)
# List all files in the extracted directory
extracted_files = os.listdir(extracted_files_path)
self.logger.info(f"Extracted files: {extracted_files}")
return extracted_files_path

def generate_temp_filename(self):
"""
Generates a temporary filename
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
import logging
Comment thread
davidgamez marked this conversation as resolved.
import os

from main import DatasetProcessor
from gcp_storage_emulator.server import create_server

HOST = "localhost"
PORT = 9023
BUCKET_NAME = "verifier"
PRODUCER_URL = "https://example.com/dataset.zip" # Replace with actual producer URL


def verify_download_content(producer_url: str):
"""
Verifies the download_content is able to retrieve the file
This is useful to simulate the download code locally and test issues related with user-agent and downloaded content.
Not supported authenticated feeds currently.
"""
logging.info("Verifying downloaded content... (not implemented)")

logging.info(f"Producer URL: {producer_url}")

processor = DatasetProcessor(
producer_url=producer_url,
feed_id=None,
feed_stable_id=None,
execution_id=None,
latest_hash=None,
bucket_name=None,
authentication_type=0,
api_key_parameter_name=None,
public_hosted_datasets_url=None,
)
tempfile = processor.generate_temp_filename()
logging.info(f"Temp filename: {tempfile}")
file_hash, is_zip = processor.download_content(tempfile)
logging.info(f"File hash: {file_hash}")


def verify_upload_dataset(producer_url: str):
"""
Verifies the upload_dataset is able to upload the dataset to the GCP storage emulator.
This is useful to simulate the upload code locally and test issues related with user-agent and uploaded content.
This function also tests the DatasetProcessor class methods for generating a temporary filename
and uploading the dataset.
:param producer_url:
:return:
"""
processor = DatasetProcessor(
producer_url=producer_url,
feed_id="feed_id",
feed_stable_id="feed_stable_id",
execution_id=None,
latest_hash="123",
bucket_name=BUCKET_NAME,
authentication_type=0,
api_key_parameter_name=None,
public_hosted_datasets_url=None,
)
tempfile = processor.generate_temp_filename()
logging.info(f"Temp filename: {tempfile}")
dataset_file = processor.upload_dataset(public=False)
logging.info(f"Dataset File: {dataset_file}")


if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
# Replace with actual producer URL
try:
os.environ["STORAGE_EMULATOR_HOST"] = f"http://{HOST}:{PORT}"
server = create_server(
host=HOST, port=PORT, in_memory=False, default_bucket=BUCKET_NAME
)
server.start()

verify_download_content(producer_url=PRODUCER_URL)
logging.info("Download content verification completed successfully.")
verify_upload_dataset(producer_url=PRODUCER_URL)
verify_upload_dataset(producer_url=PRODUCER_URL)
except Exception as e:
logging.error(f"Error verifying download content: {e}")
finally:
server.stop()
logging.info("Verification completed.")
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,9 @@ def create_cloud_event(mock_data):
class TestDatasetProcessor(unittest.TestCase):
@patch("main.DatasetProcessor.upload_file_to_storage")
@patch("main.DatasetProcessor.download_content")
@patch("main.DatasetProcessor.unzip_files")
def test_upload_dataset_diff_hash(
self, mock_download_url_content, upload_file_to_storage
self, mock_unzip_files, mock_download_url_content, upload_file_to_storage
):
"""
Test upload_dataset method of DatasetProcessor class with different hash from the latest one
Expand All @@ -57,6 +58,7 @@ def test_upload_dataset_diff_hash(
mock_blob.path = public_url
upload_file_to_storage.return_value = mock_blob, []
mock_download_url_content.return_value = file_hash, True
mock_unzip_files.return_value = [mock_blob, mock_blob]

processor = DatasetProcessor(
public_url,
Expand Down Expand Up @@ -178,6 +180,7 @@ def test_upload_dataset_download_exception(
def test_upload_file_to_storage(self):
bucket_name = "test-bucket"
source_file_path = "path/to/source/file"
extracted_file_path = "path/to/source/file"

mock_blob = Mock()
mock_blob.public_url = public_url
Expand All @@ -204,7 +207,9 @@ def test_upload_file_to_storage(self):
test_hosted_public_url,
)
dataset_id = faker.Faker().uuid4()
result, _ = processor.upload_file_to_storage(source_file_path, dataset_id)
result, _ = processor.upload_file_to_storage(
source_file_path, dataset_id, extracted_file_path
)
self.assertEqual(result.public_url, public_url)
mock_client.get_bucket.assert_called_with(bucket_name)
mock_bucket.blob.assert_called_with(
Expand Down
19 changes: 13 additions & 6 deletions functions-python/helpers/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ def filter(self, record):


lock = threading.Lock()
lock_logger = threading.Lock()
_logging_initialized = False


Expand Down Expand Up @@ -73,9 +74,15 @@ def get_logger(name: str, stable_id: str = None):
If stable_id is provided, the StableIdFilter is added.
This method can be called multiple times for the same logger name without creating a side effect.
"""
logger = logging.getLogger(name)
if stable_id and not any(
isinstance(log_filter, StableIdFilter) for log_filter in logger.filters
):
logger.addFilter(StableIdFilter(stable_id))
return logger
with lock_logger:
Comment thread
davidgamez marked this conversation as resolved.
# Create the logger with the provided name to avoid retuning the same logger instance
logger = (
logging.getLogger(f"{name}_{stable_id}")
if stable_id
else logging.getLogger(name)
)
if stable_id and not any(
isinstance(log_filter, StableIdFilter) for log_filter in logger.filters
):
logger.addFilter(StableIdFilter(stable_id))
return logger
Loading