refresh_sources.sh

#!/bin/bash
# ─────────────────────────────────────────────────────────────────────────────
#  refresh_sources.sh — Clone/pull repos + execute SQL files
#
#  Usage:
#    bash refresh_sources.sh full           # clone/pull ALL repos + run SQL
#    bash refresh_sources.sh sync_only      # clone/pull repos only (no SQL)
#    bash refresh_sources.sh sql            # only clone/pull sql_repos + run SQL
#    bash refresh_sources.sh data           # only clone/pull data_repos
#
#  NO set -e — every failure is handled explicitly so ClickHouse stays alive.
#
#  Private repo support:
#    Set these as HF Spaces Secrets (Settings → Secrets):
#      GIT_TOKEN  — GitHub Personal Access Token (for private GitHub repos)
#      HF_TOKEN   — HuggingFace token (for private HF datasets)
#    Tokens are injected into clone URLs automatically.
# ─────────────────────────────────────────────────────────────────────────────

CONFIG="/app/sources.yaml"
USER_FILES="/app/ch/user_files"
SQL_BASE="/app/data/sql_repos"
CH_URL="http://127.0.0.1:8123"

MODE="${1:-full}"

# ── Validate ─────────────────────────────────────────────────────────────────
if [ ! -f "$CONFIG" ]; then
  echo '{"error":"sources.yaml not found"}'
  exit 0   # exit 0 so callers don't crash
fi

# ── Helper: inject auth token into a git URL ─────────────────────────────────
# Converts https://github.com/user/repo → https://<token>@github.com/user/repo
# Converts https://huggingface.co/datasets/user/repo → https://hf_token@huggingface.co/...
inject_token() {
  local url="$1"
  local token=""

  # Detect HuggingFace URLs
  if echo "$url" | grep -q "huggingface.co"; then
    token="${HF_TOKEN:-}"
  else
    # GitHub / generic git host
    token="${GIT_TOKEN:-}"
  fi

  # If token is set, inject it into the URL
  if [ -n "$token" ]; then
    echo "$url" | sed "s|https://|https://${token}@|"
  else
    echo "$url"
  fi
}

# ── Helper: clone or pull a git repo ─────────────────────────────────────────
sync_repo() {
  local url="$1"
  local target="$2"
  local branch="$3"
  local depth="${4:-1}"
  local status="unknown"
  local output=""

  local branch_flag=""
  if [ -n "$branch" ] && [ "$branch" != "null" ]; then
    branch_flag="--branch $branch"
  fi

  if [ ! -d "$target/.git" ]; then
    local auth_url
    auth_url=$(inject_token "$url")
    output=$(git clone --depth "$depth" $branch_flag "$auth_url" "$target" 2>&1) && status="cloned" || status="clone_failed"
  else
    output=$(cd "$target" && git fetch --depth 1 origin 2>&1 && git reset --hard origin/$(git rev-parse --abbrev-ref HEAD) 2>&1) && status="pulled" || status="pull_failed"
  fi

  local commit=$(cd "$target" 2>/dev/null && git log -1 --format='%h %s' 2>/dev/null || echo "unknown")
  echo "{\"url\":$(echo "$url" | jq -Rs .),\"status\":\"$status\",\"commit\":$(echo "$commit" | jq -Rs .)}"
}

# ── Sync SQL repos ───────────────────────────────────────────────────────────
sync_sql_repos() {
  local count
  count=$(yq '.sql_repos | length // 0' "$CONFIG" 2>/dev/null || echo 0)
  if [ "$count" -eq 0 ] 2>/dev/null; then
    return
  fi

  mkdir -p "$SQL_BASE"

  for i in $(seq 0 $((count - 1))); do
    local url branch name target
    url=$(yq ".sql_repos[$i].url" "$CONFIG" 2>/dev/null || continue)
    branch=$(yq ".sql_repos[$i].branch // \"main\"" "$CONFIG" 2>/dev/null || echo "main")
    # Derive a directory name from the repo URL
    name=$(basename "$url" .git)
    target="$SQL_BASE/$name"

    sync_repo "$url" "$target" "$branch" 1
  done
}

