|
1 | 1 | # ============================================================================= |
2 | | -# GOOGLE DRIVE FILE EXTRACTOR |
| 2 | +# Google Drive Extractor Orchestrator |
3 | 3 | # ============================================================================= |
4 | 4 |
|
5 | 5 | import sys |
6 | 6 | import json |
7 | 7 | import uuid |
8 | 8 | from data_extract.shared.utils import ( |
9 | | - extract_file_content, |
10 | | - check_handshake, |
11 | 9 | get_drive_service, |
12 | 10 | check_gcs_marking, |
13 | 11 | upload_to_gcs, |
14 | 12 | get_target_folder_name, |
15 | 13 | ) |
| 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. |
16 | 25 |
|
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. |
21 | 30 |
|
| 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. |
22 | 36 |
|
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 | + """ |
25 | 42 |
|
26 | 43 | service = get_drive_service() |
27 | 44 | metadata_path = f"logs/{target_child_folder}_metadata.json" |
28 | 45 | archival_marking_path = f"status/{target_child_folder}.success" |
29 | 46 |
|
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 | + } |
55 | 53 |
|
56 | 54 | # Deduplication Check |
57 | 55 | if check_gcs_marking(ARCHIVAL_BUCKET, archival_marking_path): |
58 | 56 | print(f"[INFO]: {target_child_folder} already processed.") |
59 | | - |
60 | 57 | return 0 |
61 | 58 |
|
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: |
70 | 62 | return 1 |
71 | 63 |
|
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) |
79 | 66 |
|
| 67 | + # Exit if handshake failed or empty list |
| 68 | + if files_in_drive is None or len(files_in_drive) == 0: |
80 | 69 | return 1 |
81 | 70 |
|
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 | | - |
101 | 71 | for file in files_in_drive: |
102 | 72 |
|
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']}" |
120 | 75 |
|
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 | + ) |
129 | 79 |
|
130 | | - metadata["errors"].append(error_details) |
| 80 | + if ok: |
| 81 | + metadata["files_processed"].append(details) |
| 82 | + else: |
| 83 | + metadata["errors"].append(details) |
131 | 84 | metadata["status"] = "failed" |
132 | 85 |
|
133 | | - # runtime metadata |
| 86 | + # Upload failure metadata |
134 | 87 | upload_to_gcs(ARCHIVAL_BUCKET, metadata_path, json.dumps(metadata)) |
135 | | - |
136 | | - print(f"[ERROR]: Execution halted {str(e)}") |
137 | 88 | return 1 |
138 | 89 |
|
| 90 | + # Upload success metadata and marker |
139 | 91 | upload_to_gcs(ARCHIVAL_BUCKET, metadata_path, json.dumps(metadata)) |
140 | 92 | upload_to_gcs(ARCHIVAL_BUCKET, archival_marking_path, "") |
| 93 | + |
| 94 | + print(f"[SUCCESS]: Folder '{target_child_folder}' completely processed.") |
141 | 95 | return 0 |
142 | 96 |
|
143 | 97 |
|
144 | 98 | def main(): |
145 | | - |
146 | 99 | 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}") |
148 | 101 |
|
149 | | - exit_code = run_extraction(target_folder) |
| 102 | + exit_code = orchestrate_extract(target_folder) |
150 | 103 | sys.exit(exit_code) |
151 | 104 |
|
152 | 105 |
|
|
0 commit comments