From 4701cc94963fe3c19295dd221d2f295b545bb033 Mon Sep 17 00:00:00 2001 From: Fabian Reisegger Date: Tue, 16 Jan 2024 10:10:43 +0100 Subject: [PATCH] Added ability to schedule sql jobs with GC --- croud/__main__.py | 74 +++++++++++++++++++++++ croud/config/configuration.py | 21 ++++++- croud/config/schemas.py | 13 ++++ croud/gcjobs/__init__.py | 0 croud/gcjobs/commands.py | 108 ++++++++++++++++++++++++++++++++++ croud/util.py | 28 +++++++++ setup.py | 1 - 7 files changed, 243 insertions(+), 2 deletions(-) create mode 100644 croud/gcjobs/__init__.py create mode 100644 croud/gcjobs/commands.py diff --git a/croud/__main__.py b/croud/__main__.py index b8e98c7e..636cc2ae 100644 --- a/croud/__main__.py +++ b/croud/__main__.py @@ -66,6 +66,12 @@ config_set_profile, config_show, ) +from croud.gcjobs.commands import ( + create_scheduled_job, + delete_scheduled_job, + get_scheduled_job_log, + get_scheduled_jobs, +) from croud.login import login from croud.logout import logout from croud.me import me, me_edit @@ -1468,6 +1474,74 @@ }, } }, + "gc": { + "help": "Manage your Grand Central instance.", + "commands": { + "create-job": { + "help": "Create a scheduled sql job to run at specific times.", + "extra_args": [ + Argument( + "--name", type=str, required=True, help="Name of the sql job." + ), + Argument( + "--cluster-id", type=str, required=True, + help="Cluster where the job should be run." + ), + Argument( + "--cron", type=str, required=True, + help="Cron schedule of the sql job." + ), + Argument( + "--sql", type=str, required=True, + help="The sql statement the job should run." + ), + Argument( + "--enabled", type=bool, required=True, + help="Enable or disable the job." + ) + ], + "resolver": create_scheduled_job, + }, + "list-jobs": { + "help": "List all grand central scheduled jobs.", + "extra_args": [ + Argument( + "--cluster-id", type=str, required=True, + help="The cluster of which jobs should be listed." + ) + ], + "resolver": get_scheduled_jobs, + }, + "get-job-logs": { + "help": "Logs of a grand central scheduled job.", + "extra_args": [ + Argument( + "--job-id", type=str, required=True, + help="The job id of the job log to be listed." + ), + Argument( + "--cluster-id", type=str, required=True, + help="The cluster of which the job log should be listed." + ) + ], + "resolver": get_scheduled_job_log, + }, + "delete-job": { + "help": "Delete specified grand central scheduled job.", + "extra_args": [ + Argument( + "--job-id", type=str, required=True, + help="The job id of the job to be deleted." + ), + Argument( + "--cluster-id", type=str, required=True, + help="The cluster of which the job should be deleted." + ), + ], + "resolver": delete_scheduled_job, + } + } + }, "subscriptions": { "help": "Manage subscriptions.", "commands": { diff --git a/croud/config/configuration.py b/croud/config/configuration.py index c5a0a8a1..29ae9102 100644 --- a/croud/config/configuration.py +++ b/croud/config/configuration.py @@ -16,7 +16,6 @@ # However, if you have executed another commercial license agreement # with Crate these terms will supersede the license and you may use the # software solely pursuant to the terms of the relevant commercial agreement. - from pathlib import Path from typing import Any, Dict, Optional @@ -107,6 +106,14 @@ def region(self) -> Optional[str]: def organization(self) -> Optional[str]: return self.profile.get("organization-id") # type: ignore + @property + def gc_jwt_token(self) -> Optional[str]: + return self.profile.get("gc_jwt_token") + + @property + def gc_jwt_token_expiry(self) -> Optional[str]: + return self.profile.get("gc_jwt_token_expiry") + @property def profile(self) -> ProfileType: return self.profiles[self.name] # type: ignore @@ -165,6 +172,18 @@ def set_auth_token(self, profile: str, value: str) -> None: def set_current_auth_token(self, value: str) -> None: self.set_auth_token(self.name, value) + def set_gc_jwt_token(self, profile: str, value: str) -> None: + self._set_profile_option(profile, "gc_jwt_token", value) + + def set_current_gc_jwt_token(self, value: str) -> None: + self.set_gc_jwt_token(self.name, value) + + def set_gc_jwt_token_expiry(self, profile: str, value: str) -> None: + self._set_profile_option(profile, "gc_jwt_token_expiry", value) + + def set_current_gc_jwt_token_expiry(self, value: str) -> None: + self.set_gc_jwt_token_expiry(self.name, value) + def set_format(self, profile: str, value: str) -> None: self._set_profile_option(profile, "format", value) diff --git a/croud/config/schemas.py b/croud/config/schemas.py index 1239755f..a2e843e2 100644 --- a/croud/config/schemas.py +++ b/croud/config/schemas.py @@ -44,6 +44,19 @@ class ProfileSchema(Schema): attribute="organization-id", data_key="organization-id", allow_none=True ) region = fields.String(required=False) + gc_endpoint = fields.String(required=False) + gc_jwt_token = fields.String( + attribute="gc_jwt_token", + data_key="gc_jwt_token", + required=False, + allow_none=True, + ) + gc_jwt_token_expiry = fields.String( + attribute="gc_jwt_token_expiry", + data_key="gc_jwt_token_expiry", + required=False, + allow_none=True, + ) class ConfigSchema(Schema): diff --git a/croud/gcjobs/__init__.py b/croud/gcjobs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/croud/gcjobs/commands.py b/croud/gcjobs/commands.py new file mode 100644 index 00000000..f9100418 --- /dev/null +++ b/croud/gcjobs/commands.py @@ -0,0 +1,108 @@ +# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor +# license agreements. See the NOTICE file distributed with this work for +# additional information regarding copyright ownership. Crate licenses +# this file to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. You may +# obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. +# +# However, if you have executed another commercial license agreement +# with Crate these terms will supersede the license and you may use the +# software solely pursuant to the terms of the relevant commercial agreement. + +from argparse import Namespace + +from yarl import URL + +from croud.api import Client +from croud.config import CONFIG, get_output_format +from croud.printer import print_response +from croud.util import grand_central_jwt_token + + +@grand_central_jwt_token +def get_scheduled_jobs(args: Namespace) -> None: + client = _get_gc_client(args) + + data, errors = client.get("api/scheduled-jobs/") + + print_response( + data=data, + errors=errors, + keys=["name", "id", "cron", "sql", "enabled", "next_run_time"], + output_fmt=get_output_format(args), + ) + + if errors or not data: + return + + +@grand_central_jwt_token +def get_scheduled_job_log(args: Namespace) -> None: + client = _get_gc_client(args) + + data, errors = client.get(f"api/scheduled-jobs/{args.job_id}/log") + print_response( + data=data, + errors=errors, + keys=["job_id", "start", "end", "error", "statements"], + output_fmt=get_output_format(args), + ) + + if errors or not data: + return + + +@grand_central_jwt_token +def create_scheduled_job(args: Namespace) -> None: + body = { + "name": args.name, + "cron": args.cron, + "sql": args.sql, + "enabled": args.enabled, + } + + client = _get_gc_client(args) + + data, errors = client.post("api/scheduled-jobs/", body=body) + print_response( + data=data, + errors=errors, + keys=["name", "id", "cron", "sql", "enabled"], + output_fmt=get_output_format(args), + ) + + if errors or not data: + return + + +@grand_central_jwt_token +def delete_scheduled_job(args: Namespace) -> None: + client = _get_gc_client(args) + + data, errors = client.delete(f"api/scheduled-jobs/{args.job_id}") + print_response( + data=data, + errors=errors, + success_message="Scheduled job deleted.", + output_fmt=get_output_format(args), + ) + + +def _get_gc_client(args: Namespace) -> Client: + client = Client.from_args(args) + cluster, _ = client.get(f"/api/v2/clusters/{args.cluster_id}/") + + url_region_cloud = cluster.get("fqdn").split(".", 1)[1][:-1] # type: ignore + gc_url = f"https://{cluster.get('name')}.gc.{url_region_cloud}" # type: ignore + client.base_url = URL(gc_url) + client.session.cookies.set("cratedb_center_session", CONFIG.gc_jwt_token) + + return client diff --git a/croud/util.py b/croud/util.py index 4792cc11..51d472be 100644 --- a/croud/util.py +++ b/croud/util.py @@ -23,8 +23,10 @@ import subprocess import webbrowser from argparse import Namespace +from datetime import datetime, timezone from typing import Tuple +from croud.api import Client from croud.config import CONFIG from croud.printer import print_error, print_info from croud.tools.spinner import HALO @@ -126,3 +128,29 @@ def _wrapper(cmd_args: Namespace): # decorator logic cmd(cmd_args) return _wrapper + + +def grand_central_jwt_token(cmd): + @functools.wraps(cmd) + def _wrapper(cmd_args: Namespace): + # save cluster_id to config to fetch jwt_token if you specify a new cluster_id + if CONFIG.gc_jwt_token: + if ( + str(datetime.now(tz=timezone.utc).isoformat()) + > CONFIG.gc_jwt_token_expiry + ): + _set_gc_jwt(cmd_args) + else: + _set_gc_jwt(cmd_args) + + cmd(cmd_args) + + return _wrapper + + +def _set_gc_jwt(cmd_args: Namespace) -> None: + client = Client.from_args(cmd_args) + data, errors = client.get(f"/api/v2/clusters/{cmd_args.cluster_id}/jwt/") + + CONFIG.set_current_gc_jwt_token(data.get("token")) # type: ignore + CONFIG.set_current_gc_jwt_token_expiry(data.get("expiry")) # type: ignore diff --git a/setup.py b/setup.py index 29f814f2..17ca506f 100644 --- a/setup.py +++ b/setup.py @@ -64,7 +64,6 @@ "black==23.9.1", "flake8==3.8.4", "isort==5.12.0", - "mypy==0.812", ], }, python_requires=">=3.8",