fabricks 3.0.11__py3-none-any.whl

fabricks/core/dags/processor.py

@@ -0,0 +1,228 @@
+import json
+import threading
+import time
+from multiprocessing import Process
+from typing import Any, List, Union
+
+from azure.core.exceptions import AzureError
+from databricks.sdk.runtime import dbutils, spark
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+
+from fabricks.context import PATH_NOTEBOOKS
+from fabricks.core.dags.base import BaseDags
+from fabricks.core.dags.log import LOGGER
+from fabricks.core.dags.run import run
+from fabricks.core.jobs.base._types import TStep
+from fabricks.core.steps.get_step import get_step
+from fabricks.utils.azure_queue import AzureQueue
+from fabricks.utils.azure_table import AzureTable
+
+
+class DagProcessor(BaseDags):
+    def __init__(self, schedule_id: str, schedule: str, step: Union[TStep, str], notebook: bool = True):
+        self.step = get_step(step=step)
+        self.schedule = schedule
+
+        self.notebook = notebook
+
+        self._azure_queue = None
+        self._azure_table = None
+
+        super().__init__(schedule_id=schedule_id)
+
+    @property
+    def queue(self) -> AzureQueue:
+        if not self._azure_queue:
+            step = self.remove_invalid_characters(str(self.step))
+            self._azure_queue = AzureQueue(
+                f"q{step}{self.schedule_id}",
+                **self.get_connection_info(),  # type: ignore
+            )
+        return self._azure_queue
+
+    @property
+    def table(self) -> AzureTable:
+        if not self._azure_table:
+            self._azure_table = AzureTable(f"t{self.schedule_id}", **self.get_connection_info())  # type: ignore
+        return self._azure_table
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception, AzureError)),
+        reraise=True,
+    )
+    def query(self, data: Any) -> List[dict]:
+        return self.table.query(data)
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception, AzureError)),
+        reraise=True,
+    )
+    def upsert(self, data: Any) -> None:
+        self.table.upsert(data)
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception, AzureError)),
+        reraise=True,
+    )
+    def delete(self, data: Any) -> None:
+        self.table.delete(data)
+
+    def extra(self, d: dict) -> dict:
+        return {
+            "partition_key": self.schedule_id,
+            "schedule": self.schedule,
+            "schedule_id": self.schedule_id,
+            "step": str(self.step),
+            "job": d.get("Job"),
+            "target": "table",
+        }
+
+    def send(self):
+        while True:
+            scheduled = self.get_scheduled()
+            assert isinstance(scheduled, List)
+            if len(scheduled) == 0:
+                for _ in range(self.step.workers):
+                    self.queue.send_sentinel()
+                LOGGER.info("no more job to schedule", extra={"label": str(self.step)})
+                break
+
+            else:
+                sorted_scheduled = sorted(scheduled, key=lambda x: x.get("Rank"))
+                for s in sorted_scheduled:
+                    dependencies = self.table.query(f"PartitionKey eq 'dependencies' and JobId eq '{s.get('JobId')}'")
+
+                    if len(dependencies) == 0:
+                        s["Status"] = "waiting"
+                        LOGGER.debug("waiting", extra=self.extra(s))
+                        self.table.upsert(s)
+                        self.queue.send(s)
+
+            time.sleep(5)
+
+    def receive(self):
+        while True:
+            response = self.queue.receive()
+            if response == self.queue.sentinel:
+                LOGGER.info("no more job to process", extra={"label": str(self.step)})
+                break
+
+            elif response:
+                j = json.loads(response)
+
+                j["Status"] = "starting"
+                self.table.upsert(j)
+                LOGGER.info("start", extra=self.extra(j))
+
+                try:
+                    if self.notebook:
+                        dbutils.notebook.run(  # type: ignore
+                            PATH_NOTEBOOKS.joinpath("run").get_notebook_path(),  # type: ignore
+                            self.step.timeouts.job,  # type: ignore
+                            {
+                                "schedule_id": self.schedule_id,
+                                "schedule": self.schedule,  # needed to pass schedule variables to the job
+                                "step": str(self.step),
+                                "job_id": j.get("JobId"),
+                                "job": j.get("Job"),
+                            },  # type: ignore
+                        )
+
+                    else:
+                        run(
+                            step=str(self.step),
+                            job_id=j.get("JobId"),
+                            schedule_id=self.schedule_id,
+                            schedule=self.schedule,
+                        )
+
+                except Exception:
+                    LOGGER.warning("fail", extra={"label": j.get("Job")})
+
+                finally:
+                    j["Status"] = "ok"
+                    self.table.upsert(j)
+                    LOGGER.info("end", extra=self.extra(j))
+
+                dependencies = self.table.query(f"PartitionKey eq 'dependencies' and ParentId eq '{j.get('JobId')}'")
+                self.table.delete(dependencies)
+
+    def get_scheduled(self, convert: bool = False):
+        scheduled = self.table.query(f"PartitionKey eq 'statuses' and Status eq 'scheduled' and Step eq '{self.step}'")
+        if convert:
+            return spark.createDataFrame(scheduled)
+        else:
+            return scheduled
+
+    def _process(self):
+        scheduled = self.get_scheduled()
+        assert isinstance(scheduled, List)
+
+        if len(scheduled) > 0:
+            sender = threading.Thread(
+                target=self.send,
+                name=f"{str(self.step).capitalize()}Sender",
+                args=(),
+            )
+            sender.start()
+
+            receivers = []
+            for i in range(self.step.workers):
+                receiver = threading.Thread(
+                    target=self.receive,
+                    name=f"{str(self.step).capitalize()}Receiver{i}",
+                    args=(),
+                )
+                receiver.start()
+                receivers.append(receiver)
+
+            sender.join()
+            for receiver in receivers:
+                receiver.join()
+
+    def process(self):
+        scheduled = self.get_scheduled()
+        assert isinstance(scheduled, List)
+
+        if len(scheduled) > 0:
+            LOGGER.info("start", extra={"label": str(self.step)})
+
+            p = Process(target=self._process())
+            p.start()
+            p.join(timeout=self.step.timeouts.step)
+            p.terminate()
+
+            self.queue.delete()
+
+            if p.exitcode is None:
+                LOGGER.critical("timeout", extra={"label": str(self.step)})
+                raise ValueError(f"{self.step} timed out")
+
+            else:
+                df = self.get_logs(str(self.step))
+                self.write_logs(df)
+
+                LOGGER.info("end", extra={"label": str(self.step)})
+
+        else:
+            LOGGER.info("no job to schedule", extra={"label": str(self.step)})
+
+    def __str__(self) -> str:
+        return f"{str(self.step)} ({self.schedule_id})"
+
+    def __enter__(self):
+        return super().__enter__()
+
+    def __exit__(self, *args, **kwargs):
+        if self._azure_queue:
+            self._azure_queue.__exit__()
+        if self._azure_table:
+            self._azure_table.__exit__()
+
+        return super().__exit__(*args, **kwargs)

