Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(airbyte-cdk): add gzipjson decoder #20

Merged
merged 17 commits into from
Nov 14, 2024
Merged
43 changes: 43 additions & 0 deletions airbyte_cdk/sources/declarative/declarative_component_schema.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -1750,6 +1750,45 @@ definitions:
type:
type: string
enum: [XmlDecoder]
CustomDecoder:
title: Custom Decoder
description: Use this to implement custom decoder logic.
type: object
additionalProperties: true
required:
- type
- class_name
properties:
type:
type: string
enum: [CustomDecoder]
class_name:
title: Class Name
description: Fully-qualified name of the class that will be implementing the custom decoding. Has to be a sub class of Decoder. The format is `source_<name>.<package>.<class_name>`.
type: string
additionalProperties: true
examples:
- "source_amazon_ads.components.GzipJsonlDecoder"
$parameters:
type: object
additionalProperties: true
GzipJsonDecoder:
title: GzipJson Decoder
description: Use this if the response is Gzip compressed Json.
type: object
additionalProperties: true
required:
- type
properties:
type:
type: string
enum: [GzipJsonDecoder]
encoding:
type: string
default: utf-8
$parameters:
type: object
additionalProperties: true
ListPartitionRouter:
title: List Partition Router
description: A Partition router that specifies a list of attributes where each attribute describes a portion of the complete data set for a stream. During a sync, each value is iterated over and can be used as input to outbound API requests.
Expand Down Expand Up @@ -2404,10 +2443,12 @@ definitions:
title: Decoder
description: Component decoding the response so records can be extracted.
anyOf:
- "$ref": "#/definitions/CustomDecoder"
- "$ref": "#/definitions/JsonDecoder"
- "$ref": "#/definitions/JsonlDecoder"
- "$ref": "#/definitions/IterableDecoder"
- "$ref": "#/definitions/XmlDecoder"
- "$ref": "#/definitions/GzipJsonDecoder"
$parameters:
type: object
additionalProperties: true
Expand Down Expand Up @@ -2520,10 +2561,12 @@ definitions:
title: Decoder
description: Component decoding the response so records can be extracted.
anyOf:
- "$ref": "#/definitions/CustomDecoder"
- "$ref": "#/definitions/JsonDecoder"
- "$ref": "#/definitions/JsonlDecoder"
- "$ref": "#/definitions/IterableDecoder"
- "$ref": "#/definitions/XmlDecoder"
- "$ref": "#/definitions/GzipJsonDecoder"
$parameters:
type: object
additionalProperties: true
Expand Down
4 changes: 2 additions & 2 deletions airbyte_cdk/sources/declarative/decoders/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@
#

from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder, JsonlDecoder, IterableDecoder
from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder, JsonlDecoder, IterableDecoder, GzipJsonDecoder
artem1205 marked this conversation as resolved.
Show resolved Hide resolved
maxi297 marked this conversation as resolved.
Show resolved Hide resolved
from airbyte_cdk.sources.declarative.decoders.noop_decoder import NoopDecoder
from airbyte_cdk.sources.declarative.decoders.pagination_decoder_decorator import PaginationDecoderDecorator
from airbyte_cdk.sources.declarative.decoders.xml_decoder import XmlDecoder

__all__ = ["Decoder", "JsonDecoder", "JsonlDecoder", "IterableDecoder", "NoopDecoder", "PaginationDecoderDecorator", "XmlDecoder"]
__all__ = ["Decoder", "JsonDecoder", "JsonlDecoder", "IterableDecoder", "GzipJsonDecoder", "NoopDecoder", "PaginationDecoderDecorator", "XmlDecoder"]
36 changes: 26 additions & 10 deletions airbyte_cdk/sources/declarative/decoders/json_decoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@

import logging
from dataclasses import InitVar, dataclass
from typing import Any, Generator, Mapping
from gzip import decompress
from typing import Any, Generator, Mapping, MutableMapping

import requests
from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
Expand All @@ -24,24 +25,30 @@ class JsonDecoder(Decoder):
def is_stream_response(self) -> bool:
return False

def decode(self, response: requests.Response) -> Generator[Mapping[str, Any], None, None]:
def decode(self, response: requests.Response) -> Generator[MutableMapping[str, Any], None, None]:
"""
Given the response is an empty string or an emtpy list, the function will return a generator with an empty mapping.
"""
try:
body_json = response.json()
if not isinstance(body_json, list):
body_json = [body_json]
if len(body_json) == 0:
yield {}
else:
yield from body_json
yield from self.parse_body_json(body_json)
except requests.exceptions.JSONDecodeError:
logger.warning(
f"Response cannot be parsed into json: {response.status_code=}, {response.text=}"
)
yield {}

