-
Notifications
You must be signed in to change notification settings - Fork 13
Expand file tree
/
Copy pathbase_script.py
More file actions
124 lines (115 loc) · 4.58 KB
/
base_script.py
File metadata and controls
124 lines (115 loc) · 4.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import abc
import argparse
import os
import urllib.parse
from typing import Sequence
from exabel import ExabelClient
from exabel.scripts.command_line_script import CommandLineScript
class BaseScript(CommandLineScript, abc.ABC):
"""Base class for scripts using the Exabel Python SDK."""
def __init__(
self,
argv: Sequence[str],
description: str,
):
super().__init__(argv, description)
api_key = os.getenv("EXABEL_API_KEY")
access_token = os.getenv("EXABEL_ACCESS_TOKEN")
group = self.parser.add_mutually_exclusive_group(required=not api_key and not access_token)
group.add_argument(
"--api-key",
type=str,
help="The API key to use. Can also be specified in the EXABEL_API_KEY environment "
"variable.",
)
group.add_argument(
"--access-token",
type=str,
help="The access token to use. Can also be specified in the EXABEL_ACCESS_TOKEN "
"environment variable.",
)
group.add_argument("--gcp-project-number", type=str, help=argparse.SUPPRESS)
self.parser.add_argument(
"--exabel-data-api-host",
required=False,
default="data.api.exabel.com",
help="Data API host on format 'hostname[:port]'",
)
self.parser.add_argument(
"--exabel-analytics-api-host",
required=False,
default="analytics.api.exabel.com",
help="Analytics API host on format 'hostname[:port]'",
)
self.parser.add_argument(
"--exabel-management-api-host",
required=False,
default="management.api.exabel.com",
help="Management API host on format 'hostname[:port]'",
)
self.parser.add_argument(
"--exabel-export-api-host",
required=False,
default="export.api.exabel.com",
help="Export API host on format 'hostname[:port]'",
)
def run(self) -> None:
args = self.parse_arguments()
self.setup_logging()
api_key = args.api_key
access_token = args.access_token
# Command line argument takes precedence over environment variables.
if not api_key and not access_token and not args.gcp_project_number:
api_key = os.getenv("EXABEL_API_KEY")
access_token = os.getenv("EXABEL_ACCESS_TOKEN")
if api_key and access_token:
raise ValueError("Only one of EXABEL_API_KEY and EXABEL_ACCESS_TOKEN can be set.")
if not api_key and not access_token:
raise ValueError(
"No authentication information. Provide an API key or access token."
)
extra_headers = []
if args.gcp_project_number:
api_key = "NO_KEY"
extra_headers = [
("x-endpoint-api-consumer-type", "PROJECT"),
("x-endpoint-api-consumer-number", args.gcp_project_number),
]
data_api_host, data_api_port = self.parse_host_and_port(args.exabel_data_api_host)
analytics_api_host, analytics_api_port = self.parse_host_and_port(
args.exabel_analytics_api_host
)
management_api_host, management_api_port = self.parse_host_and_port(
args.exabel_management_api_host
)
export_api_host, export_api_port = self.parse_host_and_port(args.exabel_export_api_host)
client = ExabelClient(
data_api_host=data_api_host,
data_api_port=data_api_port,
analytics_api_host=analytics_api_host,
analytics_api_port=analytics_api_port,
management_api_host=management_api_host,
management_api_port=management_api_port,
export_api_host=export_api_host,
export_api_port=export_api_port,
api_key=api_key,
access_token=access_token,
extra_headers=extra_headers,
retries=getattr(args, "retries", None),
)
self.run_script(client, args)
@abc.abstractmethod
def run_script(self, client: ExabelClient, args: argparse.Namespace) -> None:
"""
Runs the script.
Args:
client: Exabel API client.
args: Parsed command line arguments.
"""
@staticmethod
def parse_host_and_port(host_and_port: str) -> tuple[str, int | None]:
"""
Parses a string on format 'hostname[:port]'.
"""
result = urllib.parse.urlsplit("//" + host_and_port)
return str(result.hostname), result.port