fabricks/core/dags/run.py

@@ -0,0 +1,39 @@
+from typing import Optional
+
+from fabricks.core.dags.log import LOGGER, TABLE_LOG_HANDLER
+from fabricks.core.jobs import get_job
+from fabricks.core.jobs.base.exception import CheckWarning, SkipRunCheckWarning
+
+
+def run(step: str, job_id: str, schedule_id: str, schedule: str, notebook_id: Optional[str] = None):
+    job = get_job(step=step, job_id=job_id)
+
+    extra = {
+        "partition_key": schedule_id,
+        "schedule_id": schedule_id,
+        "schedule": schedule,
+        "step": step,
+        "job": job,
+        "target": "buffer",
+    }
+    if notebook_id is not None:
+        extra["notebook_id"] = notebook_id
+
+    LOGGER.info("running", extra=extra)
+
+    try:
+        job.run(schedule_id=schedule_id, schedule=schedule)
+        LOGGER.info("done", extra=extra)
+
+    except SkipRunCheckWarning:
+        LOGGER.exception("skipped", extra=extra)
+
+    except CheckWarning:
+        LOGGER.exception("warned", extra=extra)
+
+    except Exception as e:
+        LOGGER.exception("failed", extra=extra)
+        raise e
+
+    finally:
+        TABLE_LOG_HANDLER.flush()

fabricks/core/dags/terminator.py

@@ -0,0 +1,25 @@
+from fabricks.context import SPARK
+from fabricks.core.dags.base import BaseDags
+from fabricks.core.dags.log import LOGGER, TABLE_LOG_HANDLER
+
+
+class DagTerminator(BaseDags):
+    def __init__(self, schedule_id: str):
+        self.schedule_id = schedule_id
+        super().__init__(schedule_id=schedule_id)
+
+    def terminate(self):
+        df = self.get_logs()
+        self.write_logs(df)
+
+        error_df = SPARK.sql("select * from {df} where status = 'failed'", df=df)
+        for row in error_df.collect():
+            LOGGER.error(f"{row['job']} failed (🔥)")
+
+        TABLE_LOG_HANDLER.table.truncate_partition(self.schedule_id)
+
+        table = self.get_table()
+        table.drop()
+
+        if not error_df.isEmpty():
+            raise ValueError(f"{error_df.count()} job(s) failed")

fabricks/core/dags/utils.py

