make multitenant by storing the azure tenant ID

Author: lilatomic

llamazure/history/__main__.py

@@ -4,11 +4,14 @@
 import os
 from dataclasses import asdict, is_dataclass
 from typing import Dict, cast
+from uuid import UUID
 
 from azure.identity import DefaultAzureCredential
 
 from llamazure.azgraph.azgraph import Graph
 from llamazure.azgraph.models import ResErr
+from llamazure.azrest.azrest import AzRest
+from llamazure.azrest.models import Req as AzReq, AzList
 from llamazure.history.data import DB, TSDB
 from llamazure.rid import mp
 from llamazure.tresource.mp import MPData, TresourceMPData
@@ -22,6 +25,8 @@ def default(self, o):
 			return asdict(o)
 		if isinstance(o, datetime.datetime):
 			return o.isoformat()
+		if isinstance(o, UUID):
+			return str(o)
 		return super().default(o)
 
 
@@ -38,7 +43,13 @@ def reformat_resources_for_tresource(resources):
 	db = DB(tsdb)
 	db.create_tables()
 
-	g = Graph.from_credential(DefaultAzureCredential())
+	credential = DefaultAzureCredential()
+	g = Graph.from_credential(credential)
+	azr = AzRest.from_credential(credential)
+
+	tenants = azr.call(AzReq.get("GetTenants", "/tenants", "2022-12-01", AzList[dict]))
+	tenant_id = UUID(tenants[0]["tenantId"])
+
 	resources = g.q("Resources")
 	if isinstance(resources, ResErr):
 		raise RuntimeError(ResErr)
@@ -48,13 +59,13 @@ def reformat_resources_for_tresource(resources):
 
 	snapshot_time = datetime.datetime.utcnow()
 
-	db.insert_snapshot(snapshot_time, ((cast(str, path), mpdata.data) for path, mpdata in tree.resources.items() if mpdata.data is not None))
+	db.insert_snapshot(snapshot_time, tenant_id, ((cast(str, path), mpdata.data) for path, mpdata in tree.resources.items() if mpdata.data is not None))
 
 	delta_q = g.q("Resources | take(1)")
 	if isinstance(delta_q, ResErr):
 		raise RuntimeError(ResErr)
 	delta = delta_q[0]
-	db.insert_delta(snapshot_time + datetime.timedelta(seconds=1), delta["id"].lower(), delta)
+	db.insert_delta(snapshot_time + datetime.timedelta(seconds=1), tenant_id, delta["id"].lower(), delta)
 
 	print(json.dumps(db.read_latest(), indent=2, cls=MyEncoder))
 	# print(json.dumps(db.read_snapshot(snapshot_time + datetime.timedelta(seconds=2)), indent=2, cls=MyEncoder))

llamazure/history/data.py

@@ -2,6 +2,7 @@
 import datetime
 from textwrap import dedent
 from typing import Any, Iterable, Optional, Tuple
+from uuid import UUID
 
 import psycopg
 from psycopg.types.json import Jsonb
@@ -49,8 +50,9 @@ def create_tables(self):
 			dedent(
 				"""\
 				CREATE TABLE IF NOT EXISTS snapshot (
-					id SERIAL PRIMARY KEY,
-					time TIMESTAMPTZ NOT NULL
+					id 				SERIAL PRIMARY KEY,
+					time			TIMESTAMPTZ NOT NULL,
+					azure_tenant 	UUID
 				)
 				"""
 			)
@@ -61,32 +63,33 @@ def create_tables(self):
 				"""\
 				CREATE TABLE IF NOT EXISTS res (
 					time TIMESTAMPTZ NOT NULL,
-					snapshot 	INTEGER,
-					rid			VARCHAR,
-					data		JSONB,
+					snapshot 		INTEGER,
+					azure_tenant 	UUID,
+					rid				VARCHAR,
+					data			JSONB,
 					FOREIGN KEY (snapshot) REFERENCES snapshot (id)
 				)
 				"""
 			)
 		)
 		self.db.create_hypertable("res", "time")
 
-	def insert_resource(self, time: datetime.datetime, snapshot_id, rid: str, data: dict):
+	def insert_resource(self, time: datetime.datetime, azure_tenant: UUID, snapshot_id, rid: str, data: dict):
 		"""Insert a resource into the DB"""
 		self.db.exec(
-			"""INSERT INTO res (time, snapshot, rid, data) VALUES (%s, %s, %s, %s)""",
-			(time, snapshot_id, rid, Jsonb(data)),
+			"""INSERT INTO res (time, snapshot, azure_tenant, rid, data) VALUES (%s, %s, %s, %s, %s)""",
+			(time, snapshot_id, azure_tenant, rid, Jsonb(data)),
 		)
 
-	def insert_snapshot(self, time: datetime.datetime, resources: Iterable[Tuple[str, dict]]):
+	def insert_snapshot(self, time: datetime.datetime, azure_tenant: UUID, resources: Iterable[Tuple[str, dict]]):
 		"""Insert a complete snapshot into the DB"""
-		snapshot_id = self.db.exec_returning("""INSERT INTO snapshot (time) VALUES (%s) RETURNING id""", (time,))
+		snapshot_id = self.db.exec_returning("""INSERT INTO snapshot (time, azure_tenant) VALUES (%s, %s) RETURNING id""", (time, azure_tenant))
 		for rid, data in resources:
-			self.insert_resource(time, snapshot_id, rid, data)
+			self.insert_resource(time, azure_tenant, snapshot_id, rid, data)
 
-	def insert_delta(self, time: datetime.datetime, rid: str, data: dict):
+	def insert_delta(self, time: datetime.datetime, azure_tenant: UUID, rid: str, data: dict):
 		"""Insert a single delta into the DB"""
-		return self.insert_resource(time, None, rid, data)
+		return self.insert_resource(time, azure_tenant, None, rid, data)
 
 	def read_snapshot(self, time: datetime.datetime):
 		"""Read a complete snapshot. Does not include any deltas"""