# ── Sync data repos (parallel) ───────────────────────────────────────────────
sync_data_repos() {
  local count
  count=$(yq '.data_repos | length // 0' "$CONFIG" 2>/dev/null || echo 0)
  if [ "$count" -eq 0 ] 2>/dev/null; then
    return
  fi

  local pids=()

  for i in $(seq 0 $((count - 1))); do
    local url branch target_dir depth target
    url=$(yq ".data_repos[$i].url" "$CONFIG" 2>/dev/null || continue)
    branch=$(yq ".data_repos[$i].branch // \"main\"" "$CONFIG" 2>/dev/null || echo "main")
    target_dir=$(yq ".data_repos[$i].target_dir" "$CONFIG" 2>/dev/null || echo "")
    depth=$(yq ".data_repos[$i].clone_depth // 1" "$CONFIG" 2>/dev/null || echo 1)

    # If no target_dir specified, derive from repo URL
    if [ -z "$target_dir" ] || [ "$target_dir" = "null" ]; then
      target_dir=$(basename "$url" .git)
    fi

    target="$USER_FILES/$target_dir"
    mkdir -p "$(dirname "$target")"

    # Clone/pull in background for parallelism
    (
      sync_repo "$url" "$target" "$branch" "$depth"
    ) &
    pids+=($!)
  done

  # Wait for all parallel clones to finish
  for pid in "${pids[@]}"; do
    wait "$pid" 2>/dev/null || true
  done
}

# ── Sync HuggingFace datasets (parallel) ─────────────────────────────────────
sync_hf_datasets() {
  local count
  count=$(yq '.hf_datasets | length // 0' "$CONFIG" 2>/dev/null || echo 0)
  if [ "$count" -eq 0 ] 2>/dev/null; then
    return
  fi

  local pids=()

  for i in $(seq 0 $((count - 1))); do
    local repo_id branch target_dir depth target url
    repo_id=$(yq ".hf_datasets[$i].repo_id" "$CONFIG" 2>/dev/null || continue)
    branch=$(yq ".hf_datasets[$i].branch // \"main\"" "$CONFIG" 2>/dev/null || echo "main")
    target_dir=$(yq ".hf_datasets[$i].target_dir" "$CONFIG" 2>/dev/null || echo "")
    depth=$(yq ".hf_datasets[$i].clone_depth // 1" "$CONFIG" 2>/dev/null || echo 1)

    # Build the HF clone URL
    url="https://huggingface.co/datasets/${repo_id}"

    # If no target_dir, derive from repo_id (e.g., "user/dataset" → "dataset")
    if [ -z "$target_dir" ] || [ "$target_dir" = "null" ]; then
      target_dir=$(echo "$repo_id" | awk -F'/' '{print $NF}')
    fi

    target="$USER_FILES/$target_dir"
    mkdir -p "$(dirname "$target")"

    # Clone with git-lfs to actually download the Parquet files
    (
      export GIT_LFS_SKIP_SMUDGE=0
      sync_repo "$url" "$target" "$branch" "$depth"
    ) &
    pids+=($!)
  done

  for pid in "${pids[@]}"; do
    wait "$pid" 2>/dev/null || true
  done
}

# ── Execute all SQL files (recursive discovery) ──────────────────────────────
execute_sql_files() {
  if [ ! -d "$SQL_BASE" ]; then
    echo '{"sql_execution":"no_sql_repos_found"}'
    return
  fi

  local total=0
  local success=0
  local failed=0
  local errors=""

  # Find all .sql files, sorted alphabetically (supports 01_xxx.sql ordering)
  while IFS= read -r sql_file; do
    total=$((total + 1))
    local relative="${sql_file#$SQL_BASE/}"

    if curl -sf "$CH_URL" --data @"$sql_file" 2>/dev/null; then
      success=$((success + 1))
      echo "[sql] OK: $relative"
    else
      failed=$((failed + 1))
      echo "[sql] FAILED: $relative"
      errors="$errors $(echo "$relative" | jq -Rs .)"
    fi
  done < <(find "$SQL_BASE" -name '*.sql' -type f | sort)

  echo "{\"sql_total\":$total,\"sql_success\":$success,\"sql_failed\":$failed}"
}

