use_multi_tenancy_endpoints: use datastore endpoint for directory (#207)

Author: DanielElisenberg

job_executor/adapter/datastore_api/__init__.py

@@ -8,6 +8,7 @@
 from urllib3 import Retry
 
 from job_executor.adapter.datastore_api.models import (
+    DatastoreResponse,
     Job,
     JobQueryResult,
     JobStatus,
@@ -72,10 +73,7 @@ def get_jobs(
         request_url += f"?{'&'.join(query_fields)}"
 
     response = execute_request("GET", request_url, True)
-    return [
-        Job.model_validate({"datastoreRdn": environment.datastore_rdn, **job})
-        for job in response.json()
-    ]
+    return [Job.model_validate(job) for job in response.json()]
 
 
 def update_job_status(
@@ -108,7 +106,9 @@ def is_system_paused() -> bool:
 
 
 def get_datastore_directory(rdn: str) -> Path:
-    return Path(environment.datastore_dir)
+    request_url = f"{DATASTORE_API_URL}/datastores/{rdn}"
+    response = execute_request("GET", request_url, True)
+    return Path(DatastoreResponse.model_validate(response.json()).directory)
 
 
 def query_for_jobs() -> JobQueryResult:

job_executor/adapter/datastore_api/models.py

@@ -9,6 +9,15 @@
 from job_executor.common.models import CamelModel
 
 
+class DatastoreResponse(CamelModel):
+    datastore_id: int
+    name: str
+    rdn: str
+    description: str
+    directory: str
+    bump_enabled: bool
+
+
 class MaintenanceStatus(CamelModel):
     paused: bool
     msg: str

tests/unit/worker/test_build_dataset_worker.py

@@ -91,6 +91,11 @@
 }
 
 EXPECTED_REQUESTS_PARTITIONED = [
+    {
+        "json": None,
+        "method": "GET",
+        "url": f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+    },
     {
         "json": {"status": "decrypting"},
         "method": "PUT",
@@ -186,6 +191,11 @@
 ]
 
 EXPECTED_REQUESTS_IMPORT = [
+    {
+        "json": None,
+        "method": "GET",
+        "url": f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+    },
     {
         "json": {"status": "decrypting"},
         "method": "PUT",
@@ -247,6 +257,14 @@
     and "mock.pseudonym.service" not in request["url"]
 ]
 JOB = SimpleNamespace(job_id=JOB_ID, datastore_rdn=DATASTORE_RDN)
+DATASTORE_RESPONSE = {
+    "datastore_id": 1,
+    "rdn": DATASTORE_RDN,
+    "name": "",
+    "description": "",
+    "bump_enabled": True,
+    "directory": DATASTORE_DIR,
+}
 
 
 def setup_function():
@@ -273,6 +291,10 @@ def test_build_partitioned_dataset(requests_mock: RequestsMocker):
     requests_mock.put(
         f"{DATASTORE_API_URL}/jobs/{JOB_ID}", json={"message": "OK"}
     )
+    requests_mock.get(
+        f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+        json=DATASTORE_RESPONSE,
+    )
     requests_mock.post(
         f"{PSEUDONYM_SERVICE_URL}?unit_id_type=FNR&job_id={JOB_ID}",
         json=PSEUDONYM_DICT,
@@ -284,7 +306,11 @@ def test_build_partitioned_dataset(requests_mock: RequestsMocker):
         f"{WORKING_DIR}/{PARTITIONED_DATASET_NAME}__DRAFT.json"
     )
     requests_made = [
-        {"method": req.method, "json": req.json(), "url": req.url}
+        {
+            "method": req.method,
+            "json": req.json() if req.method != "GET" else None,
+            "url": req.url,
+        }
         for req in requests_mock.request_history
     ]
     assert len(requests_made) == len(EXPECTED_REQUESTS_PARTITIONED)
@@ -298,6 +324,10 @@ def test_import(requests_mock: RequestsMocker):
     requests_mock.put(
         f"{DATASTORE_API_URL}/jobs/{JOB_ID}", json={"message": "OK"}
     )
+    requests_mock.get(
+        f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+        json=DATASTORE_RESPONSE,
+    )
     requests_mock.post(
         f"{PSEUDONYM_SERVICE_URL}?unit_id_type=FNR&job_id={JOB_ID}",
         json=PSEUDONYM_DICT,
@@ -315,7 +345,11 @@ def test_import(requests_mock: RequestsMocker):
         Path(INPUT_DIR_ARCHIVE) / f"unpackaged/{DATASET_NAME}.tar"
     ).exists()
     requests_made = [
-        {"method": req.method, "json": req.json(), "url": req.url}
+        {
+            "method": req.method,
+            "json": req.json() if req.method != "GET" else None,
+            "url": req.url,
+        }
         for req in requests_mock.request_history
     ]
     assert len(requests_made) == len(EXPECTED_REQUESTS_IMPORT)
@@ -329,6 +363,10 @@ def test_import_no_pseudonymization(requests_mock: RequestsMocker):
     requests_mock.put(
         f"{DATASTORE_API_URL}/jobs/{JOB_ID}", json={"message": "OK"}
     )
+    requests_mock.get(
+        f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+        json=DATASTORE_RESPONSE,
+    )
     run_worker(JOB, NO_PSEUDONYM_DATASET_NAME, Queue())  # type: ignore
     assert not os.path.exists(f"{INPUT_DIR}/{NO_PSEUDONYM_DATASET_NAME}.tar")
     assert not os.path.exists(f"{INPUT_DIR}/{NO_PSEUDONYM_DATASET_NAME}")
@@ -343,7 +381,11 @@ def test_import_no_pseudonymization(requests_mock: RequestsMocker):
         Path(INPUT_DIR_ARCHIVE) / f"unpackaged/{NO_PSEUDONYM_DATASET_NAME}.tar"
     ).exists()
     requests_made = [
-        {"method": req.method, "json": req.json(), "url": req.url}
+        {
+            "method": req.method,
+            "json": req.json() if req.method != "GET" else None,
+            "url": req.url,
+        }
         for req in requests_mock.request_history
     ]
     assert len(requests_made) == len(EXPECTED_REQUESTS_IMPORT_NO_PSEUDONYM)
@@ -359,6 +401,10 @@ def test_import_no_pseudonymization_no_partitioning(
     requests_mock.put(
         f"{DATASTORE_API_URL}/jobs/{JOB_ID}", json={"message": "OK"}
     )
+    requests_mock.get(
+        f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+        json=DATASTORE_RESPONSE,
+    )
     run_worker(JOB, NO_PSEUDONYM_FIXED_DATASET_NAME, Queue())  # type: ignore
     assert not os.path.exists(
         f"{INPUT_DIR}/{NO_PSEUDONYM_FIXED_DATASET_NAME}.tar"
@@ -384,7 +430,11 @@ def test_import_no_pseudonymization_no_partitioning(
         / f"unpackaged/{NO_PSEUDONYM_FIXED_DATASET_NAME}.tar"
     ).exists()
     requests_made = [
-        {"method": req.method, "json": req.json(), "url": req.url}
+        {
+            "method": req.method,
+            "json": req.json() if req.method != "GET" else None,
+            "url": req.url,
+        }
         for req in requests_mock.request_history
     ]
     assert len(requests_made) == len(EXPECTED_REQUESTS_IMPORT_NO_PSEUDONYM)

tests/unit/worker/test_build_metadata_worker.py

@@ -15,6 +15,7 @@
 RSA_KEYS_DIRECTORY = Path(environment.datastore_dir) / "vault"
 
 DATASET_NAME = "KJOENN"
+DATASTORE_RDN = os.environ["DATASTORE_RDN"]
 JOB_ID = "1234-1234-1234-1234"
 DATASTORE_DIR = os.environ["DATASTORE_DIR"]
 WORKING_DIR = DATASTORE_DIR + "_working"
@@ -23,6 +24,11 @@
 EXPECTED_DIR = "tests/resources/worker/build_metadata/expected"
 DATASTORE_API_URL = os.environ["DATASTORE_API_URL"]
 EXPECTED_REQUESTS = [
+    {
+        "json": None,
+        "method": "GET",
+        "url": f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+    },
     {
         "json": {"status": "decrypting"},
         "method": "PUT",
@@ -49,8 +55,15 @@
         "url": f"{DATASTORE_API_URL}/jobs/{JOB_ID}",
     },
 ]
-DATASTORE_RDN = os.environ["DATASTORE_RDN"]
 JOB = SimpleNamespace(job_id=JOB_ID, datastore_rdn=DATASTORE_RDN)
+DATASTORE_RESPONSE = {
+    "datastore_id": 1,
+    "rdn": DATASTORE_RDN,
+    "name": "",
+    "description": "",
+    "bump_enabled": True,
+    "directory": DATASTORE_DIR,
+}
 
 
 def setup_function():
@@ -77,6 +90,10 @@ def test_import(requests_mock: RequestsMocker):
     requests_mock.put(
         f"{DATASTORE_API_URL}/jobs/{JOB_ID}", json={"message": "OK"}
     )
+    requests_mock.get(
+        f"{DATASTORE_API_URL}/datastores/{DATASTORE_RDN}",
+        json=DATASTORE_RESPONSE,
+    )
     run_worker(JOB, DATASET_NAME, Queue())  # type: ignore
     with open(
         f"{WORKING_DIR}/{DATASET_NAME}__DRAFT.json", "r", encoding="utf-8"
@@ -89,7 +106,11 @@ def test_import(requests_mock: RequestsMocker):
 
     assert actual_metadata == expected_metadata
     requests_made = [
-        {"method": req.method, "json": req.json(), "url": req.url}
+        {
+            "method": req.method,
+            "json": req.json() if req.method != "GET" else None,
+            "url": req.url,
+        }
         for req in requests_mock.request_history
     ]
     assert requests_made == EXPECTED_REQUESTS