Chore: Bump CDK dependency, DuckDB, Postgres, Pandas, and others (#383)

Author: aaronsteers

airbyte/_processors/sql/postgres.py

@@ -29,7 +29,7 @@ class PostgresConfig(SqlConfig):
     def get_sql_alchemy_url(self) -> SecretString:
         """Return the SQLAlchemy URL to use."""
         return SecretString(
-            f"postgresql+psycopg2://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
+            f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
         )
 
     @overrides

airbyte/caches/_state_backend_base.py

@@ -7,7 +7,7 @@
 import abc
 from typing import TYPE_CHECKING
 
-from airbyte_protocol.models.airbyte_protocol import AirbyteStreamState
+from airbyte_protocol.models import AirbyteStreamState
 
 
 if TYPE_CHECKING:

airbyte/datasets/_sql.py

@@ -9,7 +9,7 @@
 from overrides import overrides
 from sqlalchemy import and_, func, select, text
 
-from airbyte_protocol.models.airbyte_protocol import ConfiguredAirbyteStream
+from airbyte_protocol.models import ConfiguredAirbyteStream
 
 from airbyte.constants import DEFAULT_ARROW_MAX_CHUNK_SIZE
 from airbyte.datasets._base import DatasetBase

airbyte/shared/state_providers.py

@@ -10,8 +10,8 @@
 from airbyte_protocol.models import (
     AirbyteStateMessage,
     AirbyteStateType,
+    AirbyteStreamState,
 )
-from airbyte_protocol.models.airbyte_protocol import AirbyteStreamState
 
 from airbyte import exceptions as exc
 

airbyte/shared/state_writers.py

@@ -11,7 +11,7 @@
 
 
 if TYPE_CHECKING:
-    from airbyte_protocol.models.airbyte_protocol import AirbyteStateMessage
+    from airbyte_protocol.models import AirbyteStateMessage
 
 
 class StateWriterBase(StateProviderBase, abc.ABC):

airbyte/sources/base.py

@@ -39,7 +39,7 @@
     from collections.abc import Generator, Iterable, Iterator
 
     from airbyte_cdk import ConnectorSpecification
-    from airbyte_protocol.models.airbyte_protocol import AirbyteStream
+    from airbyte_protocol.models import AirbyteStream
 
     from airbyte._executors.base import Executor
     from airbyte.caches import CacheBase

poetry.lock

@@ -15,30 +15,29 @@ enable = true
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"
 
-airbyte-cdk = "^4.6.2"
-duckdb = "^1.0.0"
-duckdb-engine = "^0.13.0"
+airbyte-cdk = "^5.6.0"
+duckdb = "^1.1.0"
+duckdb-engine = "^0.13.2"
 google-auth = ">=2.27.0,<3.0"
 google-cloud-bigquery = ">=3.12.0,<4.0"
 google-cloud-secret-manager = "^2.17.0"
 jsonschema = ">=3.2.0,<5.0"
 numpy = "<2.0"  # v2.0 causes "ValueError: numpy.dtype size changed" error
 orjson = "^3.10"
 overrides = "^7.4.0"
-pandas = ">=1.5.3,<3.0"
+pandas = { version = ">=1.5.3,<3.0" }
 pendulum = "<=3.0.0"
-psycopg2-binary = "^2.9.9"
-# psycopg = {extras = ["binary", "pool"], version = "^3.1.16"}
-# Psycopg3 is not supported in SQLAlchemy 1.x:
+psycopg = {extras = ["binary", "pool"], version = "^3.2.2"}
+psycopg2-binary =  "^2.9.9"
 pyarrow = ">=16.1,<18.0"
 pydantic = ">=2.0,<=3.0"
 python-dotenv = "^1.0.1"
 python-ulid = "^2.2.0"
 requests = "<=2.32.2,!=3.32.0"  # 3.32.0 breaks docker tests
 rich = "^13.7.0"
-snowflake-connector-python = "^3.10.0"
-snowflake-sqlalchemy = "^1.5.1"
-sqlalchemy = "^2.0.35"
+snowflake-connector-python = "^3.12.2"
+snowflake-sqlalchemy = "^1.6.1"
+sqlalchemy = ">=1.4.51,<3.0"
 types-pyyaml = "^6.0.12.12"
 
 # TODO: Remove this arbitrary python constraint once `sqlalchemy-bigquery` has done so.
@@ -49,11 +48,12 @@ pyiceberg = "^0.6.1"
 uuid7 = "^0.1.0"
 grpcio = "^1.65.0"
 structlog = "^24.4.0"
+airbyte-protocol-models-pdv2 = "^0.13.0"
 
 [tool.poetry.group.dev.dependencies]
-docker = "^7.0.0"
+docker = "^7.1.0"
 faker = "^21.0.0"
-mypy = "^1.7.1"
+mypy = "^1.11.2"
 pandas-stubs = "^2.1.4.231218"
 pdoc = "^14.3.0"
 pytest = "^8.2.0"
@@ -91,7 +91,7 @@ markers = [
 filterwarnings = [ # syntax: "action:message_regex:category:module:line"
     # Treat python warnings as errors in pytest
     "error",
-    # Snowflake reregisters its functions, which is fine.
+    # # Snowflake reregisters its functions, which is fine.
     "ignore::sqlalchemy.exc.SAWarning",  # "GenericFunction is already registered and is going to be overridden.""
     # Ignore these, specifically on Windows because of file cleanup and locking issues:
     "ignore:unclosed file:ResourceWarning",
@@ -217,7 +217,11 @@ airbyte = "ab"
 [tool.ruff.lint.isort]
 force-sort-within-sections = false
 lines-after-imports = 2
-known-first-party = ["airbyte_cdk", "airbyte_protocol"]
+known-first-party = [
+    "airbyte_cdk",
+    "airbyte_protocol",
+    "airbyte_protocol_dataclasses",
+]
 known-local-folder = ["airbyte"]
 required-imports = ["from __future__ import annotations"]
 known-third-party = []
@@ -297,8 +301,11 @@ exclude = [
 ]
 
 [[tool.mypy.overrides]]
-module = ["airbyte_protocol", "airbyte_protocol.models"]
-ignore_missing_imports = true                            # No stubs yet (😢)
+ignore_missing_imports = true  # No stubs yet (😢)
+module = [
+    "airbyte_protocol",
+    "airbyte_protocol.models",
+]
 
 [tool.pyright]
 pythonVersion = "3.10"

pyproject.toml

@@ -15,7 +15,7 @@
 
 import airbyte
 import docker
-import psycopg2 as psycopg
+import psycopg
 import pytest
 import ulid
 from _pytest.nodes import Item

tests/conftest.py

@@ -16,7 +16,7 @@
 from airbyte.shared.catalog_providers import CatalogProvider
 from airbyte.sources.base import Source
 from airbyte.strategies import WriteStrategy
-from airbyte_cdk import AirbyteMessage, AirbyteRecordMessage, Type
+from airbyte_protocol.models import AirbyteMessage, AirbyteRecordMessage, Type
 
 
 @pytest.fixture