fabricks 3.0.11__py3-none-any.whl

fabricks/core/jobs/get_job.py

@@ -0,0 +1,127 @@
+from typing import Optional, Union, cast, overload
+
+from pyspark.sql.types import Row
+
+from fabricks.core.jobs.base._types import Bronzes, Golds, Silvers, TBronze, TGold, TSilver
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.get_job_id import get_job_id
+
+
+@overload
+def get_job(*, step: str, job_id: str) -> BaseJob: ...
+
+
+@overload
+def get_job(*, step: str, topic: str, item: str) -> BaseJob: ...
+
+
+@overload
+def get_job(*, row: Row) -> BaseJob: ...
+
+
+@overload
+def get_job(job: str) -> BaseJob: ...
+
+
+def get_job(
+    job: Optional[str] = None,
+    step: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job_id: Optional[str] = None,
+    row: Optional[Row] = None,
+) -> BaseJob:
+    """
+    Retrieve a job based on the provided parameters.
+
+    Args:
+        step (Optional[str]): The step of the job.
+        topic (Optional[str]): The topic of the job.
+        item (Optional[str]): The item of the job.
+        job_id (Optional[str]): The ID of the job.
+        job (Optional[str]): The job string.
+        row (Optional[Row]): The row object containing job information.
+
+    Returns:
+        BaseJob: The retrieved job.
+
+    Raises:
+        ValueError: If the required parameters are not provided.
+
+    """
+    if row:
+        if "step" in row and "topic" in row and "item" in row:
+            j = get_job_internal(step=row.step, topic=row.topic, item=row.item)
+        elif "step" in row and "job_id" in row:
+            j = get_job(step=row.step, job_id=row.job_id)
+        elif "job" in row:
+            parts = row.job.split(".")
+            s = parts[0]
+            job_id = get_job_id(job=row.job)
+            j = get_job_internal(step=s, job_id=job_id)
+        else:
+            raise ValueError("step, topic, item or step, job_id or job mandatory")
+
+    elif job:
+        parts = job.split(".")
+        s = parts[0]
+        job_id = get_job_id(job=job)
+        j = get_job_internal(step=s, job_id=job_id)
+
+    elif job_id:
+        assert step, "step mandatory"
+        j = get_job_internal(step=step, job_id=job_id)
+
+    else:
+        assert step, "step mandatory"
+        assert topic, "topic mandatory"
+        assert item, "item mandatory"
+        j = get_job_internal(step=step, topic=topic, item=item)
+
+    return j
+
+
+def get_job_internal(
+    step: str,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job_id: Optional[str] = None,
+    conf: Optional[Union[dict, Row]] = None,
+):
+    if step in Bronzes:
+        from fabricks.core.jobs.bronze import Bronze
+
+        step = cast(TBronze, step)
+        if job_id is not None:
+            job = Bronze.from_job_id(step=step, job_id=job_id, conf=conf)
+        else:
+            assert topic
+            assert item
+            job = Bronze.from_step_topic_item(step=step, topic=topic, item=item, conf=conf)
+
+    elif step in Silvers:
+        from fabricks.core.jobs.silver import Silver
+
+        step = cast(TSilver, step)
+        if job_id is not None:
+            job = Silver.from_job_id(step=step, job_id=job_id, conf=conf)
+        else:
+            assert topic
+            assert item
+            job = Silver.from_step_topic_item(step=step, topic=topic, item=item, conf=conf)
+
+    elif step in Golds:
+        from fabricks.core.jobs.gold import Gold
+
+        step = cast(TGold, step)
+        if job_id is not None:
+            job = Gold.from_job_id(step=step, job_id=job_id, conf=conf)
+        else:
+            assert topic
+            assert item
+            job = Gold.from_step_topic_item(step=step, topic=topic, item=item, conf=conf)
+
+    else:
+        raise ValueError(f"{step} not found")
+
+    return job

fabricks/core/jobs/get_job_conf.py

