main.py

import subprocess
import traceback
import os
import tempfile
import shutil
from rfctl.client import LambdaClient
import yaml


def lambda_handler(event: dict, ctx):
    try:
        _call_method = event.get("_call_method")
        if _call_method == "deploy-function":
            return deploy_func(event, ctx)
        if _call_method == "invoke-function":
            return invoke_func(event, ctx)
        if _call_method == "delete-function":
            return delete_func(event, ctx)
        if _call_method == "get-function":
            return get_func(event, ctx)
    except Exception as e:
        return {
            "code": 1,
            "message": "error",
            "result": "{}".format(e)
        }
    return {
        "code": 1,
        "message": "error",
        "result": "call method error"
    }


def deploy_func(event: dict, ctx):
    name, namespace = event.get("name"), event.get("namespace")
    manifest, src_code = event.get("manifest"), event.get("src_code")
    if not os.getenv("S3_SYNC_URI"):
        raise Exception("S3_SYNC_URI error")
    func_info: dict = yaml.load(manifest, yaml.FullLoader)
    # write source file
    workdir = "{}/{}/{}".format(os.path.abspath(os.getenv("WORKSPACE_DIR")), namespace, name)
    if os.path.exists(workdir):
        shutil.rmtree(workdir)
    rc, _, _ = run_cmd(["mkdir", "-p", workdir])
    if rc != 0:
        raise Exception("mkdir workspace error")
    source_items = {
        "{}/lambda.yaml".format(workdir): manifest
    }
    for path, content in src_code.items():
        source_items.update({"{}/{}".format(workdir, path): content})
    for path, content in source_items.items():
        directory = os.path.dirname(path)
        if directory and not os.path.exists(directory):
            os.makedirs(directory, exist_ok=True)
        with open(path, 'w', encoding="utf-8") as f:
            f.write(content)
    # sync files to s3
    s3_sync = ["juicefs", "sync", "--check-all", "--delete-dst"]
    s3_sync.extend(["{}/".format(workdir), "{}/{}/{}/".format(os.getenv("S3_SYNC_URI"), namespace, name)])
    rc, _, _ = run_cmd(s3_sync, cwd=workdir)
    if rc != 0:
        raise Exception("save files to s3 error")
    # create or update func code
    subcmd = "update-code"
    rc, get_stdout, get_stderr = run_cmd(["rfctl", "get"], cwd=workdir)
    if "ResourceNotFoundException" in get_stdout or "ResourceNotFoundException" in get_stderr:
        subcmd = "create"
    rc, code_stdout, code_stderr = run_cmd(["rfctl", subcmd], cwd=workdir)
    if rc != 0:
        return {
            "code": 1,
            "message": "error",
            "result": "code deploy:\n{}\n{}".format(code_stdout, code_stderr)
        }
    # update func config
    if subcmd != "create":
        rc, cfg_stdout, cfg_stderr = run_cmd(["rfctl", "update-config"], cwd=workdir)
        if rc != 0:
            return {
                "code": 1,
                "message": "error",
                "result": "code deploy:\n{}\n{}\nupdate-config:\n{}\n{}".format(code_stdout, code_stderr, cfg_stdout, cfg_stderr)
            }
    # update func url
    exist_rc, _, _ = run_cmd(["rfctl", "get-url"], cwd=workdir)
    exist_url = exist_rc == 0
    if func_info.get("spec", {}).get("enable_url") and not exist_url:
        rc, _, _ = run_cmd(["rfctl", "create-url"], cwd=workdir)
        if rc != 0:
            return {
                "code": 1,
                "message": "error",
                "result": "enable func url call error"
            }
    if not func_info.get("spec", {}).get("enable_url") and exist_url:
        rc, _, _ = run_cmd(["rfctl", "delete-url"], cwd=workdir)
        if rc != 0:
            return {
                "code": 1,
                "message": "error",
                "result": "disable func url call error"
            }
    return {
        "code": 0,
        "message": "success"
    }


def invoke_func(event: dict, ctx):
    client = LambdaClient(event.get("namespace"))
    rsp = client.invoke_function(event.get("name"), event.get("payload"), get_log=True)
    return rsp


