def __init__(
self,
host,
port,
project,
region_name,
secrets_store,
hostname_verification,
trust_store_path,
cert_folder,
api_key_file,
api_key_value,
):
"""Initializes a client in an external environment such as AWS Sagemaker."""
if not host:
raise exceptions.ExternalClientError("host")
if not project:
raise exceptions.ExternalClientError("project")
self._host = host
self._port = port
self._base_url = "https://" + self._host + ":" + str(self._port)
self._project_name = project
self._region_name = region_name or self.DEFAULT_REGION
if api_key_value is not None:
api_key = api_key_value
else:
api_key = self._get_secret(secrets_store, "api-key", api_key_file)
self._auth = auth.ApiKeyAuth(api_key)
self._session = requests.session()
self._connected = True
self._verify = self._get_verify(self._host, trust_store_path)
project_info = self._get_project_info(self._project_name)
self._project_id = str(project_info["projectId"])
if cert_folder:
# On external Spark clients (Databricks, Spark Cluster),
# certificates need to be provided before the Spark application starts.
self._cert_folder_base = cert_folder
self._cert_folder = os.path.join(cert_folder, host, project)
os.makedirs(self._cert_folder, exist_ok=True)
credentials = self._get_credentials(self._project_id)
self._write_b64_cert_to_bytes(
str(credentials["kStore"]),
path=os.path.join(self._cert_folder, "keyStore.jks"),
)
self._write_b64_cert_to_bytes(
str(credentials["tStore"]),
path=os.path.join(self._cert_folder, "trustStore.jks"),
)
self._cert_key = str(credentials["password"])
with open(os.path.join(self._cert_folder, "material_passwd"),
"w") as f:
f.write(str(credentials["password"]))
def _get_secret(self, secrets_store, secret_key=None, api_key_file=None):
"""Returns secret value from the AWS Secrets Manager or Parameter Store.
:param secrets_store: the underlying secrets storage to be used, e.g. `secretsmanager` or `parameterstore`
:type secrets_store: str
:param secret_key: key for the secret value, e.g. `api-key`, `cert-key`, `trust-store`, `key-store`, defaults to None
:type secret_key: str, optional
:param api_key_file: path to a file containing an api key, defaults to None
:type api_key_file: str optional
:raises ExternalClientError: `api_key_file` needs to be set for local mode
:raises UnkownSecretStorageError: Provided secrets storage not supported
:return: secret
:rtype: str
"""
if secrets_store == self.SECRETS_MANAGER:
return self._query_secrets_manager(secret_key)
elif secrets_store == self.PARAMETER_STORE:
return self._query_parameter_store(secret_key)
elif secrets_store == self.LOCAL_STORE:
if not api_key_file:
raise exceptions.ExternalClientError(
"api_key_file needs to be set for local mode")
with open(api_key_file) as f:
return f.readline().strip()
else:
raise exceptions.UnknownSecretStorageError("Secrets storage " +
secrets_store +
" is not supported.")
def __init__(
self,
host,
port,
project,
engine,
region_name,
secrets_store,
hostname_verification,
trust_store_path,
cert_folder,
api_key_file,
api_key_value,
):
"""Initializes a client in an external environment such as AWS Sagemaker."""
if not host:
raise exceptions.ExternalClientError("host")
if not project:
raise exceptions.ExternalClientError("project")
self._host = host
self._port = port
self._base_url = "https://" + self._host + ":" + str(self._port)
self._project_name = project
self._region_name = region_name or self.DEFAULT_REGION
if api_key_value is not None:
api_key = api_key_value
else:
api_key = self._get_secret(secrets_store, "api-key", api_key_file)
self._auth = auth.ApiKeyAuth(api_key)
self._session = requests.session()
self._connected = True
self._verify = self._get_verify(self._host, trust_store_path)
project_info = self._get_project_info(self._project_name)
self._project_id = str(project_info["projectId"])
self._cert_key = None
self._cert_folder_base = None
if engine == "python":
# On external Spark clients (Databricks, Spark Cluster),
# certificates need to be provided before the Spark application starts.
self._cert_folder_base = cert_folder
self._cert_folder = os.path.join(cert_folder, host, project)
self._trust_store_path = os.path.join(self._cert_folder, "trustStore.jks")
self._key_store_path = os.path.join(self._cert_folder, "keyStore.jks")
os.makedirs(self._cert_folder, exist_ok=True)
credentials = self._get_credentials(self._project_id)
self._write_b64_cert_to_bytes(
str(credentials["kStore"]),
path=self._get_jks_key_store_path(),
)
self._write_b64_cert_to_bytes(
str(credentials["tStore"]),
path=self._get_jks_trust_store_path(),
)
self._write_pem_file(credentials["caChain"], self._get_ca_chain_path())
self._write_pem_file(
credentials["clientCert"], self._get_client_cert_path()
)
self._write_pem_file(credentials["clientKey"], self._get_client_key_path())
self._cert_key = str(credentials["password"])
with open(os.path.join(self._cert_folder, "material_passwd"), "w") as f:
f.write(str(credentials["password"]))
elif engine == "spark":
_spark_session = SparkSession.builder.getOrCreate()
self.validate_spark_configuration(_spark_session)
with open(
_spark_session.conf.get("spark.hadoop.hops.ssl.keystores.passwd.name"),
"r",
) as f:
self._cert_key = f.read()
self._trust_store_path = _spark_session.conf.get(
"spark.hadoop.hops.ssl.trustore.name"
)
self._key_store_path = _spark_session.conf.get(
"spark.hadoop.hops.ssl.keystore.name"
)