Use api key for validation if specified (#281)

Author: keyurva

.gitignore

@@ -127,4 +127,11 @@ datacommons.RCheck
 .vscode/
 
 ## JetBrains
-.idea/
+.idea/
+
+# Gemini
+GEMINI.md
+.gemini/
+
+# Temp files
+tmp/

datacommons_client/endpoints/base.py

@@ -45,7 +45,7 @@ def __init__(
 
     if url is not None:
       # Use the given URL directly (strip trailing slash)
-      self.base_url = check_instance_is_valid(url.rstrip("/"))
+      self.base_url = check_instance_is_valid(url.rstrip("/"), api_key=api_key)
     else:
       # Resolve from dc_instance
       self.base_url = resolve_instance_url(dc_instance)

datacommons_client/tests/endpoints/test_base.py

@@ -38,7 +38,8 @@ def test_api_initialization_with_url(mock_check_instance):
       "Content-Type": "application/json",
       "x-surface": "clientlib-python"
   }
-  mock_check_instance.assert_called_once_with("https://custom_instance.api/v2")
+  mock_check_instance.assert_called_once_with("https://custom_instance.api/v2",
+                                              api_key=None)
 
 
 @patch(

datacommons_client/tests/endpoints/test_request_handling.py

@@ -69,7 +69,27 @@ def test_check_instance_is_valid_valid(mock_get):
   # Assert that the instance URL is returned if it is valid
   assert check_instance_is_valid(instance_url) == instance_url
   mock_get.assert_called_once_with(
-      f"{instance_url}/node?nodes=country%2FGTM&property=->name")
+      f"{instance_url}/node?nodes=country%2FGTM&property=->name", headers={})
+
+
+@patch("requests.get")
+def test_check_instance_is_valid_with_key_valid(mock_get):
+  """Tests that a valid instance URL is correctly validated."""
+
+  # Create a mock response object with the expected JSON data and status code
+  mock_response = MagicMock()
+  mock_response.json.return_value = {"data": {"country/GTM": {}}}
+  mock_response.status_code = 200
+  mock_get.return_value = mock_response
+
+  # Mock the instance URL to test
+  instance_url = "https://valid-instance"
+  api_key = "test-api-key"
+  # Assert that the instance URL is returned if it is valid
+  assert check_instance_is_valid(instance_url, api_key=api_key) == instance_url
+  mock_get.assert_called_once_with(
+      f"{instance_url}/node?nodes=country%2FGTM&property=->name",
+      headers={"X-API-Key": api_key})
 
 
 @patch("requests.get")

datacommons_client/utils/request_handling.py

@@ -14,15 +14,19 @@
 CUSTOM_DC_V2: str = "/core/api/v2"
 
 
-def check_instance_is_valid(instance_url: str) -> str:
+def check_instance_is_valid(instance_url: str,
+                            api_key: str | None = None) -> str:
   """Check that the given instance URL points to a valid Data Commons instance.
 
     This function attempts a GET request against a known node in Data Commons to
     validate the given instance URL. If the node is found and the response has the
     expected data, the URL is considered valid.
+    
+    If an api_key is provided, it will be included in the request headers.
 
     Args:
         instance_url: The Data Commons instance URL to validate.
+        api_key: Optional API key for authentication.
 
     Returns:
         The validated instance URL.
@@ -34,8 +38,12 @@ def check_instance_is_valid(instance_url: str) -> str:
   # Test URL for a known node in Data Commons
   test_url = f"{instance_url}/node?nodes=country%2FGTM&property=->name"
 
+  headers = {}
+  if api_key:
+    headers["X-API-Key"] = api_key
+
   try:
-    response = requests.get(test_url)
+    response = requests.get(test_url, headers=headers)
     response.raise_for_status()
   except requests.exceptions.RequestException as exc:
     raise InvalidDCInstanceError(exc.response) from exc