Skip to content

Commit

Permalink
Update for backward compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
tolik0 committed Feb 7, 2025
1 parent 824d2c6 commit 15f830c
Show file tree
Hide file tree
Showing 5 changed files with 165 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1478,7 +1478,7 @@ definitions:
description: List of matchers that define which requests this policy applies to.
type: array
items:
"$ref": "#/definitions/HttpRequestMatcher"
"$ref": "#/definitions/HttpRequestRegexMatcher"
additionalProperties: true
MovingWindowCallRatePolicy:
title: Moving Window Call Rate Policy
Expand All @@ -1503,7 +1503,7 @@ definitions:
description: List of matchers that define which requests this policy applies to.
type: array
items:
"$ref": "#/definitions/HttpRequestMatcher"
"$ref": "#/definitions/HttpRequestRegexMatcher"
additionalProperties: true
UnlimitedCallRatePolicy:
title: Unlimited Call Rate Policy
Expand All @@ -1521,7 +1521,7 @@ definitions:
description: List of matchers that define which requests this policy applies to.
type: array
items:
"$ref": "#/definitions/HttpRequestMatcher"
"$ref": "#/definitions/HttpRequestRegexMatcher"
additionalProperties: true
Rate:
title: Rate
Expand All @@ -1541,7 +1541,7 @@ definitions:
type: string
format: duration
additionalProperties: true
HttpRequestMatcher:
HttpRequestRegexMatcher:
title: HTTP Request Matcher
description: >
Matches HTTP requests based on method, base URL, URL path pattern, query parameters, and headers.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -657,7 +657,7 @@ class Config:
)


class HttpRequestMatcher(BaseModel):
class HttpRequestRegexMatcher(BaseModel):
class Config:
extra = Extra.allow