# ── Also support legacy sources format (backward compat) ─────────────────────
sync_legacy_sources() {
  local count
  count=$(yq '.sources | length // 0' "$CONFIG" 2>/dev/null || echo 0)
  if [ "$count" -eq 0 ] 2>/dev/null; then
    return
  fi

  for i in $(seq 0 $((count - 1))); do
    local name url branch depth local_dir target
    name=$(yq ".sources[$i].name" "$CONFIG" 2>/dev/null || continue)
    url=$(yq ".sources[$i].repo_url" "$CONFIG" 2>/dev/null || continue)
    branch=$(yq ".sources[$i].branch // \"\"" "$CONFIG" 2>/dev/null || echo "")
    depth=$(yq ".sources[$i].clone_depth // 1" "$CONFIG" 2>/dev/null || echo 1)
    local_dir=$(yq ".sources[$i].local_dir" "$CONFIG" 2>/dev/null || echo "$name")
    target="$USER_FILES/$local_dir"

    sync_repo "$url" "$target" "$branch" "$depth"

    # Legacy view creation (only in full mode)
    if [ "$MODE" = "full" ] || [ "$MODE" = "sql" ]; then
      local num_views
      num_views=$(yq ".sources[$i].views | length // 0" "$CONFIG" 2>/dev/null || echo 0)
      for v in $(seq 0 $((num_views - 1))); do
        local view_name file_glob format cols select_clause full_glob sql
        view_name=$(yq ".sources[$i].views[$v].view_name" "$CONFIG" 2>/dev/null || continue)
        file_glob=$(yq ".sources[$i].views[$v].file_glob" "$CONFIG" 2>/dev/null || continue)
        format=$(yq ".sources[$i].views[$v].format // \"Parquet\"" "$CONFIG" 2>/dev/null || echo "Parquet")

        full_glob="${local_dir}/${file_glob}"

        local num_cols
        num_cols=$(yq ".sources[$i].views[$v].columns | length // 0" "$CONFIG" 2>/dev/null || echo 0)
        if [ "$num_cols" -gt 0 ] 2>/dev/null; then
          cols=""
          for c in $(seq 0 $((num_cols - 1))); do
            local col
            col=$(yq ".sources[$i].views[$v].columns[$c]" "$CONFIG" 2>/dev/null || continue)
            if [ -z "$cols" ]; then cols="$col"; else cols="$cols, $col"; fi
          done
          select_clause="SELECT ${cols}"
        else
          select_clause="SELECT *"
        fi

        sql="CREATE OR REPLACE VIEW ${view_name} AS ${select_clause} FROM file('${full_glob}', ${format})"
        if curl -sf "$CH_URL" --data "$sql" 2>/dev/null; then
          echo "[view] OK: $view_name"
        else
          echo "[view] FAILED: $view_name (continuing...)"
        fi
      done
    fi
  done
}

# ── Main ─────────────────────────────────────────────────────────────────────
echo "[refresh] mode=$MODE started at $(date -u +%Y-%m-%dT%H:%M:%SZ)"

case "$MODE" in
  full)
    sync_sql_repos
    sync_data_repos
    sync_hf_datasets
    sync_legacy_sources
    execute_sql_files
    ;;
  sync_only)
    sync_sql_repos
    sync_data_repos
    sync_hf_datasets
    sync_legacy_sources
    ;;
  sql)
    sync_sql_repos
    execute_sql_files
    # Also run legacy view creation
    sync_legacy_sources
    ;;
  data)
    sync_data_repos
    sync_hf_datasets
    ;;
  *)
    echo "[refresh] Unknown mode: $MODE (using full)"
    sync_sql_repos
    sync_data_repos
    sync_hf_datasets
    sync_legacy_sources
    execute_sql_files
    ;;
esac

echo "[refresh] completed at $(date -u +%Y-%m-%dT%H:%M:%SZ)"