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)