Expand Down Expand Up @@ -1642,7 +1642,7 @@ class Config:
description="The maximum number of calls allowed within the period.",
title="Call Limit",
)
matchers: List[HttpRequestMatcher] = Field(
matchers: List[HttpRequestRegexMatcher] = Field(
...,
description="List of matchers that define which requests this policy applies to.",
title="Matchers",
Expand All @@ -1659,7 +1659,7 @@ class Config:
description="List of rates that define the call limits for different time intervals.",
title="Rates",
)
matchers: List[HttpRequestMatcher] = Field(
matchers: List[HttpRequestRegexMatcher] = Field(
...,
description="List of matchers that define which requests this policy applies to.",
title="Matchers",
Expand All @@ -1671,7 +1671,7 @@ class Config:
extra = Extra.allow

type: Literal["UnlimitedCallRatePolicy"]
matchers: List[HttpRequestMatcher] = Field(
matchers: List[HttpRequestRegexMatcher] = Field(
...,
description="List of matchers that define which requests this policy applies to.",
title="Matchers",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@
HttpRequester as HttpRequesterModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
HttpRequestMatcher as HttpRequestMatcherModel,
HttpRequestRegexMatcher as HttpRequestRegexMatcherModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
HttpResponseFilter as HttpResponseFilterModel,
Expand Down Expand Up @@ -494,7 +494,7 @@
APIBudget,
FixedWindowCallRatePolicy,
HttpAPIBudget,
HttpRequestMatcher,
HttpRequestRegexMatcher,
MovingWindowCallRatePolicy,
Rate,
UnlimitedCallRatePolicy,
Expand Down Expand Up @@ -644,7 +644,7 @@ def _init_mappings(self) -> None:
MovingWindowCallRatePolicyModel: self.create_moving_window_call_rate_policy,
UnlimitedCallRatePolicyModel: self.create_unlimited_call_rate_policy,
RateModel: self.create_rate,
HttpRequestMatcherModel: self.create_http_request_matcher,
HttpRequestRegexMatcherModel: self.create_http_request_matcher,
}

# Needed for the case where we need to perform a second parse on the fields of a custom component
Expand Down Expand Up @@ -3040,9 +3040,9 @@ def create_rate(self, model: RateModel, config: Config, **kwargs: Any) -> Rate:
)

def create_http_request_matcher(
self, model: HttpRequestMatcherModel, config: Config, **kwargs: Any
) -> HttpRequestMatcher:
return HttpRequestMatcher(
self, model: HttpRequestRegexMatcherModel, config: Config, **kwargs: Any
) -> HttpRequestRegexMatcher:
return HttpRequestRegexMatcher(
method=model.method,
url_base=model.url_base,
url_path_pattern=model.url_path_pattern,
Expand Down
63 changes: 63 additions & 0 deletions airbyte_cdk/sources/streams/call_rate.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,69 @@ def __call__(self, request: Any) -> bool:


class HttpRequestMatcher(RequestMatcher):
"""Simple implementation of RequestMatcher for http requests case"""

def __init__(
self,
method: Optional[str] = None,
url: Optional[str] = None,
params: Optional[Mapping[str, Any]] = None,
headers: Optional[Mapping[str, Any]] = None,
):
"""Constructor
:param method:
:param url:
:param params:
:param headers:
"""
self._method = method
self._url = url
self._params = {str(k): str(v) for k, v in (params or {}).items()}
self._headers = {str(k): str(v) for k, v in (headers or {}).items()}

@staticmethod
def _match_dict(obj: Mapping[str, Any], pattern: Mapping[str, Any]) -> bool:
"""Check that all elements from pattern dict present and have the same values in obj dict
:param obj:
:param pattern:
:return:
"""
return pattern.items() <= obj.items()

def __call__(self, request: Any) -> bool:
"""
:param request:
:return: True if matches the provided request object, False - otherwise
"""
if isinstance(request, requests.Request):
prepared_request = request.prepare()
elif isinstance(request, requests.PreparedRequest):
prepared_request = request
else:
return False

if self._method is not None:
if prepared_request.method != self._method:
return False
if self._url is not None and prepared_request.url is not None:
url_without_params = prepared_request.url.split("?")[0]
if url_without_params != self._url:
return False
if self._params is not None:
parsed_url = parse.urlsplit(prepared_request.url)
params = dict(parse.parse_qsl(str(parsed_url.query)))
if not self._match_dict(params, self._params):
return False
if self._headers is not None:
if not self._match_dict(prepared_request.headers, self._headers):
return False
return True


class HttpRequestRegexMatcher(RequestMatcher):
"""
Extended RequestMatcher for HTTP requests that supports matching on:
- HTTP method (case-insensitive)
Expand Down
88 changes: 88 additions & 0 deletions unit_tests/sources/streams/test_call_rate.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
CallRateLimitHit,
FixedWindowCallRatePolicy,
HttpRequestMatcher,
HttpRequestRegexMatcher,
MovingWindowCallRatePolicy,
Rate,
UnlimitedCallRatePolicy,
Expand Down Expand Up @@ -357,3 +358,90 @@ def test_with_cache(self, mocker, requests_mock):
assert next(records) == {"data": "some_data"}

assert MovingWindowCallRatePolicy.try_acquire.call_count == 1


class TestHttpRequestRegexMatcher:
"""
Tests for the new regex-based logic:
- Case-insensitive HTTP method matching
- Optional url_base (scheme://netloc)
- Regex-based path matching
- Query params (must be present)
- Headers (case-insensitive keys)
"""

def test_case_insensitive_method(self):
matcher = HttpRequestRegexMatcher(method="GET")

req_ok = Request("get", "https://example.com/test/path")
req_wrong = Request("POST", "https://example.com/test/path")

assert matcher(req_ok)
assert not matcher(req_wrong)

def test_url_base(self):
matcher = HttpRequestRegexMatcher(url_base="https://example.com")

req_ok = Request("GET", "https://example.com/test/path?foo=bar")
req_wrong = Request("GET", "https://another.com/test/path?foo=bar")

assert matcher(req_ok)
assert not matcher(req_wrong)

def test_url_path_pattern(self):
matcher = HttpRequestRegexMatcher(url_path_pattern=r"/test/")

req_ok = Request("GET", "https://example.com/test/something")
req_wrong = Request("GET", "https://example.com/other/something")

assert matcher(req_ok)
assert not matcher(req_wrong)

def test_query_params(self):
matcher = HttpRequestRegexMatcher(params={"foo": "bar"})

req_ok = Request("GET", "https://example.com/api?foo=bar&extra=123")
req_missing = Request("GET", "https://example.com/api?not_foo=bar")

assert matcher(req_ok)
assert not matcher(req_missing)

def test_headers_case_insensitive(self):
matcher = HttpRequestRegexMatcher(headers={"X-Custom-Header": "abc"})

req_ok = Request(
"GET",
"https://example.com/api?foo=bar",
headers={"x-custom-header": "abc", "other": "123"},
)
req_wrong = Request("GET", "https://example.com/api", headers={"x-custom-header": "wrong"})

assert matcher(req_ok)
assert not matcher(req_wrong)

def test_combined_criteria(self):
matcher = HttpRequestRegexMatcher(
method="GET",
url_base="https://example.com",
url_path_pattern=r"/test/",
params={"foo": "bar"},
headers={"X-Test": "123"},
)

req_ok = Request("GET", "https://example.com/test/me?foo=bar", headers={"x-test": "123"})
req_bad_base = Request(
"GET", "https://other.com/test/me?foo=bar", headers={"x-test": "123"}
)
req_bad_path = Request("GET", "https://example.com/nope?foo=bar", headers={"x-test": "123"})
req_bad_param = Request(
"GET", "https://example.com/test/me?extra=xyz", headers={"x-test": "123"}
)
req_bad_header = Request(
"GET", "https://example.com/test/me?foo=bar", headers={"some-other-header": "xyz"}
)

assert matcher(req_ok)
assert not matcher(req_bad_base)
assert not matcher(req_bad_path)
assert not matcher(req_bad_param)
assert not matcher(req_bad_header)

0 comments on commit 15f830c

Please sign in to comment.