Skip to content

Commit

Permalink
Added ability to schedule sql jobs with GC
Browse files Browse the repository at this point in the history
  • Loading branch information
Taliik committed Jan 22, 2024
1 parent cdc6dee commit 7e380f5
Show file tree
Hide file tree
Showing 10 changed files with 587 additions and 1 deletion.
2 changes: 2 additions & 0 deletions CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ Changes for croud
Unreleased
==========

- Added support for scheduling sql jobs with the ``scheduled-jobs`` commands.

1.10.1 - 2024/01/11
===================

Expand Down
74 changes: 74 additions & 0 deletions croud/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,12 @@
project_users_remove,
)
from croud.regions.commands import regions_create, regions_delete, regions_list
from croud.scheduledjobs.commands import (
create_scheduled_job,
delete_scheduled_job,
get_scheduled_job_log,
get_scheduled_jobs,
)
from croud.subscriptions.commands import (
subscription_delete,
subscriptions_create,
Expand Down Expand Up @@ -1468,6 +1474,74 @@
},
}
},
"scheduled-jobs": {
"help": "Manage your scheduled sql jobs.",
"commands": {
"create": {
"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": {
"help": "Get all scheduled sql jobs.",
"extra_args": [
Argument(
"--cluster-id", type=str, required=True,
help="The cluster of which jobs should be listed."
)
],
"resolver": get_scheduled_jobs,
},
"logs": {
"help": "Logs of a scheduled sql 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": {
"help": "Delete specified scheduled sql 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": {
Expand Down
31 changes: 30 additions & 1 deletion croud/config/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -107,6 +106,18 @@ 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 gc_cluster_id(self) -> Optional[str]:
return self.profile.get("gc_cluster_id")

@property
def profile(self) -> ProfileType:
return self.profiles[self.name] # type: ignore
Expand Down Expand Up @@ -165,6 +176,24 @@ 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_gc_cluster_id(self, profile: str, value: str) -> None:
self._set_profile_option(profile, "gc_cluster_id", value)

def set_current_gc_cluster_id(self, value: str) -> None:
self.set_gc_cluster_id(self.name, value)

def set_format(self, profile: str, value: str) -> None:
self._set_profile_option(profile, "format", value)

Expand Down
19 changes: 19 additions & 0 deletions croud/config/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,25 @@ 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,
)
gc_cluster_id = fields.String(
attribute="gc_cluster_id",
data_key="gc_cluster_id",
required=False,
allow_none=True,
)


class ConfigSchema(Schema):
Expand Down
Empty file added croud/scheduledjobs/__init__.py
Empty file.
108 changes: 108 additions & 0 deletions croud/scheduledjobs/commands.py
Original file line number Diff line number Diff line change
@@ -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
30 changes: 30 additions & 0 deletions croud/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -126,3 +128,31 @@ 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):
if CONFIG.gc_jwt_token:
if not CONFIG.gc_cluster_id == cmd_args.cluster_id:
_set_gc_jwt(cmd_args)
elif (
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_cluster_id(cmd_args.cluster_id) # type: ignore
CONFIG.set_current_gc_jwt_token_expiry(data.get("expiry")) # type: ignore
1 change: 1 addition & 0 deletions docs/commands/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ To get help for a specific command, you can append ``--help``:
users
api-keys
regions
gcjobs

.. note::

Expand Down
Loading

0 comments on commit 7e380f5

Please sign in to comment.