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
48 changes: 48 additions & 0 deletions api/src/shared/common/config_reader.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
from typing import Optional, Any
from sqlalchemy.orm import Session
from shared.database.database import with_db_session
from shared.database_gen.sqlacodegen_models import ConfigKey, ConfigValueFeed


@with_db_session
def get_config_value(
namespace: str,
key: str,
feed_id: Optional[str] = None,
db_session: Session = None,
) -> Optional[Any]:
"""
Retrieves a configuration value from the database using the provided session.

It first looks for a feed-specific override. If a feed_stable_id is provided and a
value is found, it returns that value.

If no feed-specific value is found or no feed_stable_id is provided, it looks for
the global default value in the `config_key` table.

:param namespace: The namespace of the configuration key.
:param key: The configuration key.
:param feed_stable_id: The optional feed_stable_id for a specific override.
:param db_session: The SQLAlchemy session, injected by the `with_db_session` decorator.
:return: The configuration value, or None if not found.
"""
# 1. Try to get feed-specific value if feed_id is provided
if feed_id:
feed_config = (
db_session.query(ConfigValueFeed.value)
.filter(
ConfigValueFeed.feed_id == feed_id,
ConfigValueFeed.namespace == namespace,
ConfigValueFeed.key == key,
)
.first()
)
if feed_config:
return feed_config.value

# 2. If not found or no feed_id, get the default value
default_config = (
db_session.query(ConfigKey.default_value).filter(ConfigKey.namespace == namespace, ConfigKey.key == key).first()
)

return default_config.default_value if default_config else None
18 changes: 11 additions & 7 deletions functions-python/batch_process_dataset/src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,18 +122,20 @@ def create_dataset_stable_id(feed_stable_id, timestamp):
"""
return f"{feed_stable_id}-{timestamp}"

def download_content(self, temporary_file_path):
def download_content(self, temporary_file_path, feed_id):
"""
Downloads the content of a URL and return the hash of the file
"""
file_hash = download_and_get_hash(
self.producer_url,
temporary_file_path,
file_path=temporary_file_path,
feed_id=feed_id,
authentication_type=self.authentication_type,
api_key_parameter_name=self.api_key_parameter_name,
credentials=self.feed_credentials,
logger=self.logger,
)
self.logger.info(f"hash is: {file_hash}")
is_zip = zipfile.is_zipfile(temporary_file_path)
return file_hash, is_zip

Expand Down Expand Up @@ -193,7 +195,7 @@ def upload_files_to_storage(
)
return blob, extracted_files

def upload_dataset(self, public=True) -> DatasetFile or None:
def upload_dataset(self, feed_id, 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 @@ -204,7 +206,7 @@ def upload_dataset(self, public=True) -> DatasetFile or None:
try:
self.logger.info("Accessing URL %s", self.producer_url)
temp_file_path = self.generate_temp_filename()
file_sha256_hash, is_zip = self.download_content(temp_file_path)
file_sha256_hash, is_zip = self.download_content(temp_file_path, feed_id)
if not is_zip:
self.logger.error(
f"[{self.feed_stable_id}] The downloaded file from {self.producer_url} is not a valid ZIP file."
Expand Down Expand Up @@ -417,12 +419,14 @@ def _get_unzipped_size(dataset_file):
)

@with_db_session
def process_from_producer_url(self, db_session) -> Optional[DatasetFile]:
def process_from_producer_url(
self, feed_id, db_session: Session
) -> Optional[DatasetFile]:
"""
Process the dataset and store new version in GCP bucket if any changes are detected
:return: the DatasetFile object created
"""
dataset_file = self.upload_dataset()
dataset_file = self.upload_dataset(feed_id)

if dataset_file is None:
self.logger.info(f"[{self.feed_stable_id}] No database update required.")
Expand Down Expand Up @@ -543,7 +547,7 @@ def process_dataset(cloud_event: CloudEvent):
if json_payload.get("use_bucket_latest", False):
dataset_file = processor.process_from_bucket()
else:
dataset_file = processor.process_from_producer_url()
dataset_file = processor.process_from_producer_url(json_payload["feed_id"])
except Exception as e:
# This makes sure the logger is initialized
logger = get_logger("process_dataset", stable_id if stable_id else "UNKNOWN")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,8 @@ def verify_download_content(producer_url: str):
)
tempfile = processor.generate_temp_filename()
logging.info(f"Temp filename: {tempfile}")
file_hash, is_zip = processor.download_content(tempfile)
file_hash, is_zip = processor.download_content(tempfile, "feed_id")
logging.info(f"Downloaded file from {producer_url} is a valid ZIP file: {is_zip}")
logging.info(f"File hash: {file_hash}")


Expand All @@ -48,7 +49,7 @@ def verify_upload_dataset(producer_url: str):
"""
processor = DatasetProcessor(
producer_url=producer_url,
feed_id="feed_id",
feed_id="feed_id_2126",
feed_stable_id="feed_stable_id",
execution_id=None,
latest_hash="123",
Expand All @@ -59,7 +60,7 @@ def verify_upload_dataset(producer_url: str):
)
tempfile = processor.generate_temp_filename()
logging.info(f"Temp filename: {tempfile}")
dataset_file = processor.upload_dataset(public=False)
dataset_file = processor.upload_dataset("feed_id_2126", False)
logging.info(f"Dataset File: {dataset_file}")


