diff --git a/pyproject.toml b/pyproject.toml index fb8c4e3d2..4e2c3ae25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ dynamic = [ "version" ] dependencies = [ "adlfs", "aiohttp<4", - "asana<6", + "asana>=5,<6", "click<8.5", "clickhouse-connect>=0.8,<2", "clickhouse-driver<0.3", @@ -84,7 +84,7 @@ dependencies = [ "databricks-sql-connector<5", "databricks-sqlalchemy<3", "dataclasses-json<0.7", - "dlt<1.30", + "dlt>=1.22,<1.30", "dlt-cratedb<0.2", "duckdb<1.6", "duckdb-engine<0.18", @@ -182,7 +182,7 @@ optional-dependencies.test = [ "pytest-asyncio<2", "pytest-cov<8", "pytest-xdist[psutil]<4", - "testcontainers[mysql,postgres]>=4.9.1,<4.15", + "testcontainers[azurite,mysql,postgres]>=4.9.1,<4.15", "types-requests<3", "verlib2", ] @@ -287,7 +287,9 @@ lint.per-file-ignores."src/omniload/{source,target}/*" = [ "C417", # line-too-long "E501", - # FIXME: `print` found +] +lint.per-file-ignores."src/omniload/{source,target}/**/{client,helpers}.py" = [ + # `print` found "T201", ] lint.per-file-ignores."tests/*" = [ @@ -299,7 +301,7 @@ lint.per-file-ignores."tests/*" = [ "C416", # line-too-long "E501", - # Allow use of `assert`, and `print`. + # Allow use of `assert`. "S101", # Possible SQL injection vector through string-based query construction "S608", diff --git a/src/dlt_filesystem/source/impl/remote.py b/src/dlt_filesystem/source/impl/remote.py index eae0cdf5d..afdd51829 100644 --- a/src/dlt_filesystem/source/impl/remote.py +++ b/src/dlt_filesystem/source/impl/remote.py @@ -1,8 +1,10 @@ import base64 import json -from typing import Any, Dict +from typing import Any, Dict, Type from urllib.parse import parse_qs, urlparse +from fsspec import AbstractFileSystem + from dlt_filesystem.error import InvalidBlobTableError, MissingConnectorOption from dlt_filesystem.source.base import FilesystemSource from dlt_filesystem.source.error import UnsupportedEndpointError @@ -168,7 +170,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): ) -def _azure_fs(auth: AzureBlobAuth): +def _azure_kwargs(auth: AzureBlobAuth): """Build an ``adlfs.AzureBlobFileSystem`` from resolved Azure auth params. The ingestr-style short names already match adlfs kwargs, so they pass @@ -176,7 +178,6 @@ def _azure_fs(auth: AzureBlobAuth): imported lazily so the CLI ``--help`` and every non-Azure path never load the Azure SDK (matching the s3fs/gcsfs deferred-import convention). """ - import adlfs kwargs = {"account_name": auth.account_name} if auth.account_key is not None: @@ -191,12 +192,11 @@ def _azure_fs(auth: AzureBlobAuth): kwargs["client_secret"] = auth.client_secret if auth.account_host is not None: kwargs["account_host"] = auth.account_host - - # adlfs annotates its credential params as `str` (defaulting to None) and - # mixes in non-str params (blocksize: int, ...), so ty can't check a - # conditional str-kwargs splat against the signature. The kwargs are all - # valid adlfs credential arguments by construction. - return adlfs.AzureBlobFileSystem(**kwargs) # ty: ignore[invalid-argument-type] + if auth.connection_string is not None: + kwargs["connection_string"] = auth.connection_string + if auth.api_version is not None: + kwargs["api_version"] = auth.api_version + return kwargs class AzureSource(FilesystemSource): @@ -208,6 +208,16 @@ class AzureSource(FilesystemSource): aliases onto this class. """ + @property + def fs_class(self) -> Type["AbstractFileSystem"]: + # adlfs annotates its credential params as `str` (defaulting to None) and + # mixes in non-str params (blocksize: int, ...), so ty can't check a + # conditional str-kwargs splat against the signature. The kwargs are all + # valid adlfs credential arguments by construction. + from adlfs import AzureBlobFileSystem + + return AzureBlobFileSystem + def dlt_source(self, uri: str, table: str, **kwargs): if kwargs.get("incremental_key"): raise ValueError( @@ -225,7 +235,8 @@ def dlt_source(self, uri: str, table: str, **kwargs): bucket_url = f"az://{bucket_name}" - fs = _azure_fs(auth) + kwargs.update(_azure_kwargs(auth)) + fs = self.fs_class(**kwargs) try: endpoint: str = determine_endpoint(table, path_to_file) @@ -246,7 +257,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): file_glob=path_to_file, reader_name=endpoint, storage_namespace=( - f"azure:{auth.account_name.lower()}:" + f"azure:{(auth.account_name or '').lower()}:" f"{_endpoint_namespace(auth.account_host, 'azure-public')}" ), filesystem_incremental=kwargs.get("filesystem_incremental", False), diff --git a/src/dlt_filesystem/target/remote.py b/src/dlt_filesystem/target/remote.py index 0a667803b..4ef00cd41 100644 --- a/src/dlt_filesystem/target/remote.py +++ b/src/dlt_filesystem/target/remote.py @@ -167,9 +167,9 @@ def credentials(self, params: dict) -> FileSystemCredentials: # parse_azure_blob_auth guarantees the full triplet here; the # per-field `is not None` guards mirror the account-key branch and # narrow the parsed optionals to str for ty. - sp_credentials = AzureServicePrincipalCredentials( - azure_storage_account_name=auth.account_name, - ) + sp_credentials = AzureServicePrincipalCredentials() + if auth.account_name is not None: + sp_credentials.azure_storage_account_name = auth.account_name if auth.tenant_id is not None: sp_credentials.azure_tenant_id = auth.tenant_id if auth.client_id is not None: @@ -180,9 +180,9 @@ def credentials(self, params: dict) -> FileSystemCredentials: sp_credentials.azure_account_host = auth.account_host return sp_credentials - key_credentials = AzureCredentials( - azure_storage_account_name=auth.account_name, - ) + key_credentials = AzureCredentials() + if auth.account_name is not None: + key_credentials.azure_storage_account_name = auth.account_name if auth.account_key is not None: key_credentials.azure_storage_account_key = auth.account_key if auth.sas_token is not None: diff --git a/src/dlt_filesystem/util/auth.py b/src/dlt_filesystem/util/auth.py index f1817343e..fd435d93e 100644 --- a/src/dlt_filesystem/util/auth.py +++ b/src/dlt_filesystem/util/auth.py @@ -18,13 +18,15 @@ class AzureBlobAuth: spec fields. """ - account_name: str + account_name: Optional[str] = None account_key: Optional[str] = None sas_token: Optional[str] = None tenant_id: Optional[str] = None client_id: Optional[str] = None client_secret: Optional[str] = None account_host: Optional[str] = None + connection_string: Optional[str] = None + api_version: Optional[str] = None @property def is_service_principal(self) -> bool: @@ -67,15 +69,24 @@ def parse_azure_blob_auth(params: dict) -> AzureBlobAuth: def one(key: str) -> Optional[str]: return params.get(key, [None])[0] - account_name = one("account_name") - if account_name is None: - raise MissingConnectorOption("account_name", "Azure") + connection_string = one("connection_string") + api_version = one("api_version") + + if connection_string is not None: + return AzureBlobAuth( + connection_string=connection_string, + api_version=api_version, + ) + account_name = one("account_name") account_key = one("account_key") sas_token = one("sas_token") sp_values = {field: one(field) for field in AZURE_SERVICE_PRINCIPAL_FIELDS} account_host = one("account_host") + if account_name is None: + raise MissingConnectorOption("account_name", "Azure") + if account_key is not None and sas_token is not None: raise ValueError( "Conflicting Azure credentials: supply either account_key or " @@ -109,5 +120,7 @@ def one(key: str) -> Optional[str]: account_key=account_key, sas_token=sas_token, account_host=account_host, + connection_string=connection_string, + api_version=api_version, **sp_values, ) diff --git a/src/omniload/source/docebo/adapter.py b/src/omniload/source/docebo/adapter.py index a3555706b..bb21e4582 100644 --- a/src/omniload/source/docebo/adapter.py +++ b/src/omniload/source/docebo/adapter.py @@ -122,7 +122,6 @@ def users() -> Iterator[Dict[str, Any]]: }, ) def courses() -> Iterator[Dict[str, Any]]: - print("running courses transformer") """Fetch all courses from Docebo.""" for courses_batch in client.fetch_courses(page_size=1000): for course in courses_batch: diff --git a/src/omniload/source/github/adapter.py b/src/omniload/source/github/adapter.py index 042406d4b..89701c7e5 100644 --- a/src/omniload/source/github/adapter.py +++ b/src/omniload/source/github/adapter.py @@ -14,6 +14,7 @@ """Source that load github issues, pull requests and reactions for a specific repository via customizable graphql query. Loads events incrementally.""" +import logging import urllib.parse from typing import Iterator, Optional, Sequence, cast @@ -26,6 +27,8 @@ from .helpers import get_reactions_data, get_rest_pages, get_stargazers +logger = logging.getLogger(__name__) + @dlt.source(max_table_nesting=0) def github_reactions( @@ -163,7 +166,7 @@ def repo_events( # stop requesting pages if the last element was already older than initial value # note: incremental will skip those items anyway, we just do not want to use the api limits if last_created_at.start_out_of_range: - print( + logger.info( f"Overlap with previous run created at {last_created_at.initial_value}" ) break diff --git a/tests/dlt_filesystem/test_source_incremental.py b/tests/dlt_filesystem/test_source_incremental.py index 3a0a4290e..6d327be9f 100644 --- a/tests/dlt_filesystem/test_source_incremental.py +++ b/tests/dlt_filesystem/test_source_incremental.py @@ -153,7 +153,7 @@ def test_filesystem_sources_thread_incremental_identity_without_auth_material(tm with ( patch("dlt_filesystem.source.adapter.resource_for_reader") as build, - patch("dlt_filesystem.source.impl.remote._azure_fs"), + patch("adlfs.AzureBlobFileSystem"), ): AzureSource().dlt_source( "az://?account_name=account&account_key=SECRET", @@ -254,7 +254,7 @@ def test_auth_rotation_does_not_change_incremental_resource_names(): with ( patch("dlt_filesystem.source.adapter.resource_for_reader") as first_build, - patch("dlt_filesystem.source.impl.remote._azure_fs"), + patch("adlfs.AzureBlobFileSystem"), ): AzureSource().dlt_source( "az://?account_name=account&account_key=OLD_SECRET", @@ -262,7 +262,7 @@ def test_auth_rotation_does_not_change_incremental_resource_names(): ) with ( patch("dlt_filesystem.source.adapter.resource_for_reader") as second_build, - patch("dlt_filesystem.source.impl.remote._azure_fs"), + patch("adlfs.AzureBlobFileSystem"), ): AzureSource().dlt_source( "az://?account_name=account&account_key=NEW_SECRET", diff --git a/tests/main/filesystem/test_remote_integration.py b/tests/main/filesystem/test_remote_integration.py new file mode 100644 index 000000000..da3839d6d --- /dev/null +++ b/tests/main/filesystem/test_remote_integration.py @@ -0,0 +1,94 @@ +from pathlib import Path +from urllib.parse import quote + +import pytest +from azure.storage.blob import BlobServiceClient +from testcontainers.azurite import AzuriteContainer + +from omniload import run_ingest +from tests.util.container.impl.floci import FlociContainer +from tests.warehouse.manager import FLOCI_IMAGE + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="session") +def floci(): + """Provide S3 service container for the whole test session.""" + container = FlociContainer(image=FLOCI_IMAGE) + container.start() + s3 = container.get_client("s3") # ty: ignore[invalid-argument-type, missing-argument] + s3.create_bucket(Bucket="test-bucket") + s3.upload_file( + Filename="tests/assets/create_replace.csv", + Bucket="test-bucket", + Key="path/to/create_replace.csv", + ) + yield container + container.stop() + + +@pytest.fixture(scope="session") +def azurite(): + """Provide Azure Storage service container for the whole test session.""" + container = AzuriteContainer() + container.start() + connection_string = container.get_connection_string() + client = BlobServiceClient.from_connection_string(connection_string) + client.create_container("test-container") + cc = client.get_container_client("test-container") + cc.upload_blob( + name="path/to/create_replace.csv", + data=Path("tests/assets/create_replace.csv").read_text(), + ) + yield container + container.stop() + + +def duckdb_table_cardinality(db_path: Path, table_name: str) -> int: + """Return number of records in database table.""" + import duckdb + + db = duckdb.connect(db_path) + result = db.execute(f"SELECT * FROM {table_name}").fetchall() + count = len(result) + db.close() + return count + + +def test_s3_source(floci, tmp_path): + """S3 source integration test.""" + s3_endpoint = floci.get_url() + db_path = tmp_path / "db.duckdb" + result = run_ingest( + source_uri=f"s3://?endpoint_url={s3_endpoint}&access_key_id=test&secret_access_key=test", + dest_uri=f"duckdb:///{db_path}", + source_table="test-bucket/path/to/create_replace.csv", + dest_table="testdrive.data", + ) + if result is None: + raise RuntimeError("Ingest failed") + package = result.asdict()["load_packages"][0] + assert package["state"] == "loaded" + + count = duckdb_table_cardinality(db_path, "testdrive.data") + assert count == 20, f"Wrong number of records: {count}. Expected: 20" + + +def test_azure_source(azurite, tmp_path): + """Azure Storage source integration test.""" + connection_string = azurite.get_connection_string() + db_path = tmp_path / "db.duckdb" + result = run_ingest( + source_uri=f"az://?connection_string={quote(connection_string)}&api_version=2025-11-05", + dest_uri=f"duckdb:///{db_path}", + source_table="test-container/path/to/create_replace.csv", + dest_table="testdrive.data", + ) + if result is None: + raise RuntimeError("Ingest failed") + package = result.asdict()["load_packages"][0] + assert package["state"] == "loaded" + + count = duckdb_table_cardinality(db_path, "testdrive.data") + assert count == 20, f"Wrong number of records: {count}. Expected: 20"