@@ -0,0 +1,152 @@
+from typing import Optional, Union, cast, overload
+
+from pyspark.sql.types import Row
+
+from fabricks.context import IS_JOB_CONFIG_FROM_YAML, SPARK
+from fabricks.core.jobs.base._types import Bronzes, Golds, JobConf, Silvers, TBronze, TGold, TSilver, TStep
+from fabricks.core.jobs.get_job_id import get_job_id
+
+
+def get_job_conf_internal(step: TStep, row: Union[Row, dict]) -> JobConf:
+    if isinstance(row, Row):
+        row = row.asDict(recursive=True)
+
+    options = row.get("options")
+    table_options = row.get("table_options")
+    check_options = row.get("check_options")
+    spark_options = row.get("spark_options")
+    invoker_options = row.get("invoker_options")
+    extender_options = row.get("extender_options")
+
+    job_id = row.get("job_id", get_job_id(step=step, topic=row["topic"], item=row["item"]))
+
+    if step in Bronzes:
+        from fabricks.core.jobs.base._types import JobConfBronze
+
+        assert options is not None, "no option"
+        step = cast(TBronze, step)
+        return JobConfBronze(
+            job_id=job_id,
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            parser_options=row.get("parser_options"),
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            extender_options=extender_options,
+            spark_options=spark_options,
+            tags=row.get("tags"),
+        )
+
+    elif step in Silvers:
+        from fabricks.core.jobs.base._types import JobConfSilver
+
+        assert options is not None, "no option"
+        step = cast(TSilver, step)
+        return JobConfSilver(
+            job_id=job_id,
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            extender_options=extender_options,
+            spark_options=spark_options,
+            tags=row.get("tags"),
+        )
+
+    elif step in Golds:
+        from fabricks.core.jobs.base._types import JobConfGold
+
+        assert options is not None, "no option"
+        step = cast(TGold, step)
+        return JobConfGold(
+            job_id=job_id,
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            extender_options=extender_options,
+            spark_options=spark_options,
+            tags=row.get("tags"),
+        )
+
+    else:
+        raise ValueError(f"{step} not found")
+
+
+@overload
+def get_job_conf(step: TStep, *, job_id: str, row: Optional[Union[Row, dict]] = None) -> JobConf: ...
+
+
+@overload
+def get_job_conf(step: TStep, *, topic: str, item: str, row: Optional[Union[Row, dict]] = None) -> JobConf: ...
+
+
+def get_job_conf(
+    step: TStep,
+    job_id: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    row: Optional[Union[Row, dict]] = None,
+) -> JobConf:
+    if row:
+        return get_job_conf_internal(step=step, row=row)
+
+    if IS_JOB_CONFIG_FROM_YAML:
+        from fabricks.core.steps import get_step
+
+        s = get_step(step=step)
+        if topic:
+            iter = s.get_jobs_iter(topic=topic)
+        else:
+            iter = s.get_jobs_iter()
+
+        if job_id:
+            conf = next(
+                (
+                    i
+                    for i in iter
+                    if i.get("job_id", get_job_id(step=i["step"], topic=i["topic"], item=i["item"])) == job_id
+                ),
+                None,
+            )
+            if not conf:
+                raise ValueError(f"job not found ({step}, {job_id})")
+
+            return get_job_conf_internal(step=step, row=conf)
+
+        elif topic and item:
+            conf = next(
+                (i for i in iter if i.get("topic") == topic and i.get("item") == item),
+                None,
+            )
+            if not conf:
+                raise ValueError(f"job not found ({step}, {topic}, {item})")
+
+            return get_job_conf_internal(step=step, row=conf)
+
+    else:
+        df = SPARK.sql(f"select * from fabricks.{step}_jobs")
+
+    assert df, f"{step} not found"
+
+    if job_id:
+        try:
+            row = df.where(f"job_id == '{job_id}'").collect()[0]
+        except IndexError:
+            raise ValueError(f"job not found ({step}, {job_id})")
+    else:
+        try:
+            row = df.where(f"topic == '{topic}' and item == '{item}'").collect()[0]
+        except IndexError:
+            raise ValueError(f"job not found ({step}, {topic}, {item})")
+
+    return get_job_conf_internal(step=step, row=row)

fabricks/core/jobs/get_job_id.py

@@ -0,0 +1,31 @@
+from typing import Optional, overload
+
+from fabricks.utils.helpers import md5
+
+
+@overload
+def get_job_id(step: str, topic: str, item: str) -> str: ...
+
+
+@overload
+def get_job_id(*, job: str) -> str: ...
+
+
+def get_job_id(
+    step: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job: Optional[str] = None,
+) -> str:
+    if not job:
+        assert step
+        assert topic
+        assert item
+        job = f"{step}.{topic}_{item}"
+
+    return md5(job)
+
+
+def get_dependency_id(parent: str, job_id: str) -> str:
+    base = f"{job_id}*{parent}"
+    return md5(base)