Expand All @@ -72,6 +73,7 @@ def verify_upload_dataset(producer_url: str):
# create working dir if not exists
if not os.path.exists(os.environ["WORKING_DIR"]):
os.makedirs(os.environ["WORKING_DIR"])

server = create_server(
host=HOST, port=PORT, in_memory=False, default_bucket=BUCKET_NAME
)
Expand All @@ -80,7 +82,6 @@ def verify_upload_dataset(producer_url: str):
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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def test_upload_dataset_diff_hash(
test_hosted_public_url,
)
with patch.object(processor, "date", "mocked_timestamp"):
result = processor.upload_dataset()
result = processor.upload_dataset("feed_id")

self.assertIsNotNone(result)
mock_download_url_content.assert_called_once()
Expand Down Expand Up @@ -111,7 +111,7 @@ def test_upload_dataset_same_hash(
test_hosted_public_url,
)

result = processor.upload_dataset()
result = processor.upload_dataset("feed_id")

self.assertIsNone(result)
upload_files_to_storage.blob.assert_not_called()
Expand Down Expand Up @@ -143,7 +143,7 @@ def test_upload_dataset_not_zip(
test_hosted_public_url,
)

result = processor.upload_dataset()
result = processor.upload_dataset("feed_id")

self.assertIsNone(result)
upload_files_to_storage.blob.assert_not_called()
Expand Down Expand Up @@ -176,7 +176,7 @@ def test_upload_dataset_download_exception(
)

with self.assertRaises(Exception):
processor.upload_dataset()
processor.upload_dataset("feed_id")

def test_upload_files_to_storage(self):
bucket_name = "test-bucket"
Expand Down Expand Up @@ -256,7 +256,7 @@ def test_process(self, db_session):
)
db_url = os.getenv("TEST_FEEDS_DATABASE_URL", default=default_db_url)
os.environ["FEEDS_DATABASE_URL"] = db_url
result = processor.process_from_producer_url()
result = processor.process_from_producer_url(feed_id)

self.assertIsNotNone(result)
self.assertEqual(result.file_sha256_hash, new_hash)
Expand Down Expand Up @@ -357,7 +357,7 @@ def test_process_no_change(self):

processor.upload_dataset = MagicMock(return_value=None)
processor.create_dataset_entities = MagicMock()
result = processor.process_from_producer_url()
result = processor.process_from_producer_url(feed_id)

self.assertIsNone(result)
processor.create_dataset_entities.assert_not_called()
Expand Down
30 changes: 18 additions & 12 deletions functions-python/helpers/tests/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ def test_create_bucket_already_exists(self, mock_storage_client):
)
mock_storage_client.return_value.create_bucket.assert_not_called()

def test_download_and_get_hash(self):
@patch("shared.common.config_reader.get_config_value", return_value=None)
def test_download_and_get_hash(self, mock_get_config):
mock_binary_data = b"file content data"
expected_hash = hashlib.sha256(mock_binary_data).hexdigest()
file_path = "file_path"
Expand All @@ -65,7 +66,8 @@ def test_download_and_get_hash(self):
if os.path.exists(file_path):
os.remove(file_path)