def delete_func(event: dict, ctx):
    name, namespace = event.get("name"), event.get("namespace")
    if not os.getenv("S3_SYNC_URI"):
        raise Exception("S3_SYNC_URI error")
    client = LambdaClient(namespace)
    client.lambda_client.get_function(FunctionName=name)
    # delete funcdef
    workdir = tempfile.mkdtemp()
    manifest = "{}/lambda.yaml".format(workdir)
    content = """metadata:
    name: {}
    namespace: {}
# placeholder spec
spec:
    build:
        source: .
        manifest: requirements.txt
        language: python
        architecture: x86_64
    handler: main.lambda_handler
    timeout: 120
    runtime: "python3.10"
""".format(name, namespace)
    with open(manifest, 'w', encoding="utf-8") as f:
        f.write(content)
    rc, del_stdout, del_stderr = run_cmd(["rfctl", "delete"], cwd=workdir)
    if rc != 0:
        return {
            "code": 1,
            "message": "error",
            "result": "delete function:\n{}\n{}".format(del_stdout, del_stderr)
        }
    shutil.rmtree(workdir)
    # fetch&backup code
    workdir = os.path.abspath(tempfile.mkdtemp())
    s3_sync = ["juicefs", "sync", "--check-all", "--delete-dst"]
    s3_sync.extend(["{}/{}/{}/".format(os.getenv("S3_SYNC_URI"), namespace, name), "{}/".format(workdir)])
    rc, _, _ = run_cmd(s3_sync, cwd=workdir)
    if rc != 0:
        raise Exception("sync files from s3 error")
    s3_sync = ["juicefs", "sync", "--check-all", "--delete-dst"]
    s3_sync.extend(["{}/".format(workdir), "{}/deleted/{}/{}/".format(os.getenv("S3_SYNC_URI"), namespace, name)])
    rc, _, _ = run_cmd(s3_sync, cwd=workdir)
    if rc != 0:
        raise Exception("save backup files to s3 error")
    shutil.rmtree(workdir)
    # delete func code
    workdir = os.path.abspath(tempfile.mkdtemp())
    s3_sync = ["juicefs", "sync", "--check-all", "--delete-dst"]
    s3_sync.extend(["{}/".format(workdir), "{}/{}/{}/".format(os.getenv("S3_SYNC_URI"), namespace, name)])
    rc, _, _ = run_cmd(s3_sync, cwd=workdir)
    if rc != 0:
        raise Exception("delete files from s3 error")
    shutil.rmtree(workdir)
    return {
        "code": 0,
        "message": "success"
    }


def get_func(event: dict, ctx):
    # check func exist
    name, namespace = event.get("name"), event.get("namespace")
    if not os.getenv("S3_SYNC_URI"):
        raise Exception("S3_SYNC_URI error")
    client = LambdaClient(namespace)
    client.lambda_client.get_function(FunctionName=name)
    workdir = "{}/{}/{}".format(os.path.abspath(os.getenv("WORKSPACE_DIR")), namespace, name)
    os.makedirs(workdir, exist_ok=True)
    # sync files from s3
    s3_sync = ["juicefs", "sync", "--check-all", "--delete-dst"]
    s3_sync.extend(["{}/{}/{}/".format(os.getenv("S3_SYNC_URI"), namespace, name), "{}/".format(workdir)])
    rc, _, _ = run_cmd(s3_sync, cwd=workdir)
    if rc != 0:
        raise Exception("sync files from s3 error")
    # build sources
    manifest, src_code = "", {}
    for root, _, files in os.walk(workdir):
        for file in files:
            file_path = os.path.join(root, file)
            if file_path.endswith("lambda.yaml") and not manifest:
                with open(file_path, 'r', encoding='utf-8') as f:
                    manifest = f.read()
                continue
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    _path = file_path.removeprefix(workdir).lstrip("/")
                    src_code[_path] = f.read()
            except Exception as e:
                print(f"can not read file {file_path}: {e}")
    return {
        "code": 0,
        "result": {
            "manifest": manifest,
            "src_code": src_code
        }
    }


def run_cmd(cmd, cwd=None, env: dict = None):
    try:
        if not env:
            env = os.environ.copy()
        stdout, stderr = "", ""
        with subprocess.Popen(cmd, cwd=cwd, env=env, stderr=subprocess.PIPE, stdout=subprocess.PIPE) as process:
            stdout, stderr = process.communicate()
            rc = process.wait()
        if isinstance(stdout, bytes):
            stdout = stdout.decode()
        if isinstance(stderr, bytes):
            stderr = stderr.decode()
        return rc, stdout, stderr
    except Exception as e:
        print(traceback.format_exc())
        return 1, None, None


if __name__ == "__main__":
    os.environ.update({"WORKSPACE_DIR": "./temp"})
    _manifest = '''metadata:
  name: demo-py310
  namespace: default
spec:
  description: |
    A demo python func
    func's args
    ```
    {"a":1}
    这是一个函数
    ```
  build:
    source: .
    manifest: requirements.txt
    language: "python"
    architecture: "x86_64"
  enable_url: true
  handler: "main.lambda_handler"
  timeout: 120
  runtime: "python3.10"'''
    _src_code = {
        "main.py": '''
def lambda_handler(event, ctx):
    ctx.log("this is log")
    return {
        "msg": "hello py3.10"
    }
''',
        "requirements.txt": "",
        "tests/test.py": "print('test')"
    }
    print("create==========>")
    print(deploy_func({"namespace": "default", "name": "demo-py310", "manifest": _manifest, "src_code": _src_code}, {}))
    print("invoke==========>")
    _client = LambdaClient("default")
    print(_client.invoke_function("demo-py310", {}, get_log=True))
    print("get==========>")
    print(get_func({"namespace": "default", "name": "demo-py310"}, {}))
    print("delete==========>")
    print(delete_func({"namespace": "default", "name": "demo-py310"}, {}))