@@ -0,0 +1,54 @@
+from typing import Optional, cast
+
+from fabricks.context import DBUTILS, FABRICKS_STORAGE, FABRICKS_STORAGE_CREDENTIAL, IS_UNITY_CATALOG, SECRET_SCOPE
+from fabricks.utils.azure_table import AzureTable
+
+
+def _get_access_key_from_secret_scope(storage_account: str) -> str:
+    from fabricks.context.secret import AccessKey, get_secret_from_secret_scope
+
+    secret = get_secret_from_secret_scope(SECRET_SCOPE, f"{storage_account}-access-key")
+    return cast(AccessKey, secret).key
+
+
+def _get_access_key_from_os() -> Optional[str]:
+    import os
+
+    return os.environ.get("FABRICKS_ACCESS_KEY")
+
+
+def get_connection_info(storage_account: str) -> dict:
+    credential = None
+
+    if not IS_UNITY_CATALOG:
+        access_key = _get_access_key_from_secret_scope(storage_account)
+
+    else:
+        access_key = _get_access_key_from_os()
+        if not access_key:
+            access_key = _get_access_key_from_secret_scope(storage_account)
+
+        if FABRICKS_STORAGE_CREDENTIAL:
+            assert DBUTILS
+            credential = DBUTILS.credentials.getServiceCredentialsProvider(FABRICKS_STORAGE_CREDENTIAL)  # type: ignore
+
+        assert credential or access_key
+
+    return {
+        "storage_account": storage_account,
+        "access_key": access_key,
+        "credential": credential,
+    }
+
+
+def get_table():
+    storage_account = FABRICKS_STORAGE.get_storage_account()
+
+    cx = get_connection_info(storage_account)
+
+    return AzureTable(
+        "dags",
+        storage_account=storage_account,
+        access_key=cx["access_key"],
+        credential=cx["credential"],
+    )

fabricks/core/extenders.py

@@ -0,0 +1,33 @@
+from importlib.util import module_from_spec, spec_from_file_location
+from typing import Callable
+
+from fabricks.context import IS_UNITY_CATALOG, PATH_EXTENDERS
+from fabricks.context.log import DEFAULT_LOGGER
+
+EXTENDERS: dict[str, Callable] = {}
+
+
+def get_extender(name: str) -> Callable:
+    path = PATH_EXTENDERS.joinpath(f"{name}.py")
+    if not IS_UNITY_CATALOG:
+        assert path.exists(), "no valid extender found in {path.string}"
+    else:
+        DEFAULT_LOGGER.debug(f"could not check if extender exists ({path.string})")
+
+    spec = spec_from_file_location(name, path.string)
+    assert spec, "no valid extender found in {path.string}"
+    assert spec.loader is not None
+
+    mod = module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    e = EXTENDERS[name]
+
+    return e
+
+
+def extender(name: str):
+    def decorator(fn: Callable):
+        EXTENDERS[name] = fn
+        return fn
+
+    return decorator

fabricks/core/job_schema.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import List
+
+from fabricks.core.jobs.base._types import JobConf
+from fabricks.utils.schema import get_json_schema_for_type
+
+
+def get_job_schema() -> str:
+    import json
+
+    @dataclass
+    class JobWrapper:
+        job: JobConf
+
+    sc = get_json_schema_for_type(List[JobWrapper])
+    defs: dict[str, dict] = sc["$defs"]
+    removals = [("Job", "job_id"), ("Job", "table")]
+
+    for key, defi in defs.items():
+        for ent, prop in removals:
+            if key.startswith(ent) and prop in defi["properties"]:
+                req: List[str] = defi["required"]
+                req.remove(prop)  # not defined in yaml
+                jobprops: dict = defi["properties"]
+                jobprops.pop(prop)
+
+    j = json.dumps(sc, indent=4)
+    return j
+
+
+def print_job_schema():
+    print(get_job_schema())

fabricks/core/jobs/__init__.py

@@ -0,0 +1,21 @@
+from fabricks.core.jobs.base import BaseJob, Bronzes, Golds, Silvers, Steps
+from fabricks.core.jobs.bronze import Bronze
+from fabricks.core.jobs.get_job import get_job
+from fabricks.core.jobs.get_job_id import get_job_id
+from fabricks.core.jobs.get_jobs import get_jobs
+from fabricks.core.jobs.gold import Gold
+from fabricks.core.jobs.silver import Silver
+
+__all__ = [
+    "BaseJob",
+    "Bronze",
+    "Bronzes",
+    "get_job_id",
+    "get_job",
+    "get_jobs",
+    "Gold",
+    "Golds",
+    "Silver",
+    "Silvers",
+    "Steps",
+]

fabricks/core/jobs/base/__init__.py

@@ -0,0 +1,10 @@
+from fabricks.core.jobs.base._types import Bronzes, Golds, Silvers, Steps
+from fabricks.core.jobs.base.job import BaseJob
+
+__all__ = [
+    "BaseJob",
+    "Bronzes",
+    "Golds",
+    "Silvers",
+    "Steps",
+]