From bece3c9916c76a885ad615449746cc0e76de4e0a Mon Sep 17 00:00:00 2001
From: Tilly Woodfield <22456167+tillywoodfield@users.noreply.github.com>
Date: Tue, 1 Apr 2025 13:25:27 +0300
Subject: [PATCH] lib360dataquality: support local schema files

---
 lib360dataquality/cove/schema.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/lib360dataquality/cove/schema.py b/lib360dataquality/cove/schema.py
index e5606181..ef8d88ff 100644
--- a/lib360dataquality/cove/schema.py
+++ b/lib360dataquality/cove/schema.py
@@ -28,7 +28,7 @@ class Schema360(SchemaJsonMixin):
     _pkg_schema_obj = {}
     _schema_obj = {}
 
-    def __init__(self, data_dir) -> None:
+    def __init__(self, data_dir, local_pkg_schema_path=None, local_grant_schema_path=None) -> None:
         # Create dedicated location for schema work
         self.working_dir = os.path.join(data_dir, "schema")
         try:
@@ -40,11 +40,19 @@ def __init__(self, data_dir) -> None:
         # urljoin does not discard the final part of the location.
         self.schema_host = f"{self.working_dir}/"
 
-        schema_url = urljoin(config["schema_host"], self.schema_name)
-        pkg_schema_url = urljoin(config["schema_host"], self.pkg_schema_name)
-
-        self._pkg_schema_obj = get_request(pkg_schema_url).json()
-        self._schema_obj = get_request(schema_url).json()
+        if local_pkg_schema_path:
+            with open(local_pkg_schema_path) as f:
+                self._pkg_schema_obj = json.load(f)
+        else:
+            pkg_schema_url = urljoin(config["schema_host"], self.pkg_schema_name)
+            self._pkg_schema_obj = get_request(pkg_schema_url).json()
+
+        if local_grant_schema_path:
+            with open(local_grant_schema_path) as f:
+                self._schema_obj = json.load(f)
+        else:
+            schema_url = urljoin(config["schema_host"], self.schema_name)
+            self._schema_obj = get_request(schema_url).json()
 
         # Update the pkg schema to no longer point to an external reference for the
         # grants schema.