Skip to content

Commit bb617e7

Browse files
committed
refactor: decoupled drive extractor orchestration and logic; update docstrings to design-contract
1 parent deab594 commit bb617e7

8 files changed

Lines changed: 265 additions & 143 deletions

File tree

data_extract/run_extract.py

Lines changed: 54 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -1,152 +1,105 @@
11
# =============================================================================
2-
# GOOGLE DRIVE FILE EXTRACTOR
2+
# Google Drive Extractor Orchestrator
33
# =============================================================================
44

55
import sys
66
import json
77
import uuid
88
from data_extract.shared.utils import (
9-
extract_file_content,
10-
check_handshake,
119
get_drive_service,
1210
check_gcs_marking,
1311
upload_to_gcs,
1412
get_target_folder_name,
1513
)
14+
from data_extract.shared.extract_logic import (
15+
ARCHIVAL_BUCKET,
16+
get_target_folder_id,
17+
get_valid_files,
18+
process_extraction,
19+
)
20+
21+
22+
def orchestrate_extract(target_child_folder: str) -> int:
23+
"""
24+
Main orchestrator for the Google Drive to GCS ingestion lifecycle.
1625
17-
ARCHIVAL_BUCKET = "gs://operations-archival-bucket"
18-
PIPELINE_BUCKET = "gs://operations-pipeline-bucket"
19-
PARENT_FOLDER = "operations-upload-folder"
20-
MIME_TYPE = "application/vnd.google-apps.folder"
26+
Contract:
27+
- Resolves the specific target folder ID within the 'PARENT_FOLDER' hierarchy.
28+
- Enforces a deduplication check via GCS success markers to prevent redundant extraction.
29+
- Iteratively processes file extraction, archival, and pipeline mirroring.
2130
31+
Invariants:
32+
- Atomicity: A run is marked as [SUCCESS] only if every file in the target
33+
folder is successfully processed and mirrored.
34+
- Idempotency: Uses GCS marking paths to skip previously completed folders.
35+
- Lineage: Generates a unique 'execution_id' (UUID4) for each orchestration attempt.
2236
23-
# TODO: wrap steps into functions
24-
def run_extraction(target_child_folder):
37+
Failures:
38+
- Returns 1 if any file extraction fails, the target folder is missing,
39+
or the handshake is invalid.
40+
- Returns 0 on successful completion or if a deduplication skip is triggered.
41+
"""
2542

2643
service = get_drive_service()
2744
metadata_path = f"logs/{target_child_folder}_metadata.json"
2845
archival_marking_path = f"status/{target_child_folder}.success"
2946

30-
# Root Folder
31-
parent_query = f"name = '{PARENT_FOLDER}' and mimeType = '{MIME_TYPE}'"
32-
parent_results = service.files().list(q=parent_query, fields="files(id)").execute()
33-
parents = parent_results.get("files", [])
34-
35-
if not parents:
36-
print(f"[ERROR]: Parent folder '{PARENT_FOLDER}' not found or not shared.")
37-
38-
return 1
39-
40-
parent_id = parents[0]["id"]
41-
42-
# Uploaded folder with data
43-
child_query = f"name = '{target_child_folder}' and '{parent_id}' in parents and mimeType = '{MIME_TYPE}'"
44-
child_results = service.files().list(q=child_query, fields="files(id)").execute()
45-
children = child_results.get("files", [])
46-
47-
if not children:
48-
print(
49-
f"[ERROR]: Dated folder {target_child_folder} not found inside {PARENT_FOLDER}."
50-
)
51-
52-
return 1
53-
54-
folder_id = children[0]["id"]
47+
metadata = {
48+
"execution_id": str(uuid.uuid4()),
49+
"files_processed": [],
50+
"errors": [],
51+
"status": "success",
52+
}
5553

5654
# Deduplication Check
5755
if check_gcs_marking(ARCHIVAL_BUCKET, archival_marking_path):
5856
print(f"[INFO]: {target_child_folder} already processed.")
59-
6057
return 0
6158

62-
# Find Folder ID by Name
63-
query = f"name = '{target_child_folder}' and mimeType = '{MIME_TYPE}'"
64-
folder_results = service.files().list(q=query, fields="files(id)").execute()
65-
folders = folder_results.get("files", [])
66-
67-
if not folders:
68-
print(f"[ERROR]: Folder {target_child_folder} not found.")
69-
59+
# Extract target folder id
60+
folder_id = get_target_folder_id(target_child_folder, service)
61+
if not folder_id:
7062
return 1
7163

72-
folder_id = folders[0]["id"]
73-
74-
# Validate upload instruction
75-
if not check_handshake(service, folder_id):
76-
print(
77-
f"[ERROR]: {target_child_folder} missing instruction.txt or upload not safe."
78-
)
64+
# Extract files
65+
files_in_drive = get_valid_files(folder_id, target_child_folder, service)
7966