def test_download_and_get_hash_auth_type_header(self):
@patch("shared.common.config_reader.get_config_value", return_value=None)
def test_download_and_get_hash_auth_type_header(self, mock_get_config):
"""
Test the download_and_get_hash function for authentication type 2 (headers).
This test verifies that the download_and_get_hash function correctly handles authentication type 2,
Expand All @@ -88,7 +90,11 @@ def test_download_and_get_hash_auth_type_header(self):
"urllib3.PoolManager.request", return_value=mock_response
) as mock_request:
result_hash = download_and_get_hash(
url, file_path, "sha256", 8192, 2, api_key_parameter_name, credentials
url=url,
file_path=file_path,
authentication_type=2,
api_key_parameter_name=api_key_parameter_name,
credentials=credentials,
)

self.assertEqual(
Expand All @@ -113,7 +119,8 @@ def test_download_and_get_hash_auth_type_header(self):
if os.path.exists(file_path):
os.remove(file_path)

def test_download_and_get_hash_auth_type_api_key(self):
@patch("shared.common.config_reader.get_config_value", return_value=None)
def test_download_and_get_hash_auth_type_api_key(self, mock_get_config):
"""
Test the download_and_get_hash function for authentication type 1 (API key).
"""
Expand All @@ -137,13 +144,11 @@ def test_download_and_get_hash_auth_type_api_key(self):

with patch("urllib3.PoolManager", return_value=mock_http):
result_hash = download_and_get_hash(
base_url,
file_path,
"sha256",
8192,
1,
api_key_parameter_name,
credentials,
url=base_url,
file_path=file_path,
authentication_type=1,
api_key_parameter_name=api_key_parameter_name,
credentials=credentials,
)

self.assertEqual(
Expand All @@ -163,7 +168,8 @@ def test_download_and_get_hash_auth_type_api_key(self):
if os.path.exists(file_path):
os.remove(file_path)

def test_download_and_get_hash_exception(self):
@patch("shared.common.config_reader.get_config_value", return_value=None)
def test_download_and_get_hash_exception(self, mock_get_config):
file_path = "test_file.txt"
url = "https://test.com/"

Expand Down
20 changes: 14 additions & 6 deletions functions-python/helpers/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,12 +126,15 @@ def download_and_get_hash(
file_path,
hash_algorithm="sha256",
chunk_size=8192,
feed_id=None,
authentication_type=0,
api_key_parameter_name=None,
credentials=None,
logger=None,
trusted_certs=False, # If True, disables SSL verification
):
from shared.common.config_reader import get_config_value

"""
Downloads the content of a URL and stores it in a file and returns the hash of the file
"""
Expand All @@ -145,13 +148,18 @@ def download_and_get_hash(
ctx.load_default_certs()
ctx.options |= 0x4 # ssl.OP_LEGACY_SERVER_CONNECT

headers = get_config_value(
namespace="feed_download", key="http_headers", feed_id=feed_id
)
if headers is None:
headers = {
"User-Agent": "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/126.0.0.0 Mobile Safari/537.36",
"Referer": url,
}

Comment on lines +151 to +161

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

# authentication_type == 1 -> the credentials are passed in the url
headers = {
"User-Agent": "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/126.0.0.0 Mobile Safari/537.36",
"Referer": url,
}
# Careful, some URLs may already contain a query string
# (e.g. http://api.511.org/transit/datafeeds?operator_id=CE)
if authentication_type == 1 and api_key_parameter_name and credentials:
Expand Down
1 change: 1 addition & 0 deletions liquibase/changelog.xml
Original file line number Diff line number Diff line change
Expand Up @@ -69,4 +69,5 @@
<include file="changes/feat_1333.sql" relativeToChangelogFile="true"/>
<include file="changes/feat_pt_152.sql" relativeToChangelogFile="true"/>
<include file="changes/feat_fix_geolocation_circular_dep.sql" relativeToChangelogFile="true"/>
<include file="changes/feat_1325.sql" relativeToChangelogFile="true"/>
</databaseChangeLog>
25 changes: 25 additions & 0 deletions liquibase/changes/feat_1325.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
ALTER TABLE feed ADD CONSTRAINT feed_stable_id_unique UNIQUE (stable_id);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before merging this PR, we must fix => #1061

-- 1. Catalog of allowed keys + optional global values
CREATE TABLE config_key (
namespace text NOT NULL, -- e.g. 'reverse_geolocation'
key text NOT NULL, -- e.g. 'country_fallback'
description text,
default_value jsonb, -- fallback if nothing else
updated_at timestamptz NOT NULL DEFAULT now(),
PRIMARY KEY (namespace, key)
);

-- 2. Per-feed overrides
CREATE TABLE config_value_feed (
feed_id varchar(255) NOT NULL,
feed_stable_id varchar(255) NOT NULL,
namespace text NOT NULL,
key text NOT NULL,
value jsonb NOT NULL,
updated_at timestamptz NOT NULL DEFAULT now(),
PRIMARY KEY (feed_id, namespace, key),
FOREIGN KEY (namespace, key) REFERENCES config_key(namespace, key) ON DELETE CASCADE
);

-- Helpful index
CREATE INDEX config_value_feed_gin ON config_value_feed USING GIN (value);
Loading