[DCP Ingestion] Add preprocessing stage to ingestion helper

pipeline/workflow/ingestion-helper/clients/fakeworkflow.yaml

@@ -0,0 +1,310 @@
+  main:
+    params: [input]
+    steps:
+      - init:
+          assign:
+            - project_id: 'datcom-website-dev'
+            - workflow_id: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}'
+            - version: '${"version-" + string(int(sys.now()))}'
+            - bucket_name: '${text.split(input.tempLocation, "/")[2]}'
+            - execution_error: null
+            - current_stage: "dataflow"
+            - lock_timeout: 82800
+            - run_embeddings: true
+            - run_postproc: true
+            - postprocessing_result: null
+            - embedding_result: null
+            - decoded_imports: '${json.decode(input.importList)}'
+            - num_imports: '${len(decoded_imports)}'
+            - combined_import_name: ""
+            - imports_status_list: []
+            - imports_history_list: []
+            - imports_version_list: []
+            - launch_params:
+                projectId: '${project_id}'
+                spannerInstanceId: '${input.spannerInstanceId}'
+                spannerDatabaseId: '${input.spannerDatabaseId}'
+                importList: '${input.importList}'
+                tempLocation: '${input.tempLocation}'
+                stagingLocation: '${input.tempLocation}'
+                forceCombineNodes: 'true'
+                isBaseDc: 'false'
+      - build_lists:
+          for:
+            value: imp
+            in: ${decoded_imports}
+            steps:
+              - append_lists:
+                  assign:
+                    - item_gcs_path: '${"gs://" + bucket_name + "/ingestion/internal/metadata/" + imp.importName + "/" + version}'
+                    - status_item: {}
+                    - status_item.importName: '${imp.importName}'
+                    - status_item.latestVersion: '${item_gcs_path}'
+                    - status_item.graphPath: '${imp.graphPath}'
+                    - status_item.status: "STAGING"
+                    - history_item: {}
+                    - history_item.importName: '${imp.importName}'
+                    - history_item.latestVersion: '${version}'
+                    - imports_status_list: '${list.concat(imports_status_list, status_item)}'
+                    - imports_history_list: '${list.concat(imports_history_list, history_item)}'
+                    - imports_version_list: '${list.concat(imports_version_list, imp.importName)}'
+                    - clean_item: '${text.replace_all(text.replace_all(text.to_lower(imp.importName), "/", "-"), "_", "-")}'
+                    - join_prefix: '${if(combined_import_name == "", "", "-")}'
+                    - combined_import_name: '${combined_import_name + join_prefix + clean_item}'
+      - set_dataflow_job_name:
+          assign:
+            - import_name_len: '${len(combined_import_name)}'
+            - substring_end: '${if(import_name_len < 35, import_name_len, 35)}'
+            - sanitized_short_import: '${text.substring(combined_import_name, 0, substring_end)}'
+            - dataflow_job_name: '${"gmechali-june22-" + sanitized_short_import + "-" + string(int(sys.now()))}'
+      - acquire_lock:
+          try:
+            call: http.post
+            args:
+              url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/database/lock/acquire'
+              auth:
+                type: OIDC
+              body:
+                workflowId: '${workflow_id}'
+                timeout: '${lock_timeout}'
+            result: lock_result
+          retry:
+            predicate: '${http.default_retry_predicate}'
+            max_retries: 60 # Approx 5 hours
+            backoff:
+              initial_delay: 60
+              max_delay: 300 # Max 5 minutes retry interval
+              multiplier: 2
+      - update_history_pending:
+          call: http.post
+          args:
+            url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+            auth:
+              type: OIDC
+            body:
+              workflowId: '${workflow_id}'
+              status: 'PENDING'
+              stage: 'dataflow'
+              importList: '${imports_history_list}'
+          result: history_pending_result
+      - process_ingestion:
+          try:
+            steps:
+              - set_import_staging:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/status'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      imports: '${imports_status_list}'
+                      status: "STAGING"
+                  result: staging_result
+              - run_flex_template:
+                  call: googleapis.dataflow.v1b3.projects.locations.flexTemplates.launch
+                  args:
+                    projectId: '${project_id}'
+                    location: '${input.region}'
+                    body:
+                      launchParameter:
+                        jobName: '${dataflow_job_name}'
+                        containerSpecGcsPath: 'gs://datcom-templates/templates/flex/ingestion-stable.json'
+                        parameters: '${launch_params}'
+                        environment:
+                          serviceAccountEmail: 'gmechali-june22-dc-ing-df-sa@datcom-website-dev.iam.gserviceaccount.com'
+                          tempLocation: '${input.tempLocation}'
+                          stagingLocation: '${input.tempLocation}'
+                  result: launch_result
+              - get_job_id:
+                  assign:
+                    - job_id: '${launch_result.job.id}'
+              - update_history_running:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      status: 'RUNNING'
+                      stage: 'dataflow'
+                      jobId: '${job_id}'
+                  result: history_running_result
+              - poll_job:
+                  steps:
+                    - get_status:
+                        call: googleapis.dataflow.v1b3.projects.locations.jobs.get
+                        args:
+                          projectId: '${project_id}'
+                          location: '${input.region}'
+                          jobId: '${job_id}'
+                        result: job_status
+                    - check_terminal:
+                        switch:
+                          - condition: '${job_status.currentState in ["JOB_STATE_DONE", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED", "JOB_STATE_UPDATED", "JOB_STATE_DRAINED"]}'
+                            next: check_success
+                    - wait_and_retry:
+                        call: sys.sleep
+                        args:
+                          seconds: 60
+                        next: poll_job
+              - check_success:
+                  switch:
+                    - condition: '${job_status.currentState == "JOB_STATE_DONE"}'
+                      next: update_history_dataflow_success
+              - update_history_dataflow_success:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      status: 'SUCCESS'
+                      stage: 'dataflow'
+                      jobId: '${job_id}'
+                      importList: '${imports_history_list}'
+                  result: history_dataflow_success_result
+                  next: ingestion_postprocessing
+              - fail_on_job_status:
+                  raise: '${ "Dataflow job failed with state: " + job_status.currentState }'
+              - ingestion_postprocessing:
+                  steps:
+                    - set_stage_postproc:
+                        assign:
+                          - current_stage: "postprocessing"
+                    - update_history_postproc_running:
+                        call: http.post
+                        args:
+                          url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+                          auth:
+                            type: OIDC
+                          body:
+                            workflowId: '${workflow_id}'
+                            status: 'RUNNING'
+                            stage: 'postprocessing'
+                        result: history_postproc_running_result
+                    - run_postprocessing_parallel:
+                        parallel:
+                          shared: [postprocessing_result, embedding_result]
+                          branches:
+                            - postprocessing_branch:
+                                steps:
+                                  - check_postproc_flag:
+                                      switch:
+                                        - condition: '${not run_postproc}'
+                                          next: end_postproc_branch
+                                  - run_postprocessings:
+                                      call: http.post
+                                      args:
+                                        url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/aggregation/run'
+                                        auth:
+                                          type: OIDC
+                                        body:
+                                          importList: '${json.decode(input.importList)}'
+                                      result: postprocessing_result
+                                  - end_postproc_branch:
+                                      assign:
+                                        - dummy: true
+
+                            - embeddings_branch:
+                                steps:
+                                  - check_embeddings_flag:
+                                      switch:
+                                        - condition: '${not run_embeddings}'
+                                          next: end_embeddings_branch
+                                  - run_embeddings:
+                                      call: http.post
+                                      args:
+                                        url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/embeddings/ingest'
+                                        auth:
+                                          type: OIDC
+                                        body:
+                                          enableEmbeddings: true
+                                        timeout: 1800
+                                      result: embedding_result
+                                  - end_embeddings_branch:
+                                      assign:
+                                        - dummy: true
+              - promote_version:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/version'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      imports: '${imports_version_list}'
+                      version: '${version}'
+                      comment: '${"Auto-promoted by workflow " + workflow_id}'
+                  result: promote_result
+              - update_ingestion_history_step:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-status'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      jobId: '${job_id}'
+                      status: 'SUCCESS'
+                      importList: '${imports_history_list}'
+                  result: history_result
+              - update_history_postproc_success:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      status: 'SUCCESS'
+                      stage: null
+                  result: history_postproc_success_result
+          except:
+            as: e
+            steps:
+              - update_history_failure:
+                  call: http.post
+                  args:
+                    url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/imports/ingestion-history'
+                    auth:
+                      type: OIDC
+                    body:
+                      workflowId: '${workflow_id}'
+                      status: 'FAILURE'
+                      stage: '${current_stage}'
+                      jobId: '${default(job_id, "N/A")}'
+                      importList: '${imports_history_list}'
+                  result: history_failure_result
+              - capture_error:
+                  assign:
+                    - execution_error: '${e}'
+      - release_lock_step:
+          call: http.post
+          args:
+            url: 'https://gmechali-june22-dc-ingestion-helper-kqb7thiuka-uc.a.run.app/database/lock/release'
+            auth:
+              type: OIDC
+            body:
+              workflowId: '${workflow_id}'
+          result: release_lock_result
+      - fail_workflow:
+          switch:
+            - condition: '${execution_error != null}'
+              raise: '${execution_error}'
+
+      - restart_service:
+          call: googleapis.run.v2.projects.locations.services.patch
+          args:
+            name: "projects/datcom-website-dev/locations/us-central1/services/gmechali-june22-dc-datacommons-service"
+            updateMask: "template.labels"
+            body:
+              template:
+                labels:
+                  restarted-at: '${string(int(sys.now()))}'
+
+
+      - return_result:
+          return: '${launch_result}'

