From 6c9dab824f2f0cbab209732dd4e88fdcd5ebaf1c Mon Sep 17 00:00:00 2001 From: "Aaron (\"AJ\") Steers" Date: Thu, 11 Apr 2024 20:00:30 -0700 Subject: [PATCH] Chore: Bump `airbyte-api` (#186) --- airbyte/_util/api_imports.py | 42 ++++++ airbyte/_util/api_util.py | 122 +++++++++--------- airbyte/cloud/_destination_util.py | 2 +- airbyte/cloud/connections.py | 3 +- airbyte/cloud/sync_results.py | 3 +- airbyte/cloud/workspaces.py | 3 +- poetry.lock | 30 ++--- pyproject.toml | 9 +- .../cloud/test_cloud_api_util.py | 2 +- 9 files changed, 120 insertions(+), 96 deletions(-) create mode 100644 airbyte/_util/api_imports.py diff --git a/airbyte/_util/api_imports.py b/airbyte/_util/api_imports.py new file mode 100644 index 00000000..19935b6f --- /dev/null +++ b/airbyte/_util/api_imports.py @@ -0,0 +1,42 @@ +# Copyright (c) 2024 Airbyte, Inc., all rights reserved. +"""Imported classes from the Airbyte API. + +Any classes that are imported from the Airbyte API should be imported here. +This allows for easy access to these classes in other modules, especially +for type hinting purposes. + +Design Guidelines: +- No modules except `api_util` and `api_imports` should import from `airbyte_api`. +- If a module needs to import from `airbyte_api`, it should import from `api_imports` (this module) + instead. +- This module is divided into two sections: internal-use classes and public-use classes. +- Public-use classes should be carefully reviewed to ensure that they are necessary for public use + and that we are willing to support them as part of PyAirbyte. +""" +# Ignore import sorting in this file. Manual grouping is more important. +# ruff: noqa: I001 + +from __future__ import annotations + +# Internal-Use Classes + +# These classes are used internally to cache API responses. +from airbyte_api.models import ( + ConnectionResponse, + DestinationResponse, + JobResponse, +) + +# Public-Use Classes + +# This class is used to represent the status of a job. It may be used in +# type hints for public functions that return a job status. +from airbyte_api.models import JobStatusEnum # Alias not needed + + +__all__: list[str] = [ + "ConnectionResponse", + "DestinationResponse", + "JobResponse", + "JobStatusEnum", +] diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 7b0edd30..6cc5411a 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -17,9 +17,7 @@ from typing import Any import airbyte_api -from airbyte_api.models import operations as api_operations -from airbyte_api.models import shared as api_models -from airbyte_api.models.shared.jobcreaterequest import JobCreateRequest, JobTypeEnum +from airbyte_api import api, models from airbyte.exceptions import ( AirbyteConnectionSyncError, @@ -47,8 +45,8 @@ def get_airbyte_server_instance( api_root: str, ) -> airbyte_api.Airbyte: """Get an Airbyte instance.""" - return airbyte_api.Airbyte( - security=api_models.Security( + return airbyte_api.AirbyteAPI( + security=models.Security( bearer_auth=api_key, ), server_url=api_root, @@ -63,14 +61,14 @@ def get_workspace( *, api_root: str, api_key: str, -) -> api_models.WorkspaceResponse: +) -> models.WorkspaceResponse: """Get a connection.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) response = airbyte_instance.workspaces.get_workspace( - api_operations.GetWorkspaceRequest( + api.GetWorkspaceRequest( workspace_id=workspace_id, ), ) @@ -94,7 +92,7 @@ def list_connections( *, api_root: str, api_key: str, -) -> list[api_models.ConnectionResponse]: +) -> list[api.ConnectionResponse]: """Get a connection.""" _ = workspace_id # Not used (yet) airbyte_instance = get_airbyte_server_instance( @@ -102,7 +100,7 @@ def list_connections( api_root=api_root, ) response = airbyte_instance.connections.list_connections( - api_operations.ListConnectionsRequest()( + api.ListConnectionsRequest()( workspace_ids=[workspace_id], ), ) @@ -124,7 +122,7 @@ def get_connection( *, api_root: str, api_key: str, -) -> api_models.ConnectionResponse: +) -> api.ConnectionResponse: """Get a connection.""" _ = workspace_id # Not used (yet) airbyte_instance = get_airbyte_server_instance( @@ -132,7 +130,7 @@ def get_connection( api_root=api_root, ) response = airbyte_instance.connections.get_connection( - api_operations.GetConnectionRequest( + api.GetConnectionRequest( connection_id=connection_id, ), ) @@ -148,7 +146,7 @@ def run_connection( *, api_root: str, api_key: str, -) -> api_models.ConnectionResponse: +) -> api.ConnectionResponse: """Get a connection. If block is True, this will block until the connection is finished running. @@ -161,9 +159,9 @@ def run_connection( api_root=api_root, ) response = airbyte_instance.jobs.create_job( - JobCreateRequest( + models.JobCreateRequest( connection_id=connection_id, - job_type=JobTypeEnum.SYNC, + job_type=models.JobTypeEnum.SYNC, ), ) if status_ok(response.status_code) and response.job_response: @@ -188,14 +186,14 @@ def get_job_logs( *, api_root: str, api_key: str, -) -> list[api_models.JobResponse]: +) -> list[api.JobResponse]: """Get a job's logs.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) - response: api_operations.ListJobsResponse = airbyte_instance.jobs.list_jobs( - api_operations.ListJobsRequest( + response: api.ListJobsResponse = airbyte_instance.jobs.list_jobs( + api.ListJobsRequest( workspace_ids=[workspace_id], connection_id=connection_id, limit=limit, @@ -219,14 +217,14 @@ def get_job_info( *, api_root: str, api_key: str, -) -> api_models.JobResponse: +) -> api.JobResponse: """Get a job.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) response = airbyte_instance.jobs.get_job( - api_operations.GetJobRequest( + api.GetJobRequest( job_id=job_id, ), ) @@ -246,14 +244,14 @@ def create_source( config: dict[str, Any], api_root: str, api_key: str, -) -> api_models.SourceResponse: +) -> api.SourceResponse: """Get a connection.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) - response: api_operations.CreateSourceResponse = airbyte_instance.sources.create_source( - api_models.SourceCreateRequest( + response: api.CreateSourceResponse = airbyte_instance.sources.create_source( + models.SourceCreateRequest( name=name, workspace_id=workspace_id, configuration=config, # TODO: wrap in a proper configuration object @@ -275,14 +273,14 @@ def get_source( *, api_root: str, api_key: str, -) -> api_models.SourceResponse: +) -> api.SourceResponse: """Get a connection.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) response = airbyte_instance.sources.get_source( - api_operations.GetSourceRequest( + api.GetSourceRequest( source_id=source_id, ), ) @@ -306,7 +304,7 @@ def delete_source( api_root=api_root, ) response = airbyte_instance.sources.delete_source( - api_operations.DeleteSourceRequest( + api.DeleteSourceRequest( source_id=source_id, ), ) @@ -329,20 +327,18 @@ def create_destination( config: dict[str, Any], api_root: str, api_key: str, -) -> api_models.DestinationResponse: +) -> api.DestinationResponse: """Get a connection.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) - response: api_operations.CreateDestinationResponse = ( - airbyte_instance.destinations.create_destination( - api_models.DestinationCreateRequest( - name=name, - workspace_id=workspace_id, - configuration=config, # TODO: wrap in a proper configuration object - ), - ) + response: api.CreateDestinationResponse = airbyte_instance.destinations.create_destination( + models.DestinationCreateRequest( + name=name, + workspace_id=workspace_id, + configuration=config, # TODO: wrap in a proper configuration object + ), ) if status_ok(response.status_code) and response.destination_response: return response.destination_response @@ -358,14 +354,14 @@ def get_destination( *, api_root: str, api_key: str, -) -> api_models.DestinationResponse: +) -> api.DestinationResponse: """Get a connection.""" airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) response = airbyte_instance.destinations.get_destination( - api_operations.GetDestinationRequest( + api.GetDestinationRequest( destination_id=destination_id, ), ) @@ -376,19 +372,19 @@ def get_destination( raw_configuration: dict[str, Any] = raw_response["configuration"] destination_type = raw_response.get("destinationType") if destination_type == "snowflake": - response.destination_response.configuration = api_models.DestinationSnowflake.from_dict( + response.destination_response.configuration = models.DestinationSnowflake.from_dict( raw_configuration, ) if destination_type == "bigquery": - response.destination_response.configuration = api_models.DestinationBigquery.from_dict( + response.destination_response.configuration = models.DestinationBigquery.from_dict( raw_configuration, ) if destination_type == "postgres": - response.destination_response.configuration = api_models.DestinationPostgres.from_dict( + response.destination_response.configuration = models.DestinationPostgres.from_dict( raw_configuration, ) if destination_type == "duckdb": - response.destination_response.configuration = api_models.DestinationDuckdb.from_dict( + response.destination_response.configuration = models.DestinationDuckdb.from_dict( raw_configuration, ) @@ -411,7 +407,7 @@ def delete_destination( api_root=api_root, ) response = airbyte_instance.destinations.delete_destination( - api_operations.DeleteDestinationRequest( + api.DeleteDestinationRequest( destination_id=destination_id, ), ) @@ -437,23 +433,23 @@ def create_connection( workspace_id: str | None = None, prefix: str, selected_stream_names: list[str], -) -> api_models.ConnectionResponse: +) -> models.ConnectionResponse: _ = workspace_id # Not used (yet) airbyte_instance = get_airbyte_server_instance( api_key=api_key, api_root=api_root, ) - stream_configurations: list[api_models.StreamConfiguration] = [] + stream_configurations: list[models.StreamConfiguration] = [] if selected_stream_names: for stream_name in selected_stream_names: - stream_configuration = api_models.StreamConfiguration( + stream_configuration = models.StreamConfiguration( name=stream_name, ) stream_configurations.append(stream_configuration) - stream_configurations = api_models.StreamConfigurations(stream_configurations) + stream_configurations = models.StreamConfigurations(stream_configurations) response = airbyte_instance.connections.create_connection( - api_models.ConnectionCreateRequest( + models.ConnectionCreateRequest( name=name, source_id=source_id, destination_id=destination_id, @@ -479,14 +475,14 @@ def get_connection_by_name( *, api_root: str, api_key: str, -) -> api_models.ConnectionResponse: +) -> models.ConnectionResponse: """Get a connection.""" connections = list_connections( workspace_id=workspace_id, api_key=api_key, api_root=api_root, ) - found: list[api_models.ConnectionResponse] = [ + found: list[models.ConnectionResponse] = [ connection for connection in connections if connection.name == connection_name ] if len(found) == 0: @@ -519,7 +515,7 @@ def delete_connection( api_root=api_root, ) response = airbyte_instance.connections.delete_connection( - api_operations.DeleteConnectionRequest( + api.DeleteConnectionRequest( connection_id=connection_id, ), ) @@ -535,17 +531,17 @@ def delete_connection( # Not yet implemented -def check_source( - source_id: str, - *, - api_root: str, - api_key: str, - workspace_id: str | None = None, -) -> api_models.SourceCheckResponse: - """Check a source. - - # TODO: Need to use legacy Configuration API for this: - # https://airbyte-public-api-docs.s3.us-east-2.amazonaws.com/rapidoc-api-docs.html#post-/v1/sources/check_connection - """ - _ = source_id, workspace_id, api_root, api_key - raise NotImplementedError +# def check_source( +# source_id: str, +# *, +# api_root: str, +# api_key: str, +# workspace_id: str | None = None, +# ) -> api.SourceCheckResponse: +# """Check a source. + +# # TODO: Need to use legacy Configuration API for this: +# # https://airbyte-public-api-docs.s3.us-east-2.amazonaws.com/rapidoc-api-docs.html#post-/v1/sources/check_connection +# """ +# _ = source_id, workspace_id, api_root, api_key +# raise NotImplementedError diff --git a/airbyte/cloud/_destination_util.py b/airbyte/cloud/_destination_util.py index 08366398..eee8b9a8 100644 --- a/airbyte/cloud/_destination_util.py +++ b/airbyte/cloud/_destination_util.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from airbyte_api.models.shared import ( +from airbyte_api.models import ( DestinationBigquery, DestinationDuckdb, DestinationPostgres, diff --git a/airbyte/cloud/connections.py b/airbyte/cloud/connections.py index 52003264..9d1bf2e2 100644 --- a/airbyte/cloud/connections.py +++ b/airbyte/cloud/connections.py @@ -10,8 +10,7 @@ if TYPE_CHECKING: - from airbyte_api.models.shared.connectionresponse import ConnectionResponse - from airbyte_api.models.shared.jobresponse import JobResponse + from airbyte_api.models import ConnectionResponse, JobResponse from airbyte.cloud.workspaces import CloudWorkspace diff --git a/airbyte/cloud/sync_results.py b/airbyte/cloud/sync_results.py index 72ba7547..75eab5f3 100644 --- a/airbyte/cloud/sync_results.py +++ b/airbyte/cloud/sync_results.py @@ -9,9 +9,8 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, final -from airbyte_api.models.shared import ConnectionResponse, JobResponse, JobStatusEnum - from airbyte._util import api_util +from airbyte._util.api_imports import ConnectionResponse, JobResponse, JobStatusEnum from airbyte.cloud._destination_util import create_cache_from_destination_config from airbyte.datasets import CachedDataset from airbyte.exceptions import AirbyteConnectionSyncError, AirbyteConnectionSyncTimeoutError diff --git a/airbyte/cloud/workspaces.py b/airbyte/cloud/workspaces.py index d9c556c2..b9bb6a2f 100644 --- a/airbyte/cloud/workspaces.py +++ b/airbyte/cloud/workspaces.py @@ -28,8 +28,7 @@ if TYPE_CHECKING: - from airbyte_api.models.shared.destinationresponse import DestinationResponse - + from airbyte._util.api_imports import DestinationResponse from airbyte.caches.base import CacheBase diff --git a/poetry.lock b/poetry.lock index 9c67f0d9..885a063a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2,17 +2,19 @@ [[package]] name = "airbyte-api" -version = "0.47.3" +version = "0.49.2" description = "Python Client SDK for Airbyte API" optional = false python-versions = ">=3.8" -files = [] -develop = false +files = [ + {file = "airbyte-api-0.49.2.tar.gz", hash = "sha256:33432602a84f54700fd42d8a2e25574803a13543b467ae7e4948191fb0c17a9f"}, + {file = "airbyte_api-0.49.2-py3-none-any.whl", hash = "sha256:e41d3ed0247589570cdf4e3ef125b78569f15823cb6530e31bae6b4b064aed1d"}, +] [package.dependencies] certifi = ">=2023.7.22" charset-normalizer = ">=3.2.0" -dataclasses-json-speakeasy = ">=0.5.11" +dataclasses-json = ">=0.6.4" idna = ">=3.4" jsonpath-python = ">=1.0.6" marshmallow = ">=3.19.0" @@ -21,18 +23,12 @@ packaging = ">=23.1" python-dateutil = ">=2.8.2" requests = ">=2.31.0" six = ">=1.16.0" -typing_extensions = ">=4.7.1" +typing-extensions = ">=4.7.1" typing-inspect = ">=0.9.0" urllib3 = ">=1.26.18" [package.extras] -dev = ["pylint (==2.16.2)"] - -[package.source] -type = "git" -url = "https://github.com/airbytehq/airbyte-api-python-sdk.git" -reference = "856599a4861ee1f0ee4e994feff22e44ffb4cbd4" -resolved_reference = "856599a4861ee1f0ee4e994feff22e44ffb4cbd4" +dev = ["pylint (==3.1.0)"] [[package]] name = "airbyte-cdk" @@ -439,14 +435,14 @@ test = ["pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest test-randomorder = ["pytest-randomly"] [[package]] -name = "dataclasses-json-speakeasy" -version = "0.5.11" +name = "dataclasses-json" +version = "0.6.4" description = "Easily serialize dataclasses to and from JSON." optional = false python-versions = ">=3.7,<4.0" files = [ - {file = "dataclasses_json_speakeasy-0.5.11-py3-none-any.whl", hash = "sha256:ac52a069a01e8521015d682f37849bfdf056c36fa3f81497055e201fec684104"}, - {file = "dataclasses_json_speakeasy-0.5.11.tar.gz", hash = "sha256:418a987cea2ccf4e4be662f39faa5cc79b47b147c9d1a69d6928d6a27e0c17e8"}, + {file = "dataclasses_json-0.6.4-py3-none-any.whl", hash = "sha256:f90578b8a3177f7552f4e1a6e535e84293cd5da421fcce0642d49c0d7bdf8df2"}, + {file = "dataclasses_json-0.6.4.tar.gz", hash = "sha256:73696ebf24936560cca79a2430cbc4f3dd23ac7bf46ed17f38e5e5e7657a6377"}, ] [package.dependencies] @@ -2870,4 +2866,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "13b6f429df688ba505ffc513a714af167da17dc2acb34cd0749cda8d54183a73" +content-hash = "d1130167f74444477bee041a41686b12bd1d13d263cdf7e2b635d6ce9432d2ae" diff --git a/pyproject.toml b/pyproject.toml index bb48efa4..332a6db9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,14 +41,7 @@ ulid = "^1.1" # TODO: Remove this arbitrary python constraint once `sqlalchemy-bigquery` has done so. sqlalchemy-bigquery = { version = "1.9.0", python = "<3.13" } - -[tool.poetry.dependencies.airbyte-api] -git = "https://github.com/airbytehq/airbyte-api-python-sdk.git" -# Pinned to a specific commit to avoid breaking changes. -# TODO: Use a PyPi version of this after this resolves: -# https://github.com/airbytehq/airbyte-api-python-sdk/issues/67 -# rev = "aj/manual_rename_dir" This is the branch, but the commit is: -rev = "856599a4861ee1f0ee4e994feff22e44ffb4cbd4" +airbyte-api = "^0.49.2" [tool.poetry.group.dev.dependencies] docker = "^7.0.0" diff --git a/tests/integration_tests/cloud/test_cloud_api_util.py b/tests/integration_tests/cloud/test_cloud_api_util.py index 83dea4f6..0c94b186 100644 --- a/tests/integration_tests/cloud/test_cloud_api_util.py +++ b/tests/integration_tests/cloud/test_cloud_api_util.py @@ -9,7 +9,7 @@ import ulid from airbyte._util import api_util -from airbyte_api.models.shared import SourceFaker, DestinationDuckdb +from airbyte_api.models import SourceFaker, DestinationDuckdb def test_create_and_delete_source(