Skip to content

Commit

Permalink
Merge pull request #10 from chaos-genius/chore-removed-unused-code
Browse files Browse the repository at this point in the history
fix: readded job list
  • Loading branch information
kartikay-bagla authored Oct 7, 2024
2 parents bfe3164 + b6f1964 commit 5ce0426
Showing 1 changed file with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions databricks/sdk/chaosgenius/data_puller.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from databricks.sdk.chaosgenius.cg_config import CGConfig
from databricks.sdk.service.compute import InstancePoolAndStats
from databricks.sdk.service.iam import User
from databricks.sdk.service.jobs import BaseJob
from databricks.sdk.service.sql import EndpointInfo


Expand Down Expand Up @@ -67,6 +68,10 @@ def __init__(
self._wh_list = self._get_full_warehouse_info()
self._logger.info(f"Total warehouses: {len(self._wh_list)}")

self._logger.info("Getting jobs list")
self._job_list = self._get_full_jobs_info()
self._logger.info(f"Total jobs: {len(self._job_list)}")

self._logger.info("Getting users list")
self._user_list = self._get_full_user_info()
self._logger.info(f"Total users: {len(self._user_list)}")
Expand Down Expand Up @@ -156,6 +161,15 @@ def _get_full_warehouse_info(self) -> list[EndpointInfo]:
"id",
)

def _get_full_jobs_info(self) -> list[BaseJob]:
self._logger.info("Getting workspace jobs.")
return self._generic_get_full_list(
"job",
self._workspace_client.jobs.list,
self._workspace_client.jobs.get,
"job_id",
)

def _get_full_user_info(self) -> list[User]:
self._logger.info("Getting workspace users.")
return self._generic_get_full_list(
Expand All @@ -165,6 +179,25 @@ def _get_full_user_info(self) -> list[User]:
"id",
)

def get_jobs_list(self) -> bool:
self._logger.info("Saving jobs list.")
self._add_status_entry("jobs", "initializing", {})
try:
jobs_df = pd.DataFrame(
[
{"job_id": job.job_id, "data": json.dumps(job.as_dict())}
for job in self._job_list
]
)
if not jobs_df.empty:
self._write_to_table(jobs_df, "jobs_list")
self._add_status_entry("jobs", "success", {})
return True
except Exception:
self._logger.exception("Saving jobs failed :(")
self._add_status_entry("jobs", "failed", {})
return False

def _add_status_entry(self, module: str, status: str, data: dict):
self._write_to_table(
df=pd.DataFrame(
Expand Down Expand Up @@ -269,6 +302,7 @@ def get_all(self) -> list[tuple[str, bool]]:
data = [
("instance pools", self.get_instance_pools_list),
("warehouses list", self.get_sql_warehouses_list),
("jobs list", self.get_jobs_list),
("users list", self.get_users_list),
]

Expand Down

0 comments on commit 5ce0426

Please sign in to comment.