pipeline/workflow/ingestion-helper/clients/spanner.py

@@ -38,6 +38,7 @@ class IngestionState(str, Enum):
 class IngestionStage(str, Enum):
     DATAFLOW = "dataflow"
     POSTPROCESSING = "postprocessing"
+    PREPROCESSING = "preprocessing"
 
 
 class SpannerClient:
@@ -323,7 +324,7 @@ def update_ingestion_history(self,
         Args:
             workflow_id: The ID of the workflow.
             status: The status of the ingestion stage (PENDING, RUNNING, SUCCESS, FAILURE).
-            stage: The stage of the ingestion (dataflow, postprocessing, etc.).
+            stage: The stage of the ingestion (preprocessing, dataflow, postprocessing, etc.).
             job_id: The Dataflow job ID.
             ingested_imports: List of ingested import names.
             metrics: A dictionary containing metrics about the ingestion.
@@ -335,26 +336,39 @@ def update_ingestion_history(self,
 
         def _update(transaction: Transaction):
             status_str = status.value if hasattr(status, 'value') else status
+            stage_str = stage.value if stage and hasattr(stage, 'value') else stage
+
             columns = ["WorkflowExecutionID", "Status"]
             values = [workflow_id, status_str]
 
-            if stage:
-                stage_str = stage.value if hasattr(stage, 'value') else stage
+            if stage_str:
                 columns.append("Stage")
                 values.append(stage_str)
 
-            if status == IngestionState.PENDING:
-                columns.append("CreationTimestamp")
-                values.append(spanner.COMMIT_TIMESTAMP)
+            is_base_dc = os.environ.get('IS_BASE_DC', 'true').lower() == 'true'
+            if not is_base_dc:
+                if stage == IngestionStage.PREPROCESSING:
+                    columns.append("CreationTimestamp")
+                    values.append(spanner.COMMIT_TIMESTAMP)
+            else:
+                if status == IngestionState.PENDING:
+                    columns.append("CreationTimestamp")
+                    values.append(spanner.COMMIT_TIMESTAMP)
 
             # The statements below allow us to construct a partial update, only for the fields that are set.
-            if status in (IngestionState.SUCCESS, IngestionState.FAILURE,
-                          IngestionState.RETRY):
+            terminal_statuses = {
+                IngestionState.SUCCESS.value,
+                IngestionState.FAILURE.value,
+                IngestionState.RETRY.value,
+            }
+            if status_str in terminal_statuses:
                 columns.append("CompletionTimestamp")
                 values.append(spanner.COMMIT_TIMESTAMP)
                 columns.append("IngestionFailure")
-                values.append(status in (IngestionState.FAILURE,
-                                         IngestionState.RETRY))
+                values.append(status_str in {
+                    IngestionState.FAILURE.value,
+                    IngestionState.RETRY.value,
+                })
 
             if job_id:
                 columns.append("DataflowJobID")