fabricks/core/jobs/get_jobs.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass
+from typing import List, Literal, Optional, TypedDict, Union, overload
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import expr
+from pyspark.sql.types import Row
+
+from fabricks.context import IS_JOB_CONFIG_FROM_YAML, PATHS_RUNTIME, SPARK
+from fabricks.core.jobs.base._types import AllowedModes, TStep
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.get_job import get_job, get_job_internal
+from fabricks.utils.helpers import concat_dfs, run_in_parallel
+from fabricks.utils.path import Path
+from fabricks.utils.read import read_yaml
+from fabricks.utils.schema import get_schema_for_type
+
+
+class GenericOptions(TypedDict):
+    mode: AllowedModes
+
+
+@dataclass
+class JobConfGeneric:
+    step: TStep
+    job_id: str
+    topic: str
+    item: str
+    options: GenericOptions
+
+
+def _get_job(row: Row):
+    return get_job(row=row)
+
+
+def get_jobs_internal():
+    for p in PATHS_RUNTIME.values():
+        yield from read_yaml(p, root="job")
+
+
+def get_jobs_internal_df() -> DataFrame:
+    if IS_JOB_CONFIG_FROM_YAML:
+        schema = get_schema_for_type(JobConfGeneric)
+
+        def _read_yaml(path: Path):
+            df = SPARK.createDataFrame(read_yaml(path, root="job"), schema=schema)  # type: ignore
+            if df:
+                df = df.withColumn("job_id", expr("md5(concat(step,'.',topic,'_',item))"))
+                return df
+
+        dfs = run_in_parallel(_read_yaml, list(PATHS_RUNTIME.values()))
+        df = concat_dfs(dfs)
+        assert df is not None
+
+    else:
+        df = SPARK.sql("select * from fabricks.jobs")
+
+    return df
+
+
+@overload
+def get_jobs(df: Optional[DataFrame] = None, *, convert: Literal[True]) -> List[BaseJob]: ...
+
+
+@overload
+def get_jobs(df: Optional[DataFrame] = None, *, convert: Literal[False]) -> DataFrame: ...
+
+
+def get_jobs(df: Optional[DataFrame] = None, convert: Optional[bool] = False) -> Union[List[BaseJob], DataFrame]:
+    """
+    Retrieves a list of jobs or a DataFrame containing job information.
+
+    Args:
+        df (Optional[DataFrame]): Optional DataFrame containing job information.
+        convert (Optional[bool]): Flag indicating whether to convert the DataFrame to a list of jobs.
+
+    Returns:
+        Union[List[BaseJob], DataFrame]: If `convert` is False, returns a list of BaseJob objects.
+                                         If `convert` is True, returns a DataFrame with selected columns.
+
+    Raises:
+        ValueError: If the DataFrame does not contain the required columns.
+
+    """
+    if not convert:
+        return get_jobs_internal_df()
+
+    else:
+        if df is None:
+            return list(
+                get_job_internal(j["step"], j["topic"], j["item"], j.get("job_id"), conf=j)
+                for j in get_jobs_internal()
+            )
+
+        else:
+            if "step" in df.columns and "topic" in df.columns and "item" in df.columns:
+                df = df.select("step", "topic", "item")
+            elif "step" in df.columns and "job_id" in df.columns:
+                df = df.select("step", "job_id")
+            elif "job" in df.columns:
+                df = df.select("job")
+            else:
+                raise ValueError("step, topic, item or step, job_id or job mandatory")
+
+        assert df
+
+        jobs = run_in_parallel(_get_job, df)
+        return jobs

fabricks/core/jobs/get_schedule.py

@@ -0,0 +1,10 @@
+from typing import Dict
+
+from fabricks.core.jobs.get_schedules import get_schedules
+
+
+def get_schedule(name: str) -> Dict:
+    schedule = next(s for s in get_schedules() if s.get("name") == name)
+
+    assert schedule, "schedule not found"
+    return schedule

fabricks/core/jobs/get_schedules.py

@@ -0,0 +1,32 @@
+from typing import List, Optional, TypedDict
+
+from pyspark.sql import DataFrame
+
+from fabricks.context import PATH_SCHEDULES, SPARK
+from fabricks.core.jobs.base._types import TStep
+from fabricks.utils.read.read_yaml import read_yaml
+from fabricks.utils.schema import get_schema_for_type
+
+
+class Options(TypedDict):
+    steps: Optional[List[TStep]]
+    tag: Optional[str]
+    view: Optional[str]
+    variables: Optional[dict[str, str]]
+
+
+class Schedule(TypedDict):
+    name: str
+    options: Options
+
+
+def get_schedules():
+    return read_yaml(PATH_SCHEDULES, root="schedule")
+
+
+def get_schedules_df() -> DataFrame:
+    schema = get_schema_for_type(Schedule)
+    df = SPARK.createDataFrame(list(get_schedules()), schema=schema)  # type: ignore
+
+    assert df, "no schedules found"
+    return df