def test_read_invalid_missing_fields(self):
# missing project
config_raw = """
# How to label uploaded data in CDF
external-id-prefix: "test_"
"""
with self.assertRaises(InvalidConfigError):
load_yaml(config_raw, CogniteConfig)
def test_read_invalid_wrong_type(self):
config_raw = """
# API key to CDF
api-key: 123
# CDF project (also known as tenant name)
project: tenant-name
# How to label uploaded data in CDF
external-id-prefix: "test_"
"""
with self.assertRaises(InvalidConfigError):
load_yaml(config_raw, CogniteConfig)
def test_read_cognite_config(self):
config_raw = """
# API key to CDF
api-key: COGNITE_API_KEY
# CDF project (also known as tenant name)
project: tenant-name
# How to label uploaded data in CDF
external-id-prefix: "test_"
"""
config = load_yaml(config_raw, CogniteConfig)
self.assertIsInstance(config, CogniteConfig)
self.assertEqual(config.host, "https://api.cognitedata.com")
self.assertEqual(config.project, "tenant-name")
self.assertEqual(config.external_id_prefix, "test_")
self.assertEqual(config.api_key, "COGNITE_API_KEY")
client = config.get_cognite_client("test-client")
self.assertIsInstance(client, CogniteClient)
self.assertEqual(client.config.base_url, "https://api.cognitedata.com")
self.assertEqual(client.config.project, "tenant-name")
self.assertEqual(client.config.api_key, "COGNITE_API_KEY")
self.assertEqual(client.config.client_name, "test-client")
def test_read_invalid_extra_fields(self):
config_raw = """
# API key to CDF
api-key: COGNITE_API_KEY
# CDF project (also known as tenant name)
project: tenant-name
# How to label uploaded data in CDF
external-id-prefix: "test_"
# Does not exist:
no-such-field: value
"""
with self.assertRaises(InvalidConfigError):
load_yaml(config_raw, CogniteConfig)
def test_get_cognite_client_no_credentials(self):
config_raw = """
project: tenant-name
external-id-prefix: "test_"
"""
config = load_yaml(config_raw, CogniteConfig)
with self.assertRaises(InvalidConfigError) as cm:
config.get_cognite_client("client_name")
self.assertEqual(str(cm.exception), "Invalid config: No CDF credentials")
def main():
logging.basicConfig(level=logging.DEBUG)
from cognite.extractorutils.configtools import load_yaml, CogniteConfig
config = load_yaml(open("config.yaml"), CogniteConfig)
cdf = config.get_cognite_client("AAD test")
print("Login status", cdf.login.status())
tss = cdf.time_series.list(limit=100)
print("Found %d time series" % len(tss))
for ts in tss[:10]:
print(" ", ts.name)
def test_get_cognite_client_from_api_key(self):
config_raw = """
api-key: COGNITE_API_KEY
project: tenant-name
external-id-prefix: "test_"
"""
config = load_yaml(config_raw, CogniteConfig)
cdf = config.get_cognite_client("client_name")
self.assertIsInstance(cdf, CogniteClient)
print("CONFIG", repr(cdf._config))
print("API_KEY", repr(cdf._config.api_key))
self.assertEqual(cdf._config.api_key, "COGNITE_API_KEY")
self.assertIsNone(cdf._config.token)
def test_get_cognite_client_from_aad(self):
config_raw = """
idp-authentication:
tenant: foo
client_id: cid
secret: scrt
scope: scp
min_ttl: 40
project: tenant-name
external-id-prefix: "test_"
"""
config = load_yaml(config_raw, CogniteConfig)
self.assertIsNone(config.api_key)
cdf = config.get_cognite_client("client_name")
self.assertIsInstance(cdf, CogniteClient)
self.assertTrue(callable(cdf._config.token))
def test_read_base_config(self):
config_raw = """
version: "1"
logger:
# Console logging
console:
level: INFO
# Information about CDF tenant
cognite:
# CDF server
host: https://greenfield.cognitedata.com
# API key to CDF
api-key: COGNITE_API_KEY
# CDF project (also known as tenant name)
project: tenant-name
# How to label uploaded data in CDF
external-id-prefix: "test_"
"""
config = load_yaml(config_raw, BaseConfig)
self.assertIsInstance(config, BaseConfig)
self.assertEqual(config.version, "1")
self.assertEqual(config.cognite.host,
"https://greenfield.cognitedata.com")
self.assertEqual(config.cognite.project, "tenant-name")
self.assertEqual(config.cognite.external_id_prefix, "test_")
self.assertEqual(config.cognite.api_key, "COGNITE_API_KEY")
self.assertEqual(config.logger.console.level, "INFO")
self.assertIsNone(config.logger.file)
created_assets = cdf.assets.create(assets)
station_to_asset_id = {}
for asset in created_assets:
weather_station = [
s for s in weather_stations if s.id == asset.metadata["station_id"]
][0]
station_to_asset_id[weather_station] = asset.id
return station_to_asset_id
if __name__ == "__main__":
with open(sys.argv[1]) as config_file:
config: WeatherConfig = load_yaml(config_file, WeatherConfig)
config.logger.setup_logging()
logger = logging.getLogger(__name__)
logger.info("Starting example Frost extractor")
frost = FrostApi(config.frost.client_id)
cdf = config.cognite.get_cognite_client("weather-extractor")
state_store = config.extractor.state_store.create_state_store(cdf)
state_store.initialize()
logger.info("Getting info about weather stations")
weather_stations = init_stations(config.locations, frost)
if config.extractor.create_assets:
def run(self, config: CsvConfig) -> None:
"""
Extract all files listed in configuration
Args:
config: Configuration parameters
"""
with ThreadPoolExecutor(max_workers=config.extractor.parallelism,
thread_name_prefix="CsvExtractor") as executor:
for file in config.files:
executor.submit(self.extract, file)
if __name__ == "__main__":
with open(sys.argv[1]) as config_file:
config = load_yaml(config_file, CsvConfig)
config.logger.setup_logging()
logger = logging.getLogger(__name__)
logger.info("Starting example CSV extractor")
metrics = Metrics()
cdf_client = config.cognite.get_cognite_client("example-csv-extractor")
if config.metrics:
config.metrics.start_pushers(cdf_client)
def upload_callback(uploaded_rows: List[Row]):
metrics.rows_uploaded.inc(len(uploaded_rows))
logging.getLogger(__name__).info(
print("An exception: {} occurred during the load balance process.".
format(error))
def upload_callback(self, uploaded_datapoints):
count = 0
for entry in uploaded_datapoints:
count += len(entry["datapoints"])
logging.getLogger(__name__).info(f"Uploaded {count} datapoints to CDF")
self.metrics.datapoints_written.inc(count)
if __name__ == "__main__":
with open(sys.argv[1]) as config_file:
config = load_yaml(config_file, IotHubConfig)
config.logger.setup_logging()
logger = logging.getLogger(__name__)
logger.info("Starting example Azure IOT Hub extractor")
cdf_client = config.cognite.get_cognite_client("azure-iot-extractor")
metrics = Metrics()
if config.metrics:
config.metrics.start_pushers(cdf_client)
extractor = IotHubExtractor(cdf_client, metrics, config)
extractor.run(config)