67+
# Exit if handshake failed or empty list
68+
if files_in_drive is None or len(files_in_drive) == 0:
8069
return 1
8170

82-
# List files in folder
83-
file_results = (
84-
service.files()
85-
.list(
86-
q=f"'{folder_id}' in parents and name != 'instruction.txt'",
87-
fields="files(id, name, mimeType)",
88-
)
89-
.execute()
90-
)
91-
92-
files_in_drive = file_results.get("files", [])
93-
94-
metadata = {
95-
"execution_id": str(uuid.uuid4()),
96-
"files_processed": [],
97-
"errors": [],
98-
"status": "success",
99-
}
100-
10171
for file in files_in_drive:
10272

103-
archival_path = f"archive/{target_child_folder}/{file['name']}.csv"
104-
pipeline_raw_path = f"raw/{file['name']}.csv"
105-
106-
try:
107-
data = extract_file_content(service, file["id"], file["mimeType"])
108-
109-
# upload directory to archival
110-
upload_to_gcs(ARCHIVAL_BUCKET, archival_path, data)
111-
112-
# Upload on pipeline/raw/
113-
upload_to_gcs(PIPELINE_BUCKET, pipeline_raw_path, data)
114-
115-
metadata["files_processed"].append(
116-
{"name": file["name"], "status": "success"}
117-
)
118-
119-
print(f"[INFO]: {target_child_folder} processed successfully.")
73+
archival_path = f"archive/{target_child_folder}/{file['name']}"
74+
pipeline_raw_path = f"raw/{file['name']}"
12075

121-
except Exception as e:
122-
123-
error_details = {
124-
"file_name": file["name"],
125-
"drive_id": file["id"],
126-
"error_type": type(e).__name__,
127-
"error_message": str(e),
128-
}
76+
ok, details = process_extraction(
77+
file, service, archival_path, pipeline_raw_path
78+
)
12979

130-
metadata["errors"].append(error_details)
80+
if ok:
81+
metadata["files_processed"].append(details)
82+
else:
83+
metadata["errors"].append(details)
13184
metadata["status"] = "failed"
13285

133-
# runtime metadata
86+
# Upload failure metadata
13487
upload_to_gcs(ARCHIVAL_BUCKET, metadata_path, json.dumps(metadata))
135-
136-
print(f"[ERROR]: Execution halted {str(e)}")
13788
return 1
13889

90+
# Upload success metadata and marker
13991
upload_to_gcs(ARCHIVAL_BUCKET, metadata_path, json.dumps(metadata))
14092
upload_to_gcs(ARCHIVAL_BUCKET, archival_marking_path, "")
93+
94+
print(f"[SUCCESS]: Folder '{target_child_folder}' completely processed.")
14195
return 0
14296

14397

14498
def main():
145-
14699
target_folder = get_target_folder_name("operations")
147-
print(f"[INFO]: Starting extraction for {target_folder}")
100+
print(f"[INFO]: Starting extraction for folder: {target_folder}")
148101

149-
exit_code = run_extraction(target_folder)
102+
exit_code = orchestrate_extract(target_folder)
150103
sys.exit(exit_code)
151104

