Cosmos DB Expired Document Retention Gate

id: cosmosdb-expired-document-retention-gate
namespace: company.team

description: |
  Query a Cosmos DB container for documents that should have expired by their own expiry_field
  but are still present, classify the run COMPLIANT, WARNING, CRITICAL or PROBE_FAILED from the
  real count returned, alert Slack with that count, and optionally delete the lingering
  documents one by one behind a dry_run gate, then re-query to verify they are gone.

inputs:
  - id: database_id
    type: STRING
    displayName: Cosmos database ID
    description: Target Cosmos database name. Maps to the Query and Delete tasks'
      databaseId field.
    defaults: "compliance_db"

  - id: container_id
    type: STRING
    displayName: Cosmos container ID
    description: Target container (collection) name inside the database. Maps to the
      Query and Delete tasks' containerId field.
    defaults: "customer_records"

  - id: partition_key_field
    type: STRING
    displayName: Partition key field name
    description: Document property used as the partition key in the target container
      (e.g. partitionKey, id, or tenantId).
    defaults: "partitionKey"

  - id: expiry_field
    type: STRING
    displayName: Expiry field name
    description: Document field holding an ISO-8601 timestamp of when the document
      should have been purged, for example by application-level retention policy
      or a TTL that did not fire.
    defaults: "expiresAt"

  - id: grace_period_minutes
    type: INT
    displayName: Grace period in minutes
    description: Minutes past expiry_field a document is still tolerated before
      counting as a violation, to absorb normal purge latency.
    defaults: 60

  - id: violation_count_critical_threshold
    type: INT
    displayName: Critical violation count threshold
    description: Violation count at or above which the run is classified CRITICAL
      instead of WARNING.
    defaults: 50

  - id: remediation_action
    type: SELECT
    displayName: Remediation action
    description: What remediate_gate does with the violating documents when it runs.
      log_only never deletes anything, delete removes each violating document
      one by one.
    values:
      - "log_only"
      - "delete"
    defaults: "log_only"

  - id: auto_remediate
    type: BOOL
    displayName: Auto remediate
    description: When true, together with dry_run false and remediation_action
      delete, remove every violating document found.
    defaults: false

  - id: dry_run
    type: BOOL
    displayName: Dry run
    description: Must be false, together with auto_remediate true, for a delete to
      actually run.
    defaults: true

