def initialize_client() -> Optional[InfluxDBClient]: influxdb_settings = InfluxdbIntegrationSettings.get_solo() if not influxdb_settings.enabled: logger.debug( 'INFLUXDB: Integration disabled in settings (or due to an error previously)' ) return None use_secure_connection = influxdb_settings.secure in ( InfluxdbIntegrationSettings.SECURE_CERT_NONE, InfluxdbIntegrationSettings.SECURE_CERT_REQUIRED, ) if use_secure_connection: server_base_url = 'https://{}:{}'.format(influxdb_settings.hostname, influxdb_settings.port) else: server_base_url = 'http://{}:{}'.format(influxdb_settings.hostname, influxdb_settings.port) logger.debug('INFLUXDB: Initializing InfluxDB client for "%s"', server_base_url) influxdb_client = InfluxDBClient( url=server_base_url, token=influxdb_settings.api_token, verify_ssl=influxdb_settings.secure == InfluxdbIntegrationSettings.SECURE_CERT_REQUIRED, timeout=settings.DSMRREADER_CLIENT_TIMEOUT * 1000, # Ms! ) # logger.debug('INFLUXDB: InfluxDB client/server status: "%s"', influxdb_client.ready().status) if influxdb_client.buckets_api().find_bucket_by_name( influxdb_settings.bucket) is None: # pragma: nocover logger.debug('INFLUXDB: Creating InfluxDB bucket "%s"', influxdb_settings.bucket) try: influxdb_client.buckets_api().create_bucket( bucket_name=influxdb_settings.bucket, org=influxdb_settings.organization) except Exception as e: InfluxdbIntegrationSettings.objects.update(enabled=False) logger.error( 'Failed to instantiate InfluxDB connection, disabling InfluxDB integration' ) raise e return influxdb_client
class BaseTest(unittest.TestCase): def setUp(self) -> None: self.conf = influxdb_client.configuration.Configuration() self.host = os.getenv('INFLUXDB_V2_URL', "http://localhost:8086") self.debug = False self.auth_token = os.getenv('INFLUXDB_V2_TOKEN', "my-token") self.org = os.getenv('INFLUXDB_V2_ORG', "my-org") self.client = InfluxDBClient(url=self.host, token=self.auth_token, debug=self.debug, org=self.org) self.api_client = self.client.api_client self.query_api = self.client.query_api() self.buckets_api = self.client.buckets_api() self.users_api = self.client.users_api() self.organizations_api = self.client.organizations_api() self.authorizations_api = self.client.authorizations_api() self.labels_api = self.client.labels_api() self.my_organization = self.find_my_org() def tearDown(self) -> None: self.client.close() def create_test_bucket(self): bucket_name = generate_bucket_name() bucket = self.buckets_api.create_bucket(bucket_name=bucket_name, org=self.my_organization, description=bucket_name + "description") return bucket def delete_test_bucket(self, bucket): return self.buckets_api.delete_bucket(bucket) def find_my_org(self) -> Organization: return self.client.organizations_api().find_organizations( org=self.org)[0] @staticmethod def log(args): print(">>>", args) @staticmethod def generate_name(prefix): assert prefix != "" or prefix is not None return prefix + str(datetime.datetime.now().timestamp()) + "-IT" @classmethod def retention_rule(cls) -> BucketRetentionRules: return BucketRetentionRules(type='expire', every_seconds=3600) def assertEqualIgnoringWhitespace(self, first, second, msg=None) -> None: whitespace_pattern = re.compile(r"\s+") self.assertEqual(whitespace_pattern.sub("", first), whitespace_pattern.sub("", second), msg=msg)
class InfluxDBLogger: def __init__( self, bucket_name=BUCKET, batch_size=LOG_BATCH_SIZE, data_retention=3600, ): self.organization = ORGANIZATION self.client = InfluxDBClient(url=INFLUXDB_URL, token=INFLUXDB_TOKEN, org=self.organization) self.batch_size = batch_size self.bucket_name = bucket_name self.write_api = self.client.write_api(write_options=WriteOptions( batch_size=self.batch_size)) self.query_api = self.client.query_api() self.buckets_api = self.client.buckets_api() bucket = self.buckets_api.find_bucket_by_name(self.bucket_name) if bucket is None: logger.warning(f"Bucket {self.bucket_name!r} not found. " f"Creating a bucket {self.bucket_name!r}.") retention_rules = None if data_retention is not None: retention_rules = BucketRetentionRules( type="expire", every_seconds=data_retention) self.buckets_api.create_bucket( bucket_name=self.bucket_name, retention_rules=retention_rules, org=self.organization, ) def send_event(self, record_type, message): point = Point(record_type) for key, value in message.items(): point = point.field(key, value) self.write_api.write(bucket=self.bucket_name, record=point) def get_events(self, record_type): query = ''' from(bucket: currentBucket) |> range(start: -5m, stop: now()) |> filter(fn: (r) => r._measurement == recordType) |> pivot(rowKey:["_time"], columnKey: ["_field"], \ valueColumn: "_value") ''' params = {"currentBucket": self.bucket_name, "recordType": record_type} tables = self.query_api.query(query=query, params=params) if len(tables) > 0: table, *_ = tables events = table.records else: events = [] return events
def _connect(self): while not self.write_api: client = InfluxDBClient(url=self.url, token=self.auth_token, org=self.org) # get the orgID from the name: try: organizations_api = client.organizations_api() orgs = organizations_api.find_organizations() except: self.client = None logging.warning('Error connecting to the InfluxDB API. ' 'Please confirm that InfluxDB is running and ' 'that the authentication token is correct.' 'Sleeping before trying again.') time.sleep(5) continue # Look up the organization id for our org our_org = next((org for org in orgs if org.name == self.org), None) if not our_org: logging.fatal('Can not find org "%s" in InfluxDB', self.org) raise RuntimeError('Can not find org "%s" in InfluxDB' % self.org) self.org_id = our_org.id # get the bucketID from the name: bucket_api = client.buckets_api() bucket = bucket_api.find_bucket_by_name(self.bucket_name) # if the bucket does not exist then try to create it if bucket: self.bucket_id = bucket.id else: try: logging.info('Creating new bucket for: %s', self.bucket_name) new_bucket = bucket_api.create_bucket( bucket_name=self.bucket_name, org_id=self.org_id) self.bucket_id = new_bucket.id except: logging.fatal('Can not create InfluxDB bucket "%s"', self.bucket_name) raise RuntimeError('Can not create InfluxDB bucket "%s"' % self.bucket_name) self.write_api = client.write_api(write_options=ASYNCHRONOUS)
def setup(request): print("Load testing data") client = InfluxDBClient(url=url, token=token, org=my_org, debug=False) write_api = client.write_api() buckets_api = client.buckets_api() query_api = client.query_api() org = find_my_org(client, my_org) bucket = buckets_api.find_bucket_by_name(bucket_name=test_bucket) if bucket is not None: buckets_api.delete_bucket(bucket=bucket) bucket = buckets_api.create_bucket(bucket_name=test_bucket, org_id=org.id) num_records = 10 num_series = 10 today = datetime.utcnow() print("*** Write test series ***") for loc in range(num_series): for i in range(num_records): time_ = today - timedelta(hours=i + 1) point = Point(measurement_name="h2o_feet") \ .time(time_) \ .field("water_level", float(i)) \ .tag("location", "location_" + str(loc)) \ .tag("country", "country_" + str(loc)) write_api.write(bucket=bucket.name, record=point) time.sleep(1) query = 'from(bucket:"{0}") |> range(start: 1970-01-01T00:00:00.000000001Z)'.format( test_bucket) flux_result = query_api.query(query) assert len(flux_result) == num_series records = flux_result[0].records assert len(records) == num_records ip.run_line_magic("load_ext", "flux") ip.run_line_magic("flux", "http://localhost:9999 --token my-token --org my-org") request.addfinalizer(cleanup)
class InfluxDB: def __init__(self, config): self._config = config self._client = None self._write_api = None self._query_api = None self._delete_api = None self._tasks_api = None self._organizations_api = None self._token = None self._org = None self._url = None self._bucket = None def start(self) -> bool: """Initialize the InfluxDB client.""" try: influxdb_options = retrieve_options(self._config, 'influxdb2', _INFLUXDB2_OPTIONS) debug_options = retrieve_options(self._config, 'debug', _DEBUG_OPTIONS) except FailedInitialization as e: _LOGGER.error(f"{e}") return False if not influxdb_options.get('enable', None): _LOGGER.warning("InfluxDB support is disabled in the YAML configuration file") return True result = False try: self._bucket = influxdb_options.get('bucket', None) self._url = influxdb_options.get('url', None) self._token = influxdb_options.get('token', None) self._org = influxdb_options.get('org', None) self._client = InfluxDBClient(url=self._url, token=self._token, org=self._org, enable_gzip=True) if not self._client: raise FailedInitialization( f"failed to get InfluxDBClient from '{self._url}' (check url, token, and/or organization)") self._write_api = self._client.write_api(write_options=SYNCHRONOUS) self._query_api = self._client.query_api() self._delete_api = self._client.delete_api() self._tasks_api = self._client.tasks_api() self._organizations_api = self._client.organizations_api() multisma2_debug = os.getenv(_DEBUG_ENV_VAR, 'False').lower() in ('true', '1', 't') try: if multisma2_debug and debug_options.get('delete_bucket', False): self.delete_bucket() _LOGGER.info(f"Deleted bucket '{self._bucket}' at '{self._url}'") except InfluxDBBucketError as e: raise FailedInitialization(f"{e}") try: if not self.connect_bucket(multisma2_debug and debug_options.get('create_bucket', False)): raise FailedInitialization(f"Unable to access (or create) bucket '{self._bucket}' at '{self._url}'") except InfluxDBBucketError as e: raise FailedInitialization(f"{e}") _LOGGER.info(f"Connected to InfluxDB: '{self._url}', bucket '{self._bucket}'") result = True except FailedInitialization as e: _LOGGER.error(f"{e}") except NewConnectionError: _LOGGER.error(f"InfluxDB client unable to connect to host at {self._url}") except ApiException as e: _LOGGER.error(f"InfluxDB client unable to access bucket '{self._bucket}' at {self._url}: {e.reason}") except Exception as e: _LOGGER.error(f"Unexpected exception: {e}") finally: return result def stop(self): if self._write_api: self._write_api.close() self._write_api = None if self._client: self._client.close() self._client = None def bucket(self): return self._bucket def org(self): return self._org def write_api(self): return self._write_api def query_api(self): return self._query_api def delete_api(self): return self._delete_api def tasks_api(self): return self._tasks_api def organizations_api(self): return self._organizations_api def write_points(self, points): if not self._write_api: return False try: self._write_api.write(bucket=self._bucket, record=points, write_precision=WritePrecision.S) return True except ApiException as e: raise InfluxDBWriteError(f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}") except Exception as e: raise InfluxDBWriteError(f"Unexpected failure in write_points(): {e}") def write_history(self, site, topic): if not self._write_api: return False lookup = LP_LOOKUP.get(topic, None) if not lookup: _LOGGER.error(f"write_history(): unknown topic '{topic}'") return False measurement = lookup.get('measurement') tags = lookup.get('tags', None) field = lookup.get('field', None) lps = [] for inverter in site: inverter_name = inverter.pop(0) name = inverter_name.get('inverter', 'sunnyboy') for history in inverter: t = history['t'] v = history['v'] if v is None: continue lp = f"{measurement}" if tags and len(tags): lp += f",{tags[0]}={name}" if isinstance(v, int): lp += f" {field}={v}i {t}" lps.append(lp) else: _LOGGER.error( f"write_history(): unanticipated type '{type(v)}' in measurement '{measurement}/{field}'") continue try: self._write_api.write(bucket=self._bucket, record=lps, write_precision=WritePrecision.S) _LOGGER.debug(f"write_history({site}, {topic}): {lps}") return True except ApiException as e: raise InfluxDBWriteError(f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}") except Exception as e: raise InfluxDBWriteError(f"Unexpected failure in write_history(): {e}") def write_sma_sensors(self, sensor, timestamp=None): if not self._client: return False ts = timestamp if timestamp is not None else int(time.time()) lps = [] for old_point in sensor: point = old_point.copy() topic = point.pop('topic', None) point.pop('precision', None) if topic: lookup = LP_LOOKUP.get(topic, None) if not lookup: _LOGGER.error(f"write_sma_sensors(): unknown topic '{topic}'") continue if not lookup.get('output', False): continue if topic == 'production/today': day = datetime.datetime.fromtimestamp(ts).date() dt = datetime.datetime.combine(day, datetime.time(0, 0)) ts = int(dt.timestamp()) elif topic == 'production/month': month = datetime.date.fromtimestamp(ts).replace(day=1) dt = datetime.datetime.combine(month, datetime.time(0, 0)) ts = int(dt.timestamp()) elif topic == 'production/year': year = datetime.date.fromtimestamp(ts).replace(month=1, day=1) dt = datetime.datetime.combine(year, datetime.time(0, 0)) ts = int(dt.timestamp()) measurement = lookup.get('measurement') tags = lookup.get('tags', None) for k, v in point.items(): field = lookup.get('field') # sample: dc_measurements lp = f'{measurement}' if tags and len(tags): # sample: dc_measurements,_inverter=sb71 lp += f',{tags[0]}={k}' if not field: field = k if isinstance(v, int): # sample: ac_measurements,_inverter=sb71 power=0.23 1556813561098 lp += f' {field}={v}i {ts}' lps.append(lp) elif isinstance(v, float): # sample: ac_measurements,_inverter=sb71 power=0.23 1556813561098 lp += f' {field}={v} {ts}' lps.append(lp) elif isinstance(v, dict): lp_prefix = f'{lp}' for k1, v1 in v.items(): # sample: dc_measurements,_inverter=sb71 lp = f'{lp_prefix}' if tags and len(tags) > 1: # sample: dc_measurements,_inverter=sb71,_string=a lp += f',{tags[1]}={k1}' if isinstance(v1, int): # sample: dc_measurements,_inverter=sb71,_string=a power=1000 1556813561098 lp += f' {field}={v1}i {ts}' lps.append(lp) elif isinstance(v1, float): # sample: dc_measurements,_inverter=sb71,_string=a current=0.23 1556813561098 lp += f' {field}={v1} {ts}' lps.append(lp) else: _LOGGER.error( f"write_sma_sensors(): unanticipated dictionary type '{type(v1)}' in measurement '{measurement}/{field}'") else: _LOGGER.error( f"write_sma_sensors(): unanticipated type '{type(v)}' in measurement '{measurement}/{field}'") continue try: self._write_api.write(bucket=self._bucket, record=lps, write_precision=WritePrecision.S) return True except ApiException as e: raise InfluxDBWriteError(f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}") except Exception as e: raise InfluxDBWriteError(f"Unexpected failure in write_sma_sensors(): {e}") def delete_bucket(self): if not self._client: return False try: buckets_api = self._client.buckets_api() found_bucket = buckets_api.find_bucket_by_name(self._bucket) if found_bucket: buckets_api.delete_bucket(found_bucket) bucket = buckets_api.find_bucket_by_name(self._bucket) if not bucket: return True return False except ApiException as e: raise InfluxDBBucketError( f"InfluxDB client unable to delete bucket '{self._bucket}' at {self._url}: {e.reason}") except Exception as e: raise InfluxDBBucketError(f"Unexpected exception in delete_bucket(): {e}") def connect_bucket(self, create_bucket=False): if not self._client: return False try: buckets_api = self._client.buckets_api() bucket = buckets_api.find_bucket_by_name(self._bucket) if bucket: return True if create_bucket: bucket = buckets_api.create_bucket( bucket_name=self._bucket, org_id=self._org, retention_rules=None, org=None) if bucket: _LOGGER.info(f"Created bucket '{self._bucket}' at {self._url}") return True return False except ApiException as e: raise InfluxDBBucketError( f"InfluxDB client unable to create bucket '{self._bucket}' at {self._url}: {e.reason}") except NewConnectionError: raise except Exception as e: raise InfluxDBBucketError(f"Unexpected exception in connect_bucket(): {e}")
class TasksApiTest(BaseTest): def setUp(self) -> None: super(TasksApiTest, self).setUp() self.organization = self.find_my_org() self.authorization = self.add_tasks_authorization(self.organization) self.client.close() self.client = InfluxDBClient(self.host, self.authorization.token, debug=self.conf.debug) self.tasks_api = self.client.tasks_api() tasks = self.tasks_api.find_tasks() for task in tasks: if task.name.endswith("-IT"): self.tasks_api.delete_task(task.id) def add_tasks_authorization(self, organization): resource = PermissionResource(org=organization.name, type="tasks") create_task = Permission(resource=resource, action="read") delete_task = Permission(resource=resource, action="write") org_resource = PermissionResource(type="orgs") create_org = Permission(resource=org_resource, action="write") read_org = Permission(resource=org_resource, action="read") user_resource = PermissionResource(type="users") create_users = Permission(resource=user_resource, action="write") label_resource = PermissionResource(type="labels") create_labels = Permission(resource=label_resource, action="write") auth_resource = PermissionResource(type="authorizations") create_auth = Permission(resource=auth_resource, action="write") bucket = self.client.buckets_api().find_bucket_by_name("my-bucket") bucket_resource = PermissionResource(org_id=organization.id, id=bucket.id, type="buckets") read_bucket = Permission(resource=bucket_resource, action="read") write_bucket = Permission(resource=bucket_resource, action="write") return self.client.authorizations_api().create_authorization( org_id=organization.id, permissions=[ create_task, delete_task, create_org, read_org, create_users, create_labels, create_auth, read_bucket, write_bucket ]) def test_create_task(self): task_name = self.generate_name("it_task") flux = \ '''option task = {{ name: "{task_name}", every: 1h }} {flux} '''.format(task_name=task_name, flux=TASK_FLUX) task = Task(id=0, name=task_name, org_id=self.organization.id, flux=flux, status="active", description="Task Description") task = self.tasks_api.create_task(task) print(task) self.assertIsNotNone(task) self.assertGreater(len(task.id), 1) self.assertEqual(task.name, task_name) self.assertEqual(task.org_id, self.organization.id) self.assertEqual(task.status, "active") self.assertEqual(task.every, "1h") self.assertEqual(task.cron, None) self.assertEqualIgnoringWhitespace(task.flux, flux) self.assertEqual(task.description, "Task Description") def test_create_task_with_offset(self): task_name = self.generate_name("it_task") flux = \ '''option task = {{ name: "{task_name}", every: 1h, offset: 30m }} {flux} '''.format(task_name=task_name, flux=TASK_FLUX) task = Task(id=0, name=task_name, org_id=self.organization.id, flux=flux, status="active", description="Task Description") task = self.tasks_api.create_task(task) print(task) self.assertIsNotNone(task) self.assertEqual(task.offset, "30m") def test_create_task_every(self): task_name = self.generate_name("it_task") task = self.tasks_api.create_task_every(task_name, TASK_FLUX, "1h", self.organization) print(task) self.assertIsNotNone(task) self.assertGreater(len(task.id), 1) self.assertEqual(task.name, task_name) self.assertEqual(task.org_id, self.organization.id) self.assertEqual(task.status, "active") self.assertEqual(task.every, "1h") self.assertEqual(task.cron, None) self.assertTrue(task.flux.endswith(TASK_FLUX)) def test_create_task_cron(self): task_name = self.generate_name("it task") task = self.tasks_api.create_task_cron(task_name, TASK_FLUX, "0 2 * * *", self.organization.id) self.assertIsNotNone(task) self.assertGreater(len(task.id), 1) self.assertEqual(task.name, task_name) self.assertEqual(task.org_id, self.organization.id) self.assertEqual(task.status, "active") self.assertEqual(task.every, None) self.assertEqual(task.cron, "0 2 * * *") # self.assertEqualIgnoringWhitespace(task.flux, flux) self.assertTrue(task.flux.endswith(TASK_FLUX)) # self.assertEqual(task.links, "active") links = task.links self.assertIsNotNone(task.links) self.assertEqual(links.logs, "/api/v2/tasks/" + task.id + "/logs") self.assertEqual(links.members, "/api/v2/tasks/" + task.id + "/members") self.assertEqual(links.owners, "/api/v2/tasks/" + task.id + "/owners") self.assertEqual(links.runs, "/api/v2/tasks/" + task.id + "/runs") self.assertEqual(links._self, "/api/v2/tasks/" + task.id) # TODO missing get labels self.assertEqual(links.labels, "/api/v2/tasks/" + task.id + "/labels") def test_find_task_by_id(self): task_name = self.generate_name("it task") task = self.tasks_api.create_task_cron(task_name, TASK_FLUX, "0 2 * * *", self.organization.id) task_by_id = self.tasks_api.find_task_by_id(task.id) self.assertEqual(task, task_by_id) @pytest.mark.skip( reason="https://github.com/influxdata/influxdb/issues/13576") @pytest.mark.skip( reason="https://github.com/influxdata/influxdb/issues/11590") def test_find_task_by_user_id(self): task_user = self.users_api.create_user(self.generate_name("TaskUser")) self.tasks_api.create_task_cron(self.generate_name("it_task"), TASK_FLUX, "0 2 * * *", self.organization.id) tasks = self.tasks_api.find_tasks_by_user(task_user_id=task_user.id) print(tasks) self.assertEquals(len(tasks), 1) def test_delete_task(self): task = self.tasks_api.create_task_cron(self.generate_name("it_task"), TASK_FLUX, "0 2 * * *", self.organization.id) self.assertIsNotNone(task) self.tasks_api.delete_task(task.id) with pytest.raises(ApiException) as e: assert self.tasks_api.find_task_by_id(task_id=task.id) assert "failed to find task" in e.value.body def test_update_task(self): task_name = self.generate_name("it task") cron_task = self.tasks_api.create_task_cron(task_name, TASK_FLUX, "0 2 * * *", self.organization.id) flux = ''' option task = {{ name: "{task_name}", every: 3m }} {flux} '''.format(task_name=task_name, flux=TASK_FLUX) cron_task.cron = None cron_task.every = "3m" cron_task.status = "inactive" cron_task.description = "Updated description" updated_task = self.tasks_api.update_task(cron_task) time.sleep(1) self.assertIsNotNone(updated_task) self.assertGreater(len(updated_task.id), 1) self.assertEqual(updated_task.name, task_name) self.assertEqual(updated_task.org_id, cron_task.org_id) self.assertEqual(updated_task.status, "inactive") self.assertEqual(updated_task.every, "3m") self.assertEqual(updated_task.cron, None) self.assertIsNotNone(updated_task.updated_at) now = datetime.datetime.now() now.astimezone() self.assertLess(updated_task.updated_at, now.astimezone(tz=datetime.timezone.utc)) self.assertEqualIgnoringWhitespace(updated_task.flux, flux) self.assertEqual(updated_task.description, "Updated description") def test_member(self): task = self.tasks_api.create_task_cron(self.generate_name("it_task"), TASK_FLUX, "0 2 * * *", self.organization.id) members = self.tasks_api.get_members(task_id=task.id) self.assertEqual(len(members), 0) user = self.users_api.create_user(self.generate_name("Luke Health")) resource_member = self.tasks_api.add_member(member_id=user.id, task_id=task.id) self.assertIsNotNone(resource_member) self.assertEqual(resource_member.id, user.id) self.assertEqual(resource_member.name, user.name) self.assertEqual(resource_member.role, "member") members = self.tasks_api.get_members(task_id=task.id) resource_member = members[0] self.assertEqual(len(members), 1) self.assertEqual(resource_member.id, user.id) self.assertEqual(resource_member.name, user.name) self.assertEqual(resource_member.role, "member") self.tasks_api.delete_member(member_id=user.id, task_id=task.id) members = self.tasks_api.get_members(task_id=task.id) self.assertEqual(len(members), 0) def test_owner(self): task = self.tasks_api.create_task_cron(self.generate_name("it_task"), TASK_FLUX, "0 2 * * *", self.organization.id) owners = self.tasks_api.get_owners(task_id=task.id) self.assertEqual(len(owners), 1) user = self.users_api.create_user(self.generate_name("Luke Health")) resource_member = self.tasks_api.add_owner(owner_id=user.id, task_id=task.id) self.assertIsNotNone(resource_member) self.assertEqual(resource_member.id, user.id) self.assertEqual(resource_member.name, user.name) self.assertEqual(resource_member.role, "owner") owners = self.tasks_api.get_owners(task_id=task.id) self.assertEqual(len(owners), 2) resource_member = owners[1] self.assertEqual(resource_member.id, user.id) self.assertEqual(resource_member.name, user.name) self.assertEqual(resource_member.role, "owner") self.tasks_api.delete_owner(owner_id=user.id, task_id=task.id) owners = self.tasks_api.get_owners(task_id=task.id) self.assertEqual(len(owners), 1) def test_runs(self): task_name = self.generate_name("it task") task = self.tasks_api.create_task_every(task_name, TASK_FLUX, "1s", self.organization) time.sleep(5) runs = self.tasks_api.get_runs(task_id=task.id, limit=10) self.assertGreater(len(runs), 2) success_runs = list(filter(lambda x: x.status == "success", runs)) run = success_runs[0] self.assertIsNotNone(run.id) self.assertEqual(run.task_id, task.id) self.assertEqual(run.status, "success") now = datetime.datetime.now() self.assertLess(run.scheduled_for, now.astimezone(tz=datetime.timezone.utc)) self.assertLess(run.started_at, now.astimezone(tz=datetime.timezone.utc)) self.assertLess(run.finished_at, now.astimezone(tz=datetime.timezone.utc)) self.assertIsNone(run.requested_at) self.assertIsNotNone(run.links) self.assertEqual( run.links.logs, "/api/v2/tasks/" + task.id + "/runs/" + run.id + "/logs") self.assertEqual( run.links.retry, "/api/v2/tasks/" + task.id + "/runs/" + run.id + "/retry") self.assertEqual(run.links._self, "/api/v2/tasks/" + task.id + "/runs/" + run.id) self.assertEqual(run.links.task, "/api/v2/tasks/" + task.id) def test_runs_not_exist(self): with pytest.raises(ApiException) as e: assert self.tasks_api.get_runs("020f755c3c082000") assert "task not found" in e.value.body def test_run_task_manually(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "1s", self.organization) run = self.tasks_api.run_manually(task_id=task.id) print(run) self.assertIsNotNone(run) self.assertTrue(run.status, "scheduled") def test_run_task_manually_not_exist(self): with pytest.raises(ApiException) as e: assert self.tasks_api.run_manually(task_id="020f755c3c082000") assert "failed to force run" in e.value.body def test_retry_run(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "1s", self.organization) time.sleep(5) runs = self.tasks_api.get_runs(task.id) self.assertGreater(len(runs), 1) run = self.tasks_api.retry_run(task_id=runs[0].task_id, run_id=runs[0].id) self.assertIsNotNone(run) self.assertEqual(run.task_id, runs[0].task_id) self.assertEqual(run.status, "scheduled") self.assertEqual(run.task_id, task.id) def test_retry_run_not_exists(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "5s", self.organization) with pytest.raises(ApiException) as e: assert self.tasks_api.retry_run(task_id=task.id, run_id="020f755c3c082000") assert "failed to retry run" in e.value.body def test_logs(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "3s", self.organization) time.sleep(6) logs = self.tasks_api.get_logs(task_id=task.id) for log in logs: self.assertIsNotNone(log.time) self.assertIsNotNone(log.message) print(log) self.tasks_api.delete_task(task_id=task.id) def test_logs_not_exist(self): with pytest.raises(ApiException) as e: assert self.tasks_api.get_logs(task_id="020f755c3c082000") assert "failed to find task logs" in e.value.body def test_run_logs(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "1s", self.organization) time.sleep(5) runs = self.tasks_api.get_runs(task_id=task.id) self.assertGreater(len(runs), 0) logs = self.tasks_api.get_run_logs(run_id=runs[0].id, task_id=task.id) self.assertGreater(len(logs), 0) success = False for log in logs: print(log) if log.message.endswith("Completed successfully"): success = True self.assertTrue(success, "Completed successfully not found in log") def test_runs_not_exists(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "1s", self.organization) with pytest.raises(ApiException) as e: assert self.tasks_api.get_run_logs(task_id=task.id, run_id="020f755c3c082000") assert "failed to find task logs" in e.value.body def test_cancel_run_not_exist(self): task = self.tasks_api.create_task_every(self.generate_name("it task"), TASK_FLUX, "1s", self.organization) time.sleep(5) runs = self.tasks_api.get_runs(task.id) with pytest.raises(ApiException) as e: assert self.tasks_api.cancel_run(task_id=task.id, run_id=runs[0].id) assert "failed to cancel run" in e.value.body assert "run not found" in e.value.body def test_cancel_task_not_exist(self): with pytest.raises(ApiException) as e: assert self.tasks_api.cancel_run("020f755c3c082000", "020f755c3c082000") assert "failed to cancel run" in e.value.body assert "task not found" in e.value.body
class InfluxBenchmark(Benchmark): def __init__(self, loopCount): super().__init__('influx', loopCount) # self._dbType = 'influx' # self._client = InfluxDBClient(host='localhost', port=8086) # def _writeBatch(self, tblName): # global value, tick # data = [] # for i in range(0, 100): # line = "{},device={} value={} {}".format( # TIME_SERIES_NAME, # tblName, # value, # tick*1000000000) # # print(line) # data.append(line) # value += 1 # tick +=1 # self._client.write(data, {'db':DB_NAME}, protocol='line') def executeWrite(self): global tick # influx tick #TODO refactor lineTemplate = TIME_SERIES_NAME + ",rack={},shelf={},barcode='barcode_{}' temperature={},pressure={} {}" batchCount = self._loopCount // BATCH_SIZE for i in range(batchCount): lineBatch = [] for j in range(BATCH_SIZE): n = i * BATCH_SIZE + j # serial number # values first # rtName = 'rt_' + str(n) # table name contains serial number, has info temperature = 20 + (n % 10) pressure = 70 + (n % 10) # tags shelf = (n // ITEMS_PER_SHELF) % MAX_SHELF # shelf number rack = n // (ITEMS_PER_SHELF * MAX_SHELF) # rack number barcode = rack + shelf # now the SQL line = lineTemplate.format( rack, shelf, barcode, # tags temperature, pressure, # values tick * 1000000000) tick += 1 lineBatch.append(line) write_api = self._client.write_api(write_options=SYNCHRONOUS) write_api.write(INFLUX_BUCKET, INFLUX_ORG, lineBatch) # self._client.write(lineBatch, {'db':DB_NAME}, protocol='line') # def doIterate(self): # tblName = Config.getConfig().target_table_name # print("Benchmarking INFLUX database (1 pass) for: {}".format(tblName)) # for i in range(0, Config.getConfig().loop_count): # self._writeBatch(tblName) def _getOrgIdByName(self, orgName): """Find org by name. """ orgApi = self._client.organizations_api() orgs = orgApi.find_organizations() for org in orgs: if org.name == orgName: return org.id raise PerfGenError("Org not found with name: {}".format(orgName)) def _fetchAuth(self): authApi = self._client.authorizations_api() auths = authApi.find_authorizations() for auth in auths: if auth.token == INFLUX_TOKEN: return auth raise PerfGenError("No proper auth found") def _verifyPermissions(self, perms: list): if list: return #OK raise PerfGenError("No permission found") def prepare(self): self._client = InfluxDBClient(url="http://127.0.0.1:8086", token=INFLUX_TOKEN, org=INFLUX_ORG) auth = self._fetchAuth() self._verifyPermissions(auth.permissions) bktApi = self._client.buckets_api() # Delete bkt = bktApi.find_bucket_by_name(INFLUX_BUCKET) if bkt: bktApi.delete_bucket(bkt) # Recreate orgId = self._getOrgIdByName(INFLUX_ORG) bktApi.create_bucket(bucket=None, bucket_name=INFLUX_BUCKET, org_id=orgId)
class Wiretap: def __init__(self, collectors=None): if not collectors: collectors = ALL_COLLECTORS self.collectors = collectors self._startup_check() self.hashes = get_hashes() self.config = read_config() self.inventory = read_inventory() self.metrics = [] self.diffs = {} self.client = InfluxDBClient(url=settings.INFLUX_HOST, token=settings.INFLUX_TOKEN) self.write_api = self.client.write_api(write_options=SYNCHRONOUS) self.query_api = self.client.query_api() self.bucket_api = self.client.buckets_api() self._db_check() self.diff_lock = threading.Lock() self.metric_lock = threading.Lock() self.threads = list() self.RUNMODE = MODE.NORMAL welcome = "\nWiretap:\n\nYour inventory:\n" for item in self.inventory: welcome += f" {item}\n" log.error(welcome) self.start_server_threads() def start_server_threads(self): for server in self.inventory: server_thread = threading.Thread( target=remote_execution, args=(server, self), name=f"thread_{server.name}", daemon=True, ) self.threads.append(server_thread) server_thread.start() def aggregate_diff(self, key, value): last = self.diffs.get(key) with self.diff_lock: self.diffs[key] = value # Trust no one if last: return value - last def add_metric(self, server, metric): log.debug(f"add_metric {metric}") measurement = field_name = metric.tag if name_tuple := metric.tag.split("_", 1): if len(name_tuple) == 2: measurement, field_name = name_tuple if measurement in ["network"]: metric.value = self.aggregate_diff( server.name + measurement + field_name, metric.value) if not metric.value: return None metric.name = server.name point = (Point(measurement).tag("name", metric.name).tag( "agg_type", metric.agg_type).field(field_name, metric.value).time( datetime.utcfromtimestamp(metric.time), WritePrecision.S)) if self.RUNMODE is MODE.NORMAL: if self.add_hash(hash(str(point.__dict__))): with self.metric_lock: self.add_point_to_db(point) self.metrics.append(metric.json()) else: print(metric.json())
class InfluxDB: def __init__(self, config): self._config = config self._client = None self._write_api = None self._query_api = None self._delete_api = None self._tasks_api = None self._organizations_api = None self._token = None self._org = None self._url = None self._bucket = None def start(self): """Initialize the InfluxDB client.""" try: influxdb_options = retrieve_options(self._config, 'influxdb2', _INFLUXDB2_OPTIONS) debug_options = retrieve_options(self._config, 'debug', _DEBUG_OPTIONS) except FailedInitialization as e: _LOGGER.error(f"{e}") return False if len(influxdb_options.keys()) == 0: raise FailedInitialization("missing 'influxdb2' options") result = False try: self._bucket = influxdb_options.get('bucket', None) self._url = influxdb_options.get('url', None) self._token = influxdb_options.get('token', None) self._org = influxdb_options.get('org', None) self._client = InfluxDBClient(url=self._url, token=self._token, org=self._org, enable_gzip=True) if not self._client: raise FailedInitialization( f"failed to get InfluxDBClient from '{self._url}' (check url, token, and/or organization)" ) self._write_api = self._client.write_api(write_options=SYNCHRONOUS) self._query_api = self._client.query_api() self._delete_api = self._client.delete_api() self._tasks_api = self._client.tasks_api() self._organizations_api = self._client.organizations_api() cs_esphome_debug = os.getenv(_DEBUG_ENV_VAR, 'False').lower() in ('true', '1', 't') try: if cs_esphome_debug and debug_options.get( 'delete_bucket', False): self.delete_bucket() _LOGGER.info( f"Deleted bucket '{self._bucket}' at '{self._url}'") except InfluxDBBucketError as e: raise FailedInitialization(f"{e}") try: if not self.connect_bucket( cs_esphome_debug and debug_options.get('create_bucket', False)): raise FailedInitialization( f"Unable to access (or create) bucket '{self._bucket}' at '{self._url}'" ) except InfluxDBBucketError as e: raise FailedInitialization(f"{e}") _LOGGER.info( f"Connected to InfluxDB: '{self._url}', bucket '{self._bucket}'" ) result = True except FailedInitialization as e: _LOGGER.error(f" client {e}") self._client = None except NewConnectionError: _LOGGER.error( f"InfluxDB client unable to connect to host at {self._url}") except ApiException as e: _LOGGER.error( f"InfluxDB client unable to access bucket '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: _LOGGER.error(f"Unexpected exception: {e}") finally: return result def stop(self): if self._write_api: self._write_api.close() self._write_api = None if self._client: self._client.close() self._client = None def bucket(self): return self._bucket def org(self): return self._org def write_api(self): return self._write_api def query_api(self): return self._query_api def delete_api(self): return self._delete_api def tasks_api(self): return self._tasks_api def organizations_api(self): return self._organizations_api def write_point(self, measurement, tags, field, value, timestamp=None): """Write a single sensor to the database.""" timestamp = timestamp if timestamp is not None else int(time.time()) lp_tags = '' separator = '' for tag in tags: lp_tags += f"{separator}{tag.get('t')}={tag.get('v')}" separator = ',' lp = f"{measurement}," + lp_tags + f" {field}={value} {timestamp}" try: self._write_api.write(bucket=self._bucket, record=lp, write_precision=WritePrecision.S) except ApiException as e: raise InfluxDBWriteError( f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: raise InfluxDBWriteError( f"Unexpected failure in write_point(): {e}") def write_points(self, points): """Write a list of points to the database.""" try: self._write_api.write(bucket=self._bucket, record=points, write_precision=WritePrecision.S) except ApiException as e: raise InfluxDBWriteError( f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: raise InfluxDBWriteError( f"Unexpected failure in write_points(): {e}") def write_batch_sensors(self, batch_sensors, timestamp=None): """Write a batch of sensors to the database.""" if len(batch_sensors) == 0: return timestamp = timestamp if timestamp is not None else int(time.time()) batch = [] for record in batch_sensors: sensor = record.get('sensor', None) state = record.get('state', None) measurement = sensor.get('measurement', None) device = sensor.get('device', None) location = sensor.get('location', None) precision = sensor.get('precision', None) if measurement is None or device is None: raise InfluxDBFormatError( "'measurement' and/or 'device' are required") location_tag = '' if not location or not len( location) else f',_location={location}' device_tag = f',_device={device}' value = round( state, precision) if ((precision is not None) and isinstance(state, float)) else state lp = f'{measurement}{device_tag}{location_tag} sample={value} {timestamp}' batch.append(lp) try: self._write_api.write(bucket=self._bucket, record=batch, write_precision=WritePrecision.S) except ApiException as e: raise InfluxDBWriteError( f"InfluxDB client unable to write to '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: raise InfluxDBWriteError( f"Unexpected failure in write_batch_sensors(): {e}") def delete_bucket(self): try: buckets_api = self._client.buckets_api() found_bucket = buckets_api.find_bucket_by_name(self._bucket) if found_bucket: buckets_api.delete_bucket(found_bucket) bucket = buckets_api.find_bucket_by_name(self._bucket) if not bucket: return True return False except ApiException as e: raise InfluxDBBucketError( f"InfluxDB client unable to delete bucket '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: raise InfluxDBBucketError( f"Unexpected exception in delete_bucket(): {e}") def connect_bucket(self, create_bucket=False): try: buckets_api = self._client.buckets_api() bucket = buckets_api.find_bucket_by_name(self._bucket) if bucket: return True if create_bucket: bucket = buckets_api.create_bucket(bucket_name=self._bucket, org_id=self._org, retention_rules=None, org=None) if bucket: _LOGGER.info( f"Created bucket '{self._bucket}' at {self._url}") return True return False except ApiException as e: raise InfluxDBBucketError( f"InfluxDB client unable to create bucket '{self._bucket}' at {self._url}: {e.reason}" ) except Exception as e: raise InfluxDBBucketError( f"Unexpected exception in connect_bucket(): {e}")
class InfluxDBClientTestIT(BaseTest): httpRequest = [] def tearDown(self) -> None: super(InfluxDBClientTestIT, self).tearDown() if hasattr(self, 'httpd'): self.httpd.shutdown() if hasattr(self, 'httpd_thread'): self.httpd_thread.join() InfluxDBClientTestIT.httpRequest = [] def test_proxy(self): self._start_proxy_server() self.client.close() self.client = InfluxDBClient( url=self.host, token=self.auth_token, proxy=f"http://localhost:{self.httpd.server_address[1]}", proxy_headers={'ProxyHeader': 'Val'}) ready = self.client.ready() self.assertEqual(ready.status, "ready") self.assertEqual(1, len(InfluxDBClientTestIT.httpRequest)) self.assertEqual( 'Val', InfluxDBClientTestIT.httpRequest[0].headers.get('ProxyHeader')) def test_ping(self): ping = self.client.ping() self.assertTrue(ping) def test_ping_not_running_instance(self): client_not_running = InfluxDBClient("http://localhost:8099", token="my-token", debug=True) ping = client_not_running.ping() self.assertFalse(ping) client_not_running.close() def test_version(self): version = self.client.version() self.assertTrue(len(version) > 0) def test_version_not_running_instance(self): client_not_running = InfluxDBClient("http://localhost:8099", token="my-token", debug=True) with self.assertRaises(NewConnectionError): client_not_running.version() client_not_running.close() def test_username_password_authorization(self): self.client.close() self.client = InfluxDBClient(url=self.host, username="******", password="******", debug=True) self.client.query_api().query("buckets()", "my-org") def test_query_and_debug(self): self.client.close() self.client = InfluxDBClient(url=self.host, token="my-token", debug=True) # Query results = self.client.query_api().query("buckets()", "my-org") self.assertIn( "my-bucket", list(map(lambda record: record["name"], results[0].records))) # Query RAW results = self.client.query_api().query_raw("buckets()", "my-org") self.assertIn("my-bucket", codecs.decode(results.data)) # Bucket API results = self.client.buckets_api().find_buckets() self.assertIn("my-bucket", list(map(lambda bucket: bucket.name, results.buckets))) def _start_proxy_server(self): import http.server import urllib.request class ProxyHTTPRequestHandler(http.server.SimpleHTTPRequestHandler): def do_GET(self): InfluxDBClientTestIT.httpRequest.append(self) self.send_response(200) self.send_header('Content-type', 'application/json') self.end_headers() self.copyfile(urllib.request.urlopen(self.path), self.wfile) self.httpd = http.server.HTTPServer(('localhost', 0), ProxyHTTPRequestHandler) self.httpd_thread = threading.Thread(target=self.httpd.serve_forever) self.httpd_thread.start()
class InfluxDBWriter(Writer): """Write to the specified file. If filename is empty, write to stdout.""" def __init__(self, bucket_name): """ Write data records to the InfluxDB. ``` bucket_name the name of the bucket in InfluxDB. If the bucket does not exists then this writer will try to create it. ``` """ super().__init__(input_format=Text) if not INFLUXDB_SETTINGS_FOUND: raise RuntimeError('File database/settings.py not found. ' 'InfluxDB functionality is not available. Have ' 'you copied over database/settings.py.dist ' 'to database/settings.py and followed the ' 'configuration instructions in it?') if not INFLUXDB_CLIENT_FOUND: raise RuntimeError('Python module influxdb_client not found. Please ' 'install using "pip install influxdb_client" prior ' 'to using InfluxDBWriter.') self.client = InfluxDBClient(url=INFLUXDB_URL, token=INFLUXDB_AUTH_TOKEN, org=INFLUXDB_ORG) # get the orgID from the name: try: self.organizations_api = self.client.organizations_api() orgs = self.organizations_api.find_organizations() except: raise RuntimeError('Error connecting to the InfluxDB API. ' 'Please confirm that InfluxDB is running and ' 'that the authentication token is correct.') our_org = next((org for org in orgs if org.name == INFLUXDB_ORG), None) if not our_org: raise RuntimeError('Can not find the organization "' + INFLUXDB_ORG + '" in InfluxDB') self.org_id = our_org.id # get the bucketID from the name: self.bucket_api = self.client.buckets_api() bucket = self.bucket_api.find_bucket_by_name(bucket_name) # if the bucket does not exist then try to create it if not bucket: try: new_bucket = self.bucket_api.create_bucket(bucket_name=bucket_name, org_id=self.org_id) logging.info('Creating new bucket for: %s', bucket_name) self.bucket_id = new_bucket.id except: raise RuntimeError('Can not create bucket in InfluxDB for ' + bucket_name) else: self.bucket_id = bucket.id self.write_api = self.client.write_api(write_options=ASYNCHRONOUS) ############################ def write(self, record): """ Note: Assume record is a dict or list of dict. Each dict contains a list of "fields" and float "timestamp" (UTC epoch seconds) """ if record is None: return logging.info('InfluxDBWriter writing record: %s', record) if type(record) is not dict and type(record) is not list: logging.warning('InfluxDBWriter could not ingest record ' 'type %s: %s', type(record), str(record)) try: if type(record) is list: influxDB_record = map(lambda single_record: {"measurement": single_record['data_id'], "tags": {"sensor": single_record['data_id'] }, "fields": single_record['fields'], "time": int(single_record['timestamp']*1000000000) }, record) else: influxDB_record = {"measurement": record['data_id'], "tags": {"sensor": record['data_id'] }, "fields": record['fields'], "time": int(record['timestamp']*1000000000) } self.write_api.write(self.bucket_id, self.org_id, influxDB_record) return except: logging.warning('InfluxDBWriter could not ingest record ' 'type %s: %s', type(record), str(record))
class GzipSupportTest(BaseTest): def setUp(self) -> None: super(GzipSupportTest, self).setUp() httpretty.enable() httpretty.reset() def tearDown(self) -> None: self.client.close() httpretty.disable() def test_gzip_disabled(self): query_response = \ "#datatype,string,long,dateTime:RFC3339,dateTime:RFC3339,dateTime:RFC3339,long,string,string,string\n" \ "#group,false,false,false,false,false,false,false,false,false,true\n#default,_result,,,,,,,,\n" \ ",result,table,_start,_stop,_time,_value,_field,_measurement,host\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,121,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,122,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,123,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,124,free,mem,A\n" httpretty.register_uri( httpretty.GET, uri="http://localhost/api/v2/me", status=200, body="{\"name\":\"Tom\"}", adding_headers={'Content-Type': 'application/json'}) httpretty.register_uri(httpretty.POST, uri="http://localhost/api/v2/write", status=204) httpretty.register_uri(httpretty.POST, uri="http://localhost/api/v2/query", status=200, body=query_response) self.client = InfluxDBClient("http://localhost", "my-token", org="my-org", enable_gzip=False) _user = self.client.users_api().me() self.assertEqual("Tom", _user._name) _response = self.client.write_api(write_options=SYNCHRONOUS) \ .write("my-bucket", "my-org", "h2o_feet,location=coyote_creek water_level=1 1") self.assertEqual(None, _response) _tables = self.client.query_api() \ .query('from(bucket:"my-bucket") |> range(start: 1970-01-01T00:00:00.000000001Z) |> last()', "my-org") self.assertEqual(1, len(_tables)) self.assertEqual(4, len(_tables[0].records)) self.assertEqual(121, _tables[0].records[0].get_value()) self.assertEqual(122, _tables[0].records[1].get_value()) self.assertEqual(123, _tables[0].records[2].get_value()) self.assertEqual(124, _tables[0].records[3].get_value()) _requests = httpretty.httpretty.latest_requests self.assertEqual(3, len(_requests)) # Unsupported self.assertEqual("/api/v2/me", _requests[0].path) self.assertEqual(None, _requests[0].headers['Content-Encoding']) self.assertEqual("identity", _requests[0].headers['Accept-Encoding']) # Write self.assertEqual( "/api/v2/write?org=my-org&bucket=my-bucket&precision=ns", _requests[1].path) self.assertEqual("identity", _requests[1].headers['Content-Encoding']) self.assertEqual("identity", _requests[1].headers['Accept-Encoding']) self.assertEqual("h2o_feet,location=coyote_creek water_level=1 1", _requests[1].parsed_body) # Query self.assertEqual("/api/v2/query?org=my-org", _requests[2].path) self.assertEqual(None, _requests[2].headers['Content-Encoding']) self.assertEqual("identity", _requests[2].headers['Accept-Encoding']) self.assertTrue( 'from(bucket:"my-bucket") |> range(start: 1970-01-01T00:00:00.000000001Z) |> last()' in str(_requests[2].parsed_body)) def test_gzip_enabled(self): query_response = \ "#datatype,string,long,dateTime:RFC3339,dateTime:RFC3339,dateTime:RFC3339,long,string,string,string\n" \ "#group,false,false,false,false,false,false,false,false,false,true\n#default,_result,,,,,,,,\n" \ ",result,table,_start,_stop,_time,_value,_field,_measurement,host\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,121,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,122,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,123,free,mem,A\n" \ ",,0,1970-01-01T00:00:10Z,1970-01-01T00:00:20Z,1970-01-01T00:00:10Z,124,free,mem,A\n" httpretty.register_uri( httpretty.GET, uri="http://localhost/api/v2/me", status=200, body="{\"name\":\"Tom\"}", adding_headers={'Content-Type': 'application/json'}) httpretty.register_uri(httpretty.POST, uri="http://localhost/api/v2/write", status=204) httpretty.register_uri(httpretty.POST, uri="http://localhost/api/v2/query", status=200, body=gzip.compress( bytes(query_response, "utf-8")), adding_headers={'Content-Encoding': 'gzip'}) self.client = InfluxDBClient("http://localhost", "my-token", org="my-org", enable_gzip=True) _user = self.client.users_api().me() self.assertEqual("Tom", _user._name) _response = self.client.write_api(write_options=SYNCHRONOUS) \ .write("my-bucket", "my-org", "h2o_feet,location=coyote_creek water_level=1 1") self.assertEqual(None, _response) _tables = self.client.query_api() \ .query('from(bucket:"my-bucket") |> range(start: 1970-01-01T00:00:00.000000001Z) |> last()', "my-org") self.assertEqual(1, len(_tables)) self.assertEqual(4, len(_tables[0].records)) self.assertEqual(121, _tables[0].records[0].get_value()) self.assertEqual(122, _tables[0].records[1].get_value()) self.assertEqual(123, _tables[0].records[2].get_value()) self.assertEqual(124, _tables[0].records[3].get_value()) _requests = httpretty.httpretty.latest_requests self.assertEqual(3, len(_requests)) # Unsupported self.assertEqual("/api/v2/me", _requests[0].path) self.assertEqual(None, _requests[0].headers['Content-Encoding']) self.assertEqual("identity", _requests[0].headers['Accept-Encoding']) # Write self.assertEqual( "/api/v2/write?org=my-org&bucket=my-bucket&precision=ns", _requests[1].path) self.assertEqual("gzip", _requests[1].headers['Content-Encoding']) self.assertEqual("identity", _requests[1].headers['Accept-Encoding']) self.assertNotEqual("h2o_feet,location=coyote_creek water_level=1 1", _requests[1].parsed_body) # Query self.assertEqual("/api/v2/query?org=my-org", _requests[2].path) self.assertEqual(None, _requests[2].headers['Content-Encoding']) self.assertEqual("gzip", _requests[2].headers['Accept-Encoding']) self.assertTrue( 'from(bucket:"my-bucket") |> range(start: 1970-01-01T00:00:00.000000001Z) |> last()' in str(_requests[2].parsed_body)) def test_write_query_gzip(self): httpretty.disable() self.client = InfluxDBClient(self.host, token="my-token", org="my-org", debug=False, enable_gzip=True) self.api_client = self.client.api_client self.buckets_client = self.client.buckets_api() self.query_client = self.client.query_api() self.org = "my-org" self.my_organization = self.find_my_org() _bucket = self.create_test_bucket() self.client.write_api(write_options=SYNCHRONOUS) \ .write(_bucket.name, self.org, "h2o_feet,location=coyote_creek water_level=111 1") _result = self.query_client.query( f"from(bucket:\"{_bucket.name}\") |> range(start: 1970-01-01T00:00:00.000000001Z)", self.org) self.assertEqual(len(_result), 1) self.assertEqual(_result[0].records[0].get_measurement(), "h2o_feet") self.assertEqual(_result[0].records[0].get_value(), 111.0) self.assertEqual(_result[0].records[0].get_field(), "water_level")
class InfluxAPI(Processor): def __init__(self, url: str, token: str, org: str, data_bucket: str, meta_bucket: str, workers: int = cpu_count()): super().__init__() self.client = InfluxDBClient(url=url, token=token, org=org) self.url = url self.token = token self.org = org if not self.check_bucket_exists(data_bucket): raise KeyError(f"Data bucket {data_bucket} as does not exist") if not self.check_bucket_exists(meta_bucket): raise KeyError(f"Meta bucket {meta_bucket} as does not exist") self.data_bucket = data_bucket self.meta_bucket = meta_bucket # write with batch api with sane looking defaults self.api = self.client.write_api( write_options=WriteOptions(batch_size=200, flush_interval=2000, jitter_interval=100, retry_interval=2000, write_scheduler=ThreadPoolScheduler( max_workers=workers))) def __del__(self): # Cleanup connection self.client.close() def check_bucket_exists(self, bucket_name: str) -> bool: return self.client.buckets_api().find_bucket_by_name( bucket_name=bucket_name) @staticmethod def _generate_data(data: List[dict]) -> List[Point]: points = [] for tags, fields, timestamp in map( itemgetter('tags', 'fields', 'timestamp'), data): collector = tags['collector'] # Every measurement goes against what collected it p = Point.measurement(collector) # which makes the collector tag not needed, remove to reduce dimensionality [ p.tag(key, value) for key, value in tags.items() if key != 'collector' and not isinstance(value, type(None)) ] [ p.field(key, value) for key, value in fields.items() if not isinstance(value, type(None)) ] p.time(datetime.fromtimestamp(timestamp, tz=timezone.utc), write_precision=WritePrecision.S) points.append(p) return points def process(self, data: dict) -> Optional[ProcessedData]: p = self._process(data) p.discard_strings() return p def send_data(self, data: List[dict], **kwargs) -> bool: """Sending specific data. Assumes its in the common format""" if not kwargs.get('bucket', None): raise KeyError("Expected to have bucket name passed in") self.api.write(bucket=kwargs['bucket'], record=self._generate_data(data)) return True def send(self, data: ProcessedData) -> bool: self.api.write(bucket=self.data_bucket, record=self._generate_data(data.data)) self.api.write(bucket=self.data_bucket, record=self._generate_data(data.all_meta_data)) return True
class TimeseriesClient: host: Union[str, None] port: int token: Union[str, None] organization: str def __init__( self, host: str = None, port: int = None, organization: str = "GEWV", token: str = None, client: InfluxDBClient = None, verify_ssl: bool = True, ): if client is None: if host is None: raise Exception( "Missing Host Address for Timeseries DB Client.") if port is None: raise Exception("Missing Port for Timeseries DB Client.") if token is None: raise Exception("Missing Token for Timeseries DB Client.") protocol = "https" if verify_ssl else "http" self._client = InfluxDBClient( url=f"{protocol}://{host}:{port}", token=token, verify_ssl=verify_ssl, ) if len(organization) != 16: # receive id of the org and store the info self._org_api = self._client.organizations_api() self._org_id = self.get_org_id_by_name(org_name=organization) if self._org_id is None: raise Exception( f"The organization {organization} dont exists in InfluxDB. Break execution." ) self._client.org = self._org_id else: self._client.org = organization else: self._client = client self._org_api = self._client.organizations_api() self._write_api = self._client.write_api(write_options=SYNCHRONOUS) self._query_api = self._client.query_api() self._bucket_api = self._client.buckets_api() self._grafana_api = GrafanaApi(host=host, port=3000, use_tls=False) @staticmethod def from_env_properties(): client = InfluxDBClient.from_env_properties() return TimeseriesClient(client=client) def health(self): return self._client.health() def get_org_id_by_name(self, org_name: str) -> Union[str, None]: orgs: List[Organization] = self._org_api.find_organizations() for org in orgs: if org.name == org_name: return org.id return None def create_bucket(self, bucket: str): try: self._bucket_api.create_bucket(bucket_name=bucket) except ApiException as err: if err.status != 422: raise def exist_bucket(self, bucket: str): return self._bucket_api.find_bucket_by_name(bucket_name=bucket) def get_bucket_by_name(self, bucket_name: str): return self._bucket_api.find_bucket_by_name(bucket_name=bucket_name) def delete_bucket(self, bucket: str): bucket_id = self.get_bucket_by_name(bucket_name=bucket) return self._bucket_api.delete_bucket(bucket=bucket_id) def get_grafana_orgs(self) -> List[GrafanaOrganization]: return self._grafana_api.get_organizations() def get_grafana_org(self, org_name: str) -> GrafanaOrganization: return self._grafana_api.get_organization_by_name(org_name=org_name) def create_grafana_org(self, org_name: str): return self._grafana_api.create_organization(org_name=org_name) def delete_grafana_org(self, org_name: str): org = self.get_grafana_org(org_name=org_name) if org is None: raise Exception( f"Cant delete grafana org {org_name}. Org not exist!") return self._grafana_api.delete_organization(org["id"]) def create_project(self, project_name: str): # Steps # 1. create new bucket # 2. create token for bucket # 3. create new org in grafana # 4. create new source in grafana pass def get_points( self, **kwargs, ) -> List[FluxTable]: if not self.health: raise Exception("Influx DB is not reachable or unhealthy.") tables = self._query_api.query(query=self.build_query(**kwargs)) return tables def get_dataframe(self, **kwargs): return self.query_dataframe(flux_query=self.build_query(**kwargs)) def query_dataframe( self, flux_query: str, ): """ with this function you can send a own query to InfluxDB and you will get back a dataframe with datetimeindex """ if not self.health: raise Exception("Influx DB is not reachable or unhealthy.") df = cast( DataFrame, self._query_api.query_data_frame(query=flux_query), ) if "_time" in df.columns: df = df.set_index(pd.to_datetime(df["_time"])) return df def write_points(self, project: str, points: List[Point]): self._write_api.write(bucket=project, record=points) def write_a_dataframe( self, project: str, measurement_name: str, dataframe: pd.DataFrame, tag_columns: List[str] = [], additional_tags: Dict[str, str] = None, ): """ Write a pandas dataframe to the influx db. You can define some tags, that are appended to every entry. """ if additional_tags is None: self._write_api.write( bucket=project, record=dataframe, data_frame_measurement_name=measurement_name, data_frame_tag_columns=tag_columns, ) return tags_dataframe = pd.DataFrame(index=dataframe.index) # create the dataframe with the tags for tag_name, tag in additional_tags.items(): tag_columns.append(tag_name) tags_dataframe[tag_name] = [tag] * len(dataframe) combined_frames = pd.concat([dataframe, tags_dataframe], axis=1) self._write_api.write( bucket=project, record=combined_frames, data_frame_measurement_name=measurement_name, data_frame_tag_columns=tag_columns, ) def build_query( self, project: str, fields: Dict[str, str] = {}, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None, precision: str = "5m", ) -> str: query = f""" from(bucket: "{project}") """ if start_time is not None and end_time is not None: self.test_datetime(start_time) self.test_datetime(end_time) query += f""" |> range(start: {start_time.isoformat()}, stop: {end_time.isoformat()}) """ elif start_time is not None: self.test_datetime(start_time) query += f""" |> range(start: {start_time.isoformat()}) """ elif end_time is not None: self.test_datetime(end_time) query += f""" |> range(stop: {end_time.isoformat()}) """ for f, v in fields.items(): query += f""" |> filter(fn: (r) => r["{f}"] == "{v}") """ query += f""" |> aggregateWindow(every: {precision}, fn: mean, createEmpty: true) |> yield(name: "mean") """ return query @staticmethod def test_datetime(dt: datetime): if not isinstance(dt, datetime): raise Exception( f"The delivered datetime {dt} is not from type datetime.") if dt.tzinfo is None: raise Exception( f"The time {dt.isoformat()} has no timezone info. That is necassary." )
TMP_FILE_PATH = data["tmp_path"] IP_LOOKUP_TABLE_PATH = data["ip_path"] PREV_PULLED_DATA_PATH = data["cache_path"] SYS_LOG_PATH = data["syslog"] OPENVPNLOG_PATH = data["vpn_status"] BUCKET = platform.uname()[1] + "-VPN" IPDB_PATH = data["ipdbpath"] except: print("Issue with Config, please run the command init command") ###influxdb Parameters try: start_time = time.perf_counter() client = InfluxDBClient(url=URL, token=TOKEN, org=ORG) write_api = client.write_api(write_options=SYNCHRONOUS) bucket_api = client.buckets_api() except: print("Issue with Config, please run the command init command") ### Fucntions def main(): if sys.argv[1].casefold() == "init": init_environment() return try: bucket_api.create_bucket(bucket=Bucket( name=BUCKET, org_id=ORG_ID, retention_rules=[BucketRetentionRules(every_seconds=604800)]))