152105

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
# =============================================================================
2+
# Google Drive Extractor Logic
3+
# =============================================================================
4+
5+
6+
from typing import List, Dict
7+
from data_extract.shared.utils import (
8+
extract_file_content,
9+
check_handshake,
10+
upload_to_gcs,
11+
GoogleDriveService,
12+
)
13+
14+
ARCHIVAL_BUCKET = "gs://operations-archival-bucket"
15+
PIPELINE_BUCKET = "gs://operations-pipeline-bucket"
16+
PARENT_FOLDER = "operations-upload-folder"
17+
MIME_TYPE = "application/vnd.google-apps.folder"
18+
19+
20+
def get_target_folder_id(
21+
target_folder: str, drive_api: GoogleDriveService
22+
) -> str | None:
23+
"""
24+
Resolves a specific Drive Folder ID within the strict system hierarchy.
25+
26+
Contract:
27+
- Locates the system root 'PARENT_FOLDER'.
28+
- Performs a scoped search for 'target_folder' strictly within that root.
29+
30+
Invariants:
31+
- Namespace Protection: Prevents global Drive collisions by enforcing the
32+
parent-child relationship.
33+
34+
Failures:
35+
- Returns None if the parent root or the target subfolder is missing/inaccessible.
36+
"""
37+
38+
# Find Root Folder
39+
parent_query = f"name = '{PARENT_FOLDER}' and mimeType = '{MIME_TYPE}'"
40+
parent_results = (
41+
drive_api.files().list(q=parent_query, fields="files(id)").execute()
42+
)
43+
parents = parent_results.get("files", [])
44+
45+
if not parents:
46+
print(f"[ERROR]: Parent folder '{PARENT_FOLDER}' not found or not shared.")
47+
return None
48+
49+
parent_id = parents[0]["id"]
50+
51+
# Find Target folder inside Parent
52+
child_query = f"name = '{target_folder}' and '{parent_id}' in parents and mimeType = '{MIME_TYPE}'"
53+
child_results = drive_api.files().list(q=child_query, fields="files(id)").execute()
54+
children = child_results.get("files", [])
55+
56+
if not children:
57+
print(f"[ERROR]: Folder '{target_folder}' not found inside '{PARENT_FOLDER}'.")
58+
return None
59+
60+
return children[0]["id"]
61+
62+
63+
def get_valid_files(
64+
folder_id: str, target_folder: str, drive_api: GoogleDriveService
65+
) -> List[Dict] | None:
66+
"""
67+
Retrieves the candidate file manifest from the target Drive folder.
68+
69+
Contract:
70+
- Fetches metadata (ID, Name, MimeType) for all files in the parent folder.
71+
- Filters out system-reserved files (e.g., 'instruction.txt').
72+
73+
Invariants:
74+
- Scope: Only direct children of 'folder_id' are returned (non-recursive).
75+
76+
Failures:
77+
- Returns an empty list if the folder is empty or contains only reserved files.
78+
"""
79+
80+
# Validate instruction inside the folder
81+
if not check_handshake(drive_api, folder_id):
82+
print(f"[ERROR]: '{target_folder}' missing instruction.txt or upload not safe.")
83+
return None
84+
85+
# List files inside folder (excluding instruction.txt)
86+
file_results = (
87+
drive_api.files()
88+
.list(
89+
q=f"'{folder_id}' in parents and name != 'instruction.txt'",
90+
fields="files(id, name, mimeType)",
91+
)
92+
.execute()
93+
)
94+
95+
files_in_drive = file_results.get("files", [])
96+
97+
# Warn if folder is valid but empty
98+
if not files_in_drive:
99+
print(f"[WARNING]: '{target_folder}' is empty. Nothing to process.")
100+
101+
return files_in_drive
102+
103+
104+
def process_extraction(
105+
file: dict,
106+
drive_api: GoogleDriveService,
107+
archival_path: str,
108+
pipeline_path: str,
109+
) -> tuple[bool, dict]:
110+
"""
111+
Executes the dual-mirroring extraction for a single file.
112+
113+
Contract:
114+
- Extracts binary content from Google Drive via the service provider.
115+
- Performs a synchronous dual-upload to the Archival and Pipeline GCS buckets.
116+
117+
Invariants:
118+
- Mirroring: Every successful extraction results in two identical GCS
119+
artifacts in separate administrative zones.
120+
- Integrity: The extraction fails if either GCS upload fails.
121+
122+
Returns:
123+
tuple: (Success Boolean, Metadata/Error Detail Dictionary).
124+
"""
125+
126+
try:
127+
data = extract_file_content(drive_api, file["id"], file["mimeType"])
128+
129+
upload_to_gcs(ARCHIVAL_BUCKET, archival_path, data)
130+
upload_to_gcs(PIPELINE_BUCKET, pipeline_path, data)
131+
132+
print(f"[INFO]: File '{file['name']}' extracted successfully.")
133+
success_details = {"name": file["name"], "status": "success"}
134+
135+
return True, success_details
136+
137+
except Exception as e:
138+
print(f"[ERROR]: Execution halted on file '{file['name']}': {str(e)}")
139+
140+
error_details = {
141+
"file_name": file["name"],
142+
"drive_id": file["id"],
143+
"error_type": type(e).__name__,
144+
"error_message": str(e),
145+
}
146+
147+
return False, error_details

data_extract/shared/utils.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# =============================================================================
2-
# EXTRACTOR UTILS
2+
# Google Drive Extractor Utils
33
# =============================================================================
44

55

@@ -33,8 +33,7 @@ def extract_file_content(
3333
Args:
3434
service: Authorized Google Drive API service instance.
3535
file_id: Unique Drive file identifier.
36-
mime_type: The MIME type of the source file used to determine
37-
the extraction method.
36+
mime_type: The MIME type of the source file used to determine the extraction method.
3837
3938
Behavior:
4039
- If Google Sheet: Uses files().export_media to perform a server-side

data_pipeline/run_pipeline.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# =============================================================================
2-
# PIPELINE EXECUTOR
2+
# Pipeline Orchestrator
33
# =============================================================================
44

55

0 commit comments

Comments
 (0)