tasks:
  - id: probe_violations
    type: io.kestra.plugin.azure.storage.cosmosdb.Query
    description: Read-only. Finds every document whose expiry_field is in the past
      by more than grace_period_minutes. allowFailure keeps the flow going so an
      unreachable account can still be classified as PROBE_FAILED below.
    allowFailure: true
    tenantId: "{{ secret('AZURE_TENANT_ID') }}"
    connectionString: "{{ secret('COSMOS_CONNECTION_STRING') }}"
    databaseId: "{{ inputs.database_id }}"
    containerId: "{{ inputs.container_id }}"
    query: |
      SELECT c.id, c.{{ inputs.partition_key_field }}, c.{{ inputs.expiry_field }} FROM c WHERE IS_DEFINED(c.{{ inputs.expiry_field }}) AND c.{{ inputs.expiry_field }} < "{{ now() | dateAdd(-inputs.grace_period_minutes, 'MINUTES') | date("yyyy-MM-dd'T'HH:mm:ss'Z'") }}"

  - id: classify_retention_status
    type: io.kestra.plugin.core.debug.Return
    description: PROBE_FAILED when the probe produced no output at all, CRITICAL or
      WARNING when it found violations past or under the threshold, COMPLIANT
      when it found none.
    format: "{{ outputs.probe_violations is not defined ? 'PROBE_FAILED' :
      ((outputs.probe_violations.queryResults | length) == 0 ? 'COMPLIANT' :
      ((outputs.probe_violations.queryResults | length) >=
      inputs.violation_count_critical_threshold ? 'CRITICAL' : 'WARNING')) }}"

  - id: route_by_retention_status
    type: io.kestra.plugin.core.flow.Switch
    description: Branch A logs a compliant container. Branch B alerts Slack when
      documents lingered past their own expiry or the probe itself failed.
    value: "{{ outputs.classify_retention_status.value }}"
    cases:
      "COMPLIANT":
        - id: log_compliant
          type: io.kestra.plugin.core.log.Log
          message: "No expired documents found in {{ inputs.database_id }}/{{
            inputs.container_id }} past the {{ inputs.grace_period_minutes }}
            minute grace period."
      "WARNING":
        - id: alert_warning
          type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
          url: "{{ secret('SLACK_WEBHOOK') }}"
          payload: |
            {{ {"text": "Cosmos retention WARNING on " ~ inputs.database_id ~ "/" ~ inputs.container_id ~ ". " ~ (outputs.probe_violations.queryResults | length) ~ " document(s) are past their " ~ inputs.expiry_field ~ " plus a " ~ inputs.grace_period_minutes ~ " minute grace period and are still present."} | toJson }}
      "CRITICAL":
        - id: alert_critical
          type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
          url: "{{ secret('SLACK_WEBHOOK') }}"
          payload: |
            {{ {"text": "Cosmos retention CRITICAL on " ~ inputs.database_id ~ "/" ~ inputs.container_id ~ ". " ~ (outputs.probe_violations.queryResults | length) ~ " document(s) are past their " ~ inputs.expiry_field ~ " plus grace period, at or above the " ~ inputs.violation_count_critical_threshold ~ " threshold."} | toJson }}
      "PROBE_FAILED":
        - id: alert_probe_failed
          type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
          url: "{{ secret('SLACK_WEBHOOK') }}"
          payload: |
            {{ {"text": "Cosmos retention probe on " ~ inputs.database_id ~ "/" ~ inputs.container_id ~ " produced no result. Check COSMOS_CONNECTION_STRING and that the database and container exist."} | toJson }}
    defaults:
      - id: log_unhandled_status
        type: io.kestra.plugin.core.log.Log
        message: "Unhandled retention status: {{ outputs.classify_retention_status.value
          }}"

  - id: remediate_gate
    type: io.kestra.plugin.core.flow.If
    description: Runs only when the probe found violations, auto_remediate is true,
      dry_run is false, and remediation_action is delete.
    condition: "{{ (outputs.classify_retention_status.value == 'WARNING' or
      outputs.classify_retention_status.value == 'CRITICAL') and
      inputs.auto_remediate and not inputs.dry_run and inputs.remediation_action
      == 'delete' }}"
    then:
      - id: delete_expired_documents
        type: io.kestra.plugin.core.flow.Loop
        description: Deletes each violating document one at a time, supplying both id
          and the container partition key.
        concurrencyLimit: 1
        values: "{{ outputs.probe_violations.queryResults | toJson }}"
        tasks:
          - id: delete_expired_document
            type: io.kestra.plugin.azure.storage.cosmosdb.Delete
            allowFailure: true
            tenantId: "{{ secret('AZURE_TENANT_ID') }}"
            connectionString: "{{ secret('COSMOS_CONNECTION_STRING') }}"
            databaseId: "{{ inputs.database_id }}"
            containerId: "{{ inputs.container_id }}"
            item:
              id: "{{ fromJson(item.value).id }}"
              partitionKey: "{{ fromJson(item.value).partitionKey ??
                fromJson(item.value)[inputs.partition_key_field] ??
                fromJson(item.value).id }}"
              "{{ inputs.partition_key_field }}": "{{
                fromJson(item.value)[inputs.partition_key_field] ??
                fromJson(item.value).id }}"

      - id: reprobe_violations
        type: io.kestra.plugin.azure.storage.cosmosdb.Query
        allowFailure: true
        tenantId: "{{ secret('AZURE_TENANT_ID') }}"
        connectionString: "{{ secret('COSMOS_CONNECTION_STRING') }}"
        databaseId: "{{ inputs.database_id }}"
        containerId: "{{ inputs.container_id }}"
        query: |
          SELECT c.id, c.{{ inputs.partition_key_field }}, c.{{ inputs.expiry_field }} FROM c WHERE IS_DEFINED(c.{{ inputs.expiry_field }}) AND c.{{ inputs.expiry_field }} < "{{ now() | dateAdd(-inputs.grace_period_minutes, 'MINUTES') | date("yyyy-MM-dd'T'HH:mm:ss'Z'") }}"

      - id: reclassify_retention_status
        type: io.kestra.plugin.core.debug.Return
        format: "{{ outputs.reprobe_violations is not defined ? 'PROBE_FAILED' :
          ((outputs.reprobe_violations.queryResults | length) == 0 ? 'COMPLIANT'
          : 'WARNING') }}"

      - id: verify_remediation
        type: io.kestra.plugin.core.flow.If
        description: Fails the flow when violations still remain after the delete loop,
          so a partial delete is never silently reported as a fix.
        condition: "{{ outputs.reclassify_retention_status.value != 'COMPLIANT' }}"
        then:
          - id: remediation_failed
            type: io.kestra.plugin.core.execution.Fail
            errorMessage: "Delete loop ran, but {{ inputs.database_id }}/{{
              inputs.container_id }} still has expired documents."
        else:
          - id: remediated
            type: io.kestra.plugin.core.log.Log
            message: "Deleted and verified. 0 expired documents remain in {{
              inputs.database_id }}/{{ inputs.container_id }}."

  - id: log_audit_summary
    type: io.kestra.plugin.core.log.Log
    description: Final audit log confirming the gate ran, regardless of the branch
      taken above.
    message: "Cosmos retention gate completed. database={{ inputs.database_id }}
      container={{ inputs.container_id }} status={{
      outputs.classify_retention_status.value }} auto_remediate={{
      inputs.auto_remediate }} dry_run={{ inputs.dry_run }}"

errors:
  - id: alert_flow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK') }}"
    payload: |
      {{ {"text": "Cosmos retention gate flow failed (execution " ~ execution.id ~ "). Check the execution logs."} | toJson }}

triggers:
  - id: six_hourly
    type: io.kestra.plugin.core.trigger.Schedule
    description: Shipped disabled, and always in safe mode by default.
    cron: "0 */6 * * *"
    disabled: true
    inputs:
      database_id: "compliance_db"
      container_id: "customer_records"
      partition_key_field: "partitionKey"
      expiry_field: "expiresAt"
      grace_period_minutes: 60
      violation_count_critical_threshold: 50
      remediation_action: "log_only"
      auto_remediate: false
      dry_run: true