Add cost callback for rest calls

pyproject.toml

@@ -9,10 +9,10 @@ license = "Apache 2.0"
 PyJWT = "2.4.0"
 python = "<3.11,>=3.7.2"
 requests = "^2.25.1"
-singer-sdk = "^0.4.9"
+# singer-sdk = "^0.4.9"
 # For local SDK dev:
 # singer-sdk = {path = "../singer-sdk", develop = true}
-# singer-sdk = {git = "https://gitlab.com/meltano/singer-sdk.git", rev = "97-hierarchical-streams"}
+singer-sdk = {git = "https://github.com/oviohub/meltano-sdk.git", rev = "040a8383f3994f47d5a324da81e2cb7b559ee23c"}
 types-simplejson = "^3.17.2"
 types-python-dateutil = "^2.8.6"
 nested-lookup = "^0.2.23"

tap_github/client.py

@@ -258,6 +258,15 @@ def backoff_handler(self, details: dict) -> None:
             self.authenticator.get_next_auth_token()
             prepared_request.headers.update(self.authenticator.auth_headers or {})
 
+    def calculate_sync_cost(
+        self,
+        request: requests.PreparedRequest,
+        response: requests.Response,
+        context: Optional[dict],
+    ) -> Dict[str, int]:
+        """Return the cost of the last REST API call."""
+        return {"rest": 1, "graphql": 0, "search": 0}
+
 
 class GitHubGraphqlStream(GraphQLStream, GitHubRestStream):
     """GitHub Graphql stream class."""
@@ -354,3 +363,14 @@ def get_url_params(
             params["since"] = str(since)
 
         return params
+
+    def calculate_sync_cost(
+        self,
+        request: requests.PreparedRequest,
+        response: requests.Response,
+        context: Optional[dict],
+    ) -> Dict[str, int]:
+        """Return the cost of the last graphql API call."""
+        costgen = extract_jsonpath("$.data.rateLimit.cost", input=response.json())
+        cost = next(costgen)
+        return {"rest": 0, "graphql": int(cost), "search": 0}

tap_github/repository_streams.py

@@ -90,7 +90,7 @@ def query(self) -> str:
                         f'repo{i}: repository(name: "{repo[1]}", owner: "{repo[0]}") '
                         "{ nameWithOwner databaseId }"
                     )
-                return "query {" + " ".join(chunks) + " }"
+                return "query {" + " ".join(chunks) + " rateLimit { cost } }"
 
         repos_with_ids: list = list()
         temp_stream = TempStream(self._tap, list(repo_list))
@@ -101,6 +101,8 @@ def query(self) -> str:
         # the line.
         for record in temp_stream.request_records({}):
             for item in record.keys():
+                if item == "rateLimit":
+                    continue
                 try:
                     org, repo = record[item]["nameWithOwner"].split("/")
                 except TypeError:
@@ -1612,6 +1614,9 @@ def query(self) -> str:
                 }
               }
             }
+            rateLimit {
+              cost
+            }
           }
         """
 
@@ -2103,6 +2108,9 @@ def query(self) -> str:
                 }
               }
             }
+            rateLimit {
+              cost
+            }
           }
 
         """

tap_github/streams.py

@@ -36,6 +36,7 @@
     RepositoryStream,
     ReviewCommentsStream,
     ReviewsStream,
+    StargazersGraphqlStream,
     StargazersStream,
     StatsContributorsStream,
     WorkflowRunJobsStream,
@@ -86,6 +87,7 @@ def __init__(self, valid_queries: Set[str], streams: List[Type[Stream]]):
             RepositoryStream,
             ReviewCommentsStream,
             ReviewsStream,
+            StargazersGraphqlStream,
             StargazersStream,
             StatsContributorsStream,
             WorkflowRunJobsStream,

tap_github/user_streams.py

@@ -73,7 +73,7 @@ def query(self) -> str:
                         f'user{i}: repositoryOwner(login: "{user}") '
                         "{ login avatarUrl}"
                     )
-                return "query {" + " ".join(chunks) + " }"
+                return "query {" + " ".join(chunks) + " rateLimit { cost } }"
 
         users_with_ids: list = list()
         temp_stream = TempStream(self._tap, list(user_list))
@@ -88,6 +88,8 @@ def query(self) -> str:
         # the line.
         for record in temp_stream.request_records({}):
             for item in record.keys():
+                if item == "rateLimit":
+                    continue
                 try:
                     username = record[item]["login"]
                 except TypeError:
@@ -293,6 +295,9 @@ def query(self) -> str:
                 }
               }
             }
+            rateLimit {
+              cost
+            }
           }
         """