@staticmethod
def parse_body_json(
body_json: Mapping[str, Any] | list,
) -> Generator[MutableMapping[str, Any], None, None]:
if not isinstance(body_json, list):
body_json = [body_json]
if len(body_json) == 0:
yield {}
else:
yield from body_json


@dataclass
class IterableDecoder(Decoder):
Expand All @@ -54,7 +61,7 @@ class IterableDecoder(Decoder):
def is_stream_response(self) -> bool:
return True

def decode(self, response: requests.Response) -> Generator[Mapping[str, Any], None, None]:
def decode(self, response: requests.Response) -> Generator[MutableMapping[str, Any], None, None]:
for line in response.iter_lines():
yield {"record": line.decode()}

Expand All @@ -70,8 +77,17 @@ class JsonlDecoder(Decoder):
def is_stream_response(self) -> bool:
return True

def decode(self, response: requests.Response) -> Generator[Mapping[str, Any], None, None]:
def decode(self, response: requests.Response) -> Generator[MutableMapping[str, Any], None, None]:
# TODO???: set delimiter? usually it is `\n` but maybe it would be useful to set optional?
# https://github.com/airbytehq/airbyte-internal-issues/issues/8436
for record in response.iter_lines():
yield orjson.loads(record)


@dataclass
class GzipJsonDecoder(JsonDecoder):
artem1205 marked this conversation as resolved.
Show resolved Hide resolved
encoding: str = "utf-8"

def decode(self, response: requests.Response) -> Generator[MutableMapping[str, Any], None, None]:
raw_string = decompress(response.content).decode(encoding=self.encoding)
yield from self.parse_body_json(orjson.loads(raw_string))
Original file line number Diff line number Diff line change
Expand Up @@ -687,6 +687,29 @@ class XmlDecoder(BaseModel):
type: Literal["XmlDecoder"]


class CustomDecoder(BaseModel):
class Config:
extra = Extra.allow

type: Literal["CustomDecoder"]
class_name: str = Field(
...,
description="Fully-qualified name of the class that will be implementing the custom decoding. Has to be a sub class of Decoder. The format is `source_<name>.<package>.<class_name>`.",
examples=["source_amazon_ads.components.GzipJsonlDecoder"],
title="Class Name",
)
parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")


class GzipJsonDecoder(BaseModel):
class Config:
extra = Extra.allow

type: Literal["GzipJsonDecoder"]
encoding: Optional[str] = "utf-8"
parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")


