refresh_sources.py

#!/usr/bin/env python3
"""
refresh_sources.py — Standalone refresh script in Python

Actions:
  • git clone (first run) or git pull (subsequent) for each source
  • Inject auth token if `auth_env_var` is defined in config
  • CREATE OR REPLACE VIEW for each defined view
"""
import os
import sys
import json
import yaml
import urllib.request
import urllib.error
import subprocess

# Allow local testing gracefully
CONFIG = "/app/sources.yaml" if os.path.exists("/app/sources.yaml") else "sources.yaml"
USER_FILES = "/app/ch/user_files" if os.path.exists("/app") else "user_files"
CH_URL = "http://127.0.0.1:8123"

def run_subprocess(cmd, cwd=None, env=None):
    try:
        # We use Popen to stream stderr to the console for progress bars
        # while still capturing stdout for return values if needed.
        process = subprocess.Popen(
            cmd, cwd=cwd, env=env,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )
        
        # Stream stderr in real time
        stderr_output = []
        while True:
            line = process.stderr.readline()
            if not line and process.poll() is not None:
                break
            if line:
                sys.stderr.write(line)
                sys.stderr.flush()
                stderr_output.append(line)
        
        stdout, _ = process.communicate()
        
        if process.returncode == 0:
            return stdout.strip(), True
        else:
            return "".join(stderr_output).strip(), False
    except Exception as e:
        return str(e), False


def execute_clickhouse_query(sql):
    req = urllib.request.Request(CH_URL, data=sql.encode('utf-8'), method='POST')
    try:
        resp = urllib.request.urlopen(req, timeout=10)
        return resp.read().decode('utf-8'), True
    except urllib.error.URLError as e:
        return str(e), False

def main():
    mode = "full"
    if len(sys.argv) > 1:
        mode = sys.argv[1]

    if not os.path.exists(CONFIG):
        print(json.dumps({"error": f"Config not found at {CONFIG}"}))
        sys.exit(1)

    with open(CONFIG, "r") as f:
        config_data = yaml.safe_load(f)
    
    sources = config_data.get("sources", [])
    
    result = {
        "mode": mode,
        "source_count": len(sources),
        "sources": []
    }

    env_vars = os.environ.copy()

    for source in sources:
        name = source.get("name", "unknown")
        repo_url = source.get("repo_url")
        local_dir = source.get("local_dir")
        clone_depth = str(source.get("clone_depth", 1))
        branch = source.get("branch", "")
        auth_env_var = source.get("auth_env_var", "")
        
        target_dir = os.path.join(USER_FILES, local_dir)

        # Build git clone/pull options
        auth_flag = []
        if auth_env_var and auth_env_var in env_vars:
            token = env_vars[auth_env_var]
            auth_flag = ["-c", f"http.extraHeader=Authorization: Bearer {token}"]
        
        branch_flag = []
        if branch:
            branch_flag = ["--branch", branch]

        git_status = "unknown"
        
        # Git Sync
        if not os.path.isdir(target_dir):
            cmd = ["git", "clone", "--depth", clone_depth] + branch_flag + auth_flag + [repo_url, target_dir]
            output, success = run_subprocess(cmd)
            git_status = "cloned" if success else f"error_clone: {output}"
            
            # Retrieve real files for Git LFS pointers natively used by Hugging Face Datasets
            if success:
                cmd_lfs = ["git", "lfs", "pull"]
                out_lfs, suc_lfs = run_subprocess(cmd_lfs, cwd=target_dir)
                if not suc_lfs:
                    git_status = f"error_lfs_pull: {out_lfs}"
        else:
            # We want to fetch and hard reset
            cmd_fetch = ["git"] + auth_flag + ["fetch", "--depth", "1", "origin"]
            output_f, success_f = run_subprocess(cmd_fetch, cwd=target_dir)
            if success_f:
                # get current branch
                rev_cmd = ["git", "rev-parse", "--abbrev-ref", "HEAD"]
                curr_branch, _ = run_subprocess(rev_cmd, cwd=target_dir)
                reset_cmd = ["git", "reset", "--hard", f"origin/{curr_branch}"]
                output_r, success_r = run_subprocess(reset_cmd, cwd=target_dir)
                git_status = "pulled" if success_r else f"error_reset: {output_r}"
            else:
                git_status = f"error_fetch: {output_f}"
        
        latest_commit, _ = run_subprocess(["git", "log", "-1", "--format=%h %s"], cwd=target_dir)
        if not latest_commit:
            latest_commit = "unknown"

        source_info = {
            "name": name,
            "git_status": git_status,
            "latest_commit": latest_commit
        }

        # View Creation
        if mode == "full":
            views_info = []
            views = source.get("views", [])
            for view in views:
                view_name = view.get("view_name")
                file_glob = view.get("file_glob")
                format_ = view.get("format", "Parquet")
                columns = view.get("columns", [])

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

                if columns:
                    cols_str = ", ".join(columns)
                    select_clause = f"SELECT {cols_str}"
                else:
                    select_clause = "SELECT *"

                sql = f"CREATE OR REPLACE VIEW {view_name} AS {select_clause} FROM file('{full_glob}', {format_})"

                _, success = execute_clickhouse_query(sql)
                view_status = "ok" if success else "error"
                
                views_info.append({
                    "name": view_name,
                    "glob": full_glob,
                    "select": select_clause,
                    "status": view_status
                })
            
            source_info["views"] = views_info
        
        result["sources"].append(source_info)
    
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()