def _verify_credentials(env, values):
kwargs = {
"base_url": values["cdf_base_url"],
"client_name": "function-action-validator",
"disable_pypi_version_check": True,
}
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
client = CogniteClient(api_key=values[f"cdf_{env}_credentials"],
**kwargs)
login_status = client.login.status()
inferred_project = login_status.project
if login_status.logged_in:
logger.info(
f"{env.capitalize()} API-key successfully logged in! (User/service account: '{login_status.user}', "
f"project '{inferred_project}', base-URL '{values['cdf_base_url']}')"
)
else:
raise ValueError(
f"Can't login with {env} credentials (base-URL used: '{values['cdf_base_url']}')"
)
given_project = values["cdf_project"]
if given_project is not None and inferred_project != given_project:
raise ValueError(
f"Inferred project, {inferred_project}, from the provided {env} credentials "
f"does not match the given project: {given_project}")
return inferred_project
def test_verify_ssl_enabled_by_default(self, mock_requests, mock_findall):
c = CogniteClient()
mock_requests.get.assert_called_with(_PYPI_ADDRESS, verify=True)
assert c._api_client._request_session.verify == True
assert c._api_client._request_session_with_retry.verify == True
def get_cognite_client(
self,
client_name: str,
token_custom_args: Optional[Dict[str,
str]] = None) -> CogniteClient:
kwargs = {}
if self.api_key:
kwargs["api_key"] = self.api_key
elif self.idp_authentication:
if self.idp_authentication.token_url:
kwargs["token_url"] = self.idp_authentication.token_url
elif self.idp_authentication.tenant:
base_url = urljoin(self.idp_authentication.authority,
self.idp_authentication.tenant)
kwargs["token_url"] = f"{base_url}/oauth2/v2.0/token"
kwargs["token_client_id"] = self.idp_authentication.client_id
kwargs["token_client_secret"] = self.idp_authentication.secret
kwargs["token_scopes"] = self.idp_authentication.scopes
if token_custom_args is None:
token_custom_args = {}
if self.idp_authentication.resource:
token_custom_args[
"resource"] = self.idp_authentication.resource
kwargs["token_custom_args"] = token_custom_args
else:
raise InvalidConfigError("No CDF credentials")
return CogniteClient(project=self.project,
base_url=self.host,
client_name=client_name,
disable_pypi_version_check=True,
**kwargs)
def test_version_check_disabled(self, rsps_with_login_mock):
rsps_with_login_mock.assert_all_requests_are_fired = False
with unset_env_var("COGNITE_PROJECT"):
with set_env_var("COGNITE_DISABLE_PYPI_VERSION_CHECK", "1"):
CogniteClient()
assert len(rsps_with_login_mock.calls) == 1
assert rsps_with_login_mock.calls[0].request.url.startswith("https://greenfield.cognitedata.com")
def post_data(df):
prefix = "".join(
random.choice(string.ascii_uppercase + string.digits)
for _ in range(5))
ts_names = {
"x": "tutorial_{}_x".format(prefix),
"y": "tutorial_{}_y".format(prefix)
}
df.rename(columns=ts_names, inplace=True)
client = CogniteClient()
client.time_series.post_time_series(
[TimeSeries(name=name) for name in df.columns if name != "timestamp"])
client.time_series.post_time_series(
[TimeSeries(name="tutorial_{}_mean_x_y".format(prefix))])
created_time_series = []
while len(created_time_series) != 3:
created_time_series = client.time_series.get_time_series(
prefix="tutorial_" + prefix)
sleep(0.5)
client.datapoints.post_datapoints_frame(df)
ts_ids = {ts.name.split("_", 2)[-1]: ts.id for ts in created_time_series}
return ts_ids
def __init__(self,
me=None,
project=None,
environment=None,
ts_name=None,
ts_ext_id=None):
if not me or not project or not ts_name or not ts_ext_id:
return
self.myself = me
self.project = project
self.ts_name = ts_name
self.ts_ext_id = ts_ext_id
self.ts_id = None
self.is_ok = False
if environment:
if environment == "greenfield":
url = "https://greenfield.cognitedata.com/"
self.environment = environment
else:
url = "https://api.cognitedata.com/"
self.environment = None
self.api_key = self.myself.property.cdf_api_key
self.client = CogniteClient(api_key=self.api_key,
project=project,
client_name="actingweb_fitbit_v1",
base_url=url,
debug=True)
status = self.client.login.status()
if status.logged_in:
self.is_ok = True
self.ts_id = self.check_timeseries()
else:
self.is_ok = False
def test_client_debug_mode(self, rsps_with_login_mock):
with unset_env_var("COGNITE_PROJECT"):
CogniteClient(debug=True)
log = logging.getLogger("cognite-sdk")
assert isinstance(log.handlers[0].formatter, DebugLogFormatter)
log.handlers = []
log.propagate = False
def post_data(df):
prefix = "".join(
random.choice(string.ascii_uppercase + string.digits)
for _ in range(5))
ts_names = {
"temperature": "tutorial_{}_temperature".format(prefix),
"pressure": "tutorial_{}_pressure".format(prefix),
"friction": "tutorial_{}_friction".format(prefix),
}
df.rename(columns=ts_names, inplace=True)
client = CogniteClient()
client.time_series.create(
[TimeSeries(name=name, external_id=name) for name in df.columns])
created_time_series = []
while len(created_time_series) != 3:
created_time_series = client.time_series.search(name="tutorial_" +
prefix)
sleep(0.5)
client.datapoints.insert_dataframe(df, external_id_headers=True)
ts_ids = {ts.name.split("_", 2)[-1]: ts.id for ts in created_time_series}
return ts_ids
def _use_client_credentials(cognite_client: CogniteClient, client_credentials: Optional[Dict] = None) -> Optional[str]:
"""
If client_credentials is passed, will use those, otherwise will implicitly use those the client was instantiated
with
Args:
client_credentials: a dictionary containing:
client_id
client_secret
Returns:
nonce (optional, str): a nonce if able to obtain, otherwise returns None
"""
if client_credentials:
client_id = client_credentials["client_id"]
client_secret = client_credentials["client_secret"]
else:
client_id = cognite_client.config.token_client_id
client_secret = cognite_client.config.token_client_secret
session_url = f"/api/playground/projects/{cognite_client.config.project}/sessions"
payload = {"items": [{"clientId": f"{client_id}", "clientSecret": f"{client_secret}"}]}
try:
res = cognite_client.post(session_url, json=payload)
nonce = res.json()["items"][0]["nonce"]
return nonce
except CogniteAPIError as e:
print("Unable to get nonce using client credentials flow. The session API returned with error:", e.message)
return None
def time_series_in_cdp():
global TEST_TS_1_ID, TEST_TS_2_ID
client = CogniteClient()
try:
ts_list = [TimeSeries(name=TEST_TS_1_NAME)]
client.time_series.post_time_series(ts_list)
log.warning("Posted sdk test time series 1")
client.datapoints.post_datapoints(
name=TEST_TS_1_NAME,
datapoints=[Datapoint(timestamp=i, value=i) for i in range(TEST_TS_START, TEST_TS_END, int(3.6e6))],
)
log.warning("Posted datapoints to sdk test time series 1")
TEST_TS_1_ID = client.time_series.get_time_series(prefix=TEST_TS_1_NAME).to_json()[0]["id"]
except APIError as e:
log.warning("Posting test time series 1 failed with code {}".format(e.code))
try:
ts_list = [TimeSeries(name=TEST_TS_2_NAME)]
client.time_series.post_time_series(ts_list)
log.warning("Posted sdk test time series 2")
client.datapoints.post_datapoints(
name=TEST_TS_2_NAME,
datapoints=[Datapoint(timestamp=i, value=i) for i in range(TEST_TS_START, TEST_TS_END, int(3.6e6))],
)
log.warning("Posted datapoints to sdk test time series 2")
TEST_TS_2_ID = client.time_series.get_time_series(prefix=TEST_TS_2_NAME).to_json()[0]["id"]
except APIError as e:
log.warning("Posting test time series 2 failed with code {}".format(e.code))
TEST_TS_1_ID = client.time_series.get_time_series(prefix=TEST_TS_1_NAME).to_json()[0]["id"]
TEST_TS_2_ID = client.time_series.get_time_series(prefix=TEST_TS_2_NAME).to_json()[0]["id"]
yield TEST_TS_1_ID, TEST_TS_2_ID
def getaccess(self):
from cognite.client import CogniteClient
import os
c = CogniteClient(api_key=os.environ['PUBLICDATA_API_KEY'],
client_name=os.environ['PUBLICDATA_API_KEY'])
c.login.status()
def main(args):
_configure_logger(Path(args.log), args.live, args.log_level)
api_key = args.api_key if args.api_key else os.environ.get(
"COGNITE_EXTRACTOR_API_KEY")
args.api_key = "" # Don't log the api key if given through CLI
logger.info("Extractor configured with {}".format(args))
input_path = Path(args.input)
if not input_path.exists():
logger.fatal("Input folder does not exists: {!s}".format(input_path))
sys.exit(2)
failed_path = input_path.joinpath("failed") if args.move_failed else None
finished_path = input_path.joinpath(
"finished") if args.keep_finished else None
if failed_path:
failed_path.mkdir(parents=True, exist_ok=True)
if finished_path:
finished_path.mkdir(parents=True, exist_ok=True)
try:
client = CogniteClient(api_key=api_key,
client_name="tebis-csv-datapoint-extractor")
client.login.status()
except CogniteAPIError as exc:
logger.error("Failed to create CDP client: {!s}".format(exc))
client = CogniteClient(api_key=api_key,
client_name="tebis-csv-datapoint-extractor")
project_name = client.config.project
monitor = configure_prometheus(args.live, project_name)
try:
extract_data_points(
client,
monitor,
get_all_time_series(client),
args.live,
args.from_time,
args.until_time,
input_path,
failed_path,
finished_path,
)
except KeyboardInterrupt:
logger.warning("Extractor stopped")
def file_ids():
file_ids = {}
for file in CogniteClient().files.retrieve_multiple(external_ids=[
"test/subdir/a.txt", "test/subdir/b.txt", "test/subdir/c.txt",
"test/big.txt"
]):
file_ids[file.name] = file.id
return file_ids
def ts_ids():
ts_ids = {}
for ts in CogniteClient().time_series.search(limit=100,
name="test__constant"):
short_name = re.fullmatch(r"test__(constant_[0-9]+)_with_noise",
ts.name).group(1)
ts_ids[short_name] = ts.id
return ts_ids
def _use_token_exchange(cognite_client: CogniteClient):
session_url = f"/api/playground/projects/{cognite_client.config.project}/sessions"
payload = {"items": [{"tokenExchange": True}]}
try:
res = cognite_client.post(url=session_url, json=payload)
nonce = res.json()["items"][0]["nonce"]
return nonce
except CogniteAPIError as e:
print("Unable to get nonce using token exchange flow. The session API returned with error:", e.message)
def __init__(cls):
cls.apiKey = os.environ["TEST_API_KEY_WRITE"]
cls.project = os.environ["PROJECT"]
cls.clientName = os.environ["COGNITE_CLIENT_NAME"]
cls.baseUrl = os.environ["COGNITE_BASE_URL"]
cls.client = CogniteClient(api_key=cls.apiKey,
project=cls.project,
client_name=cls.clientName,
base_url=cls.baseUrl)
def test_parameter_config(self):
base_url = "blabla"
max_workers = 1
timeout = 10
client_name = "test-client"
project = "something"
client = CogniteClient(
project=project, base_url=base_url, max_workers=max_workers, timeout=timeout, client_name=client_name
)
self.assert_config_is_correct(client._config, base_url, max_workers, timeout, client_name, project)
def test_parameter_config(self):
base_url = "blabla"
num_of_workers = 1
timeout = 10
client = CogniteClient(project="something",
base_url=base_url,
num_of_workers=num_of_workers,
timeout=timeout)
self.assert_config_is_correct(client, base_url, num_of_workers,
timeout)
def test_invalid_api_key(self, rsps):
rsps.add(rsps.GET, _PYPI_ADDRESS, status=200, body="")
rsps.add(
rsps.GET,
BASE_URL + "/login/status",
status=200,
json={"data": {"project": "", "loggedIn": False, "user": "", "projectId": -1}},
)
with unset_env_var("COGNITE_PROJECT"):
with pytest.raises(CogniteAPIKeyError):
CogniteClient()
def main(args):
_configure_logger(args.log, args.log_level)
api_key = args.api_key if args.api_key else os.environ.get(
"COGNITE_EXTRACTOR_API_KEY")
args.api_key = "" # Don't log the api key if given through CLI
logger.info("Extractor configured with {}".format(args))
if not args.input_dir.exists():
logger.fatal("Input folder does not exists: {!s}".format(
args.input_dir))
sys.exit(2)
try:
client = CogniteClient(api_key=api_key,
client_name=COGNITE_CLIENT_NAME)
logger.info(client.login.status())
except CogniteAPIError as exc:
logger.error("Failed to create CDF client: {!s}".format(exc))
client = CogniteClient(api_key=api_key,
client_name=COGNITE_CLIENT_NAME)
try:
process_path(
client,
args.input_dir,
args.pattern,
not args.non_recursive,
args.upload_to_cdf,
args.upload_to_raw,
not args.no_overwrite,
args.ignore_existing,
args.ignore_meta,
args.raw_db,
args.raw_table,
)
except KeyboardInterrupt:
logger.warning("Extractor stopped")
except Exception as exc:
logger.error(str(exc), exc_info=exc)
def test_unknown_fields_in_api_error(self, rsps):
with set_env_var("COGNITE_DISABLE_PYPI_VERSION_CHECK", "1"):
c = CogniteClient()
rsps.add(
rsps.GET,
c.assets._get_base_url_with_base_path() + "/any",
status=400,
json={
"error": {
"message": "bla",
"extra": {
"haha": "blabla"
},
"other": "yup"
}
},
)
try:
c.get(url="/api/v1/projects/{}/any".format(c.config.project))
assert False, "Call did not raise exception"
except CogniteAPIError as e:
assert {"extra": {"haha": "blabla"}, "other": "yup"} == e.extra
def create_client_and_check_config(self, i):
from cognite._thread_local import credentials
api_key = "thread-local-api-key{}".format(i)
project = "thread-local-project{}".format(i)
credentials.api_key = api_key
credentials.project = project
sleep(random.random())
client = CogniteClient()
assert api_key == client.config.api_key
assert project == client.config.project
def get_cognite_client(self, client_name: str) -> CogniteClient:
kwargs = {}
if self.api_key:
kwargs["api_key"] = self.api_key
elif self.idp_authentication:
authorizer = Authenticator(self.idp_authentication)
kwargs["token"] = authorizer.get_token
else:
raise InvalidConfigError("No CDF credentials")
return CogniteClient(project=self.project,
base_url=self.host,
client_name=client_name,
disable_pypi_version_check=True,
**kwargs)
def _authenticate(self):
"Authenticate and instantiate the Cognite client. " ""
if self._client is None:
api_key = os.environ.get('COGNITE_API_KEY')
client_name = os.environ.get('COGNITE_CLIENT_NAME')
project = os.environ.get('COGNITE_PROJECT')
for env_var in [
'COGNITE_API_KEY', 'COGNITE_CLIENT_NAME', 'COGNITE_PROJECT'
]:
if not os.environ.get(env_var):
raise ValueError(
f'Missing environment variable {env_var}.')
self._client = CogniteClient(api_key=api_key,
project=project,
client_name=client_name,
disable_pypi_version_check=True)
def setUp(self):
self.client = CogniteClient(
client_name="extractor-utils-integration-tests", )
# Delete stuff we will use if it exists
try:
self.client.raw.tables.delete(self.database_name, self.table_name)
except CogniteAPIError:
pass
try:
self.client.time_series.delete(external_id=self.time_series1)
except CogniteNotFoundError:
pass
try:
self.client.time_series.delete(external_id=self.time_series2)
except CogniteNotFoundError:
pass
def test_mock_cognite_client():
with monkeypatch_cognite_client() as c_mock:
c = CogniteClient()
assert isinstance(c_mock, MagicMock)
assert c_mock == c
api_pairs = [
(c.time_series, c_mock.time_series),
(c.raw, c_mock.raw),
(c.assets, c_mock.assets),
(c.datapoints, c_mock.datapoints),
(c.events, c_mock.events),
(c.files, c_mock.files),
(c.iam, c_mock.iam),
(c.login, c_mock.login),
(c.three_d, c_mock.three_d),
]
for api, mock_api in api_pairs:
assert isinstance(mock_api, MagicMock)
assert api == mock_api
with pytest.raises(AttributeError):
api.does_not_exist
import time
from unittest import mock
import pytest
from cognite.client import CogniteClient, utils
from cognite.client.data_classes import Asset, AssetFilter, AssetUpdate
from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError
from tests.utils import set_request_limit
COGNITE_CLIENT = CogniteClient()
@pytest.fixture
def new_asset():
ts = COGNITE_CLIENT.assets.create(Asset(name="any"))
yield ts
COGNITE_CLIENT.assets.delete(id=ts.id)
assert COGNITE_CLIENT.assets.retrieve(ts.id) is None
def generate_asset_tree(root_external_id: str,
depth: int,
children_per_node: int,
current_depth=1):
assert 1 <= children_per_node <= 10, "children_per_node must be between 1 and 10"
assets = []
if current_depth == 1:
assets = [Asset(external_id=root_external_id, name=root_external_id)]
if depth > current_depth:
for i in range(children_per_node):
import re
import pytest
from cognite.client import CogniteClient
from cognite.client._api.iam import APIKeyList, GroupList, SecurityCategoryList, ServiceAccountList
from cognite.client.data_classes import APIKey, Group, SecurityCategory, ServiceAccount
from tests.utils import jsgz_load
IAM_API = CogniteClient().iam
@pytest.fixture
def mock_service_accounts(rsps):
response_body = {
"items": [{
"name": "*****@*****.**",
"groups": [1, 2, 3],
"id": 0,
"isDeleted": False,
"deletedTime": 0
}]
}
url_pattern = re.compile(
re.escape(IAM_API._get_base_url_with_base_path()) +
"/serviceaccounts.*")
rsps.assert_all_requests_are_fired = False
rsps.add(rsps.POST, url_pattern, status=200, json=response_body)
rsps.add(rsps.GET, url_pattern, status=200, json=response_body)
yield rsps
from random import randint
import numpy as np
import pandas as pd
import pytest
from cognite.client import APIError, CogniteClient
from cognite.client.stable.raw import RawResponse, RawRow
raw = CogniteClient().raw
DB_NAME = None
TABLE_NAME = None
ROW_KEY = None
ROW_COLUMNS = None
@pytest.fixture(autouse=True, scope="class")
def db_name():
global DB_NAME
DB_NAME = "test_db_{}".format(randint(1, 2**53 - 1))
@pytest.fixture(autouse=True, scope="class")
def table_name():
global TABLE_NAME
TABLE_NAME = "test_table_{}".format(randint(1, 2**53 - 1))
@pytest.fixture(autouse=True, scope="class")
def row_key():
import re
import pytest
from cognite.client import CogniteClient
from cognite.client._api.events import Event, EventList, EventUpdate
from cognite.client.data_classes import EventFilter, TimestampRange
from tests.utils import jsgz_load
EVENTS_API = CogniteClient().events
@pytest.fixture
def mock_events_response(rsps):
response_body = {
"items": [{
"externalId": "string",
"startTime": 0,
"endTime": 0,
"description": "string",
"metadata": {
"metadata-key": "metadata-value"
},
"assetIds": [1],
"source": "string",
"id": 1,
"lastUpdatedTime": 0,
}]
}
url_pattern = re.compile(