class MinMaxDatetime(BaseModel):
type: Literal["MinMaxDatetime"]
datetime: str = Field(
Expand Down Expand Up @@ -1620,7 +1643,16 @@ class SimpleRetriever(BaseModel):
description="PartitionRouter component that describes how to partition the stream, enabling incremental syncs and checkpointing.",
title="Partition Router",
)
decoder: Optional[Union[JsonDecoder, JsonlDecoder, IterableDecoder, XmlDecoder]] = Field(
decoder: Optional[
Union[
CustomDecoder,
JsonDecoder,
JsonlDecoder,
IterableDecoder,
XmlDecoder,
GzipJsonDecoder,
]
] = Field(
None,
description="Component decoding the response so records can be extracted.",
title="Decoder",
Expand Down Expand Up @@ -1680,7 +1712,16 @@ class AsyncRetriever(BaseModel):
description="PartitionRouter component that describes how to partition the stream, enabling incremental syncs and checkpointing.",
title="Partition Router",
)
decoder: Optional[Union[JsonDecoder, JsonlDecoder, IterableDecoder, XmlDecoder]] = Field(
decoder: Optional[
Union[
CustomDecoder,
JsonDecoder,
JsonlDecoder,
IterableDecoder,
XmlDecoder,
GzipJsonDecoder,
]
] = Field(
None,
description="Component decoding the response so records can be extracted.",
title="Decoder",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
from airbyte_cdk.sources.declarative.decoders import (
Decoder,
GzipJsonDecoder,
IterableDecoder,
JsonDecoder,
JsonlDecoder,
Expand Down Expand Up @@ -134,6 +135,9 @@
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
CustomBackoffStrategy as CustomBackoffStrategyModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
CustomDecoder as CustomDecoderModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
CustomErrorHandler as CustomErrorHandlerModel,
)
Expand Down Expand Up @@ -182,6 +186,9 @@
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
ExponentialBackoffStrategy as ExponentialBackoffStrategyModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
GzipJsonDecoder as GzipJsonDecoderModel,
)
from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
HttpRequester as HttpRequesterModel,
)
Expand Down Expand Up @@ -402,6 +409,7 @@ def _init_mappings(self) -> None:
CursorPaginationModel: self.create_cursor_pagination,
CustomAuthenticatorModel: self.create_custom_component,
CustomBackoffStrategyModel: self.create_custom_component,
CustomDecoderModel: self.create_custom_component,
CustomErrorHandlerModel: self.create_custom_component,
CustomIncrementalSyncModel: self.create_custom_component,
CustomRecordExtractorModel: self.create_custom_component,
Expand All @@ -425,6 +433,7 @@ def _init_mappings(self) -> None:
InlineSchemaLoaderModel: self.create_inline_schema_loader,
JsonDecoderModel: self.create_json_decoder,
JsonlDecoderModel: self.create_jsonl_decoder,
GzipJsonDecoderModel: self.create_gzipjson_decoder,
KeysToLowerModel: self.create_keys_to_lower_transformation,
IterableDecoderModel: self.create_iterable_decoder,
XmlDecoderModel: self.create_xml_decoder,
Expand Down Expand Up @@ -1548,6 +1557,12 @@ def create_iterable_decoder(
def create_xml_decoder(model: XmlDecoderModel, config: Config, **kwargs: Any) -> XmlDecoder:
return XmlDecoder(parameters={})

@staticmethod
def create_gzipjson_decoder(
model: GzipJsonDecoderModel, config: Config, **kwargs: Any
) -> GzipJsonDecoder:
return GzipJsonDecoder(parameters={}, encoding=model.encoding)

@staticmethod
def create_json_file_schema_loader(
model: JsonFileSchemaLoaderModel, config: Config, **kwargs: Any
Expand Down
77 changes: 77 additions & 0 deletions unit_tests/sources/declarative/decoders/test_json_decoder.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#
import gzip
import json
import os

Expand All @@ -13,6 +14,7 @@
from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
ModelToComponentFactory,
)
from airbyte_cdk.sources.declarative.decoders import GzipJsonDecoder


@pytest.mark.parametrize(
Expand Down Expand Up @@ -120,3 +122,78 @@ def get_body():
counter += 1

assert counter == lines_in_response * len(stream_slices)


@pytest.mark.parametrize(
"encoding",
[
"utf-8",
"utf",
],
ids=["utf-8", "utf"],
)
def test_gzipjson_decoder(requests_mock, encoding):
response_to_compress = json.dumps(
[
{
"campaignId": 214078428,
"campaignName": "sample-campaign-name-214078428",
"adGroupId": "6490134",
"adId": "665320125",
"targetId": "791320341",
"asin": "G000PSH142",
"advertisedAsin": "G000PSH142",
"keywordBid": "511234974",
"keywordId": "965783021",
},
{
"campaignId": 44504582,
"campaignName": "sample-campaign-name-44504582",
"adGroupId": "6490134",
"adId": "665320125",
"targetId": "791320341",
"asin": "G000PSH142",
"advertisedAsin": "G000PSH142",
"keywordBid": "511234974",
"keywordId": "965783021",
},
{
"campaignId": 509144838,
"campaignName": "sample-campaign-name-509144838",
"adGroupId": "6490134",
"adId": "665320125",
"targetId": "791320341",
"asin": "G000PSH142",
"advertisedAsin": "G000PSH142",
"keywordBid": "511234974",
"keywordId": "965783021",
},
{
"campaignId": 231712082,
"campaignName": "sample-campaign-name-231712082",
"adGroupId": "6490134",
"adId": "665320125",
"targetId": "791320341",
"asin": "G000PSH142",
"advertisedAsin": "G000PSH142",
"keywordBid": "511234974",
"keywordId": "965783021",
},
{
"campaignId": 895306040,
"campaignName": "sample-campaign-name-895306040",
"adGroupId": "6490134",
"adId": "665320125",
"targetId": "791320341",
"asin": "G000PSH142",
"advertisedAsin": "G000PSH142",
"keywordBid": "511234974",
"keywordId": "965783021",
},
]
)
body = gzip.compress(response_to_compress.encode(encoding))

requests_mock.register_uri("GET", "https://airbyte.io/", content=body)
response = requests.get("https://airbyte.io/")
assert len(list(GzipJsonDecoder(parameters={}, encoding=encoding).decode(response))) == 5
Loading