|
1 | | -import base64 |
2 | | -import json |
3 | 1 | from abc import abstractmethod |
4 | 2 | from typing import TYPE_CHECKING, Any, Dict, Type |
5 | 3 | from urllib.parse import parse_qs, urlparse |
|
16 | 14 | parse_uri, |
17 | 15 | source_selects_single_file, |
18 | 16 | ) |
19 | | -from dlt_filesystem.util.auth import AzureBlobAuth, parse_azure_blob_auth |
| 17 | +from dlt_filesystem.util.auth import ( |
| 18 | + azure_blob_filesystem_kwargs, |
| 19 | + gcs_filesystem_kwargs, |
| 20 | + parse_azure_blob_auth, |
| 21 | + s3_filesystem_kwargs, |
| 22 | +) |
20 | 23 |
|
21 | 24 | if TYPE_CHECKING: |
22 | 25 | from fsspec import AbstractFileSystem |
@@ -52,25 +55,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): |
52 | 55 |
|
53 | 56 | bucket_url = f"gs://{bucket_name}" |
54 | 57 |
|
55 | | - credentials_path = params.pop("credentials_path", [None])[0] |
56 | | - credentials_base64 = params.pop("credentials_base64", [None])[0] |
57 | | - |
58 | | - # Merge params into fs kwargs, without overriding kwargs already |
59 | | - # supplied by the caller (e.g. filesystem_incremental, column_types). |
60 | | - for key, value in params.items(): |
61 | | - kwargs.setdefault(key, value[0]) |
62 | | - |
63 | | - if "token" not in kwargs: |
64 | | - credentials = None |
65 | | - if credentials_path: |
66 | | - credentials = credentials_path |
67 | | - elif credentials_base64: |
68 | | - credentials = json.loads(base64.b64decode(credentials_base64).decode()) |
69 | | - else: |
70 | | - credentials = "anon" |
71 | | - kwargs["token"] = credentials |
72 | | - |
73 | | - fs = self.fs_class(**kwargs) |
| 58 | + fs = self.fs_class(**gcs_filesystem_kwargs(params, kwargs)) |
74 | 59 |
|
75 | 60 | try: |
76 | 61 | endpoint: str = determine_endpoint(table, path_to_file) |
@@ -133,27 +118,14 @@ def dlt_source(self, uri: str, table: str, **kwargs): |
133 | 118 |
|
134 | 119 | parsed_uri = urlparse(uri) |
135 | 120 | source_fields = parse_qs(parsed_uri.query) |
136 | | - access_key_id = source_fields.get("access_key_id") |
137 | | - if not access_key_id: |
138 | | - raise MissingConnectorOption("access_key_id", self.fs_name) |
139 | | - |
140 | | - secret_access_key = source_fields.get("secret_access_key") |
141 | | - if not secret_access_key: |
142 | | - raise MissingConnectorOption("secret_access_key", self.fs_name) |
143 | | - |
| 121 | + fs_kwargs = s3_filesystem_kwargs(source_fields, self.fs_name) |
144 | 122 | bucket_name, path_to_file = parse_uri(parsed_uri, table) |
145 | 123 | if not bucket_name or not path_to_file: |
146 | 124 | raise InvalidBlobTableError(self.fs_name) |
147 | 125 |
|
148 | 126 | bucket_url = f"{self.fs_protocol}://{bucket_name}/" |
149 | 127 |
|
150 | 128 | endpoint_url = source_fields.get("endpoint_url") |
151 | | - fs_kwargs: dict = { |
152 | | - "key": access_key_id[0], |
153 | | - "secret": secret_access_key[0], |
154 | | - } |
155 | | - if endpoint_url: |
156 | | - fs_kwargs["endpoint_url"] = endpoint_url[0] |
157 | 129 |
|
158 | 130 | fs = self.fs_class(**fs_kwargs) |
159 | 131 |
|
@@ -190,35 +162,6 @@ def fs_name(self) -> str: |
190 | 162 | return "S3" |
191 | 163 |
|
192 | 164 |
|
193 | | -def _azure_kwargs(auth: AzureBlobAuth): |
194 | | - """Return AzureBlobAuth information as dictionary. |
195 | | -
|
196 | | - The ingestr-style short names already match adlfs kwargs, so they pass |
197 | | - straight through; only the supplied ones are forwarded. ``adlfs`` is |
198 | | - imported lazily so the CLI ``--help`` and every non-Azure path never load |
199 | | - the Azure SDK (matching the s3fs/gcsfs deferred-import convention). |
200 | | - """ |
201 | | - |
202 | | - kwargs = {"account_name": auth.account_name} |
203 | | - if auth.account_key is not None: |
204 | | - kwargs["account_key"] = auth.account_key |
205 | | - if auth.sas_token is not None: |
206 | | - kwargs["sas_token"] = auth.sas_token |
207 | | - if auth.tenant_id is not None: |
208 | | - kwargs["tenant_id"] = auth.tenant_id |
209 | | - if auth.client_id is not None: |
210 | | - kwargs["client_id"] = auth.client_id |
211 | | - if auth.client_secret is not None: |
212 | | - kwargs["client_secret"] = auth.client_secret |
213 | | - if auth.account_host is not None: |
214 | | - kwargs["account_host"] = auth.account_host |
215 | | - if auth.connection_string is not None: |
216 | | - kwargs["connection_string"] = auth.connection_string |
217 | | - if auth.api_version is not None: |
218 | | - kwargs["api_version"] = auth.api_version |
219 | | - return kwargs |
220 | | - |
221 | | - |
222 | 165 | class AzureSource(FilesystemSource): |
223 | 166 | """Azure Blob Storage / ADLS Gen2 source (``az://``, ``adls://``, ``abfss://``). |
224 | 167 |
|
@@ -256,7 +199,7 @@ def dlt_source(self, uri: str, table: str, **kwargs): |
256 | 199 |
|
257 | 200 | bucket_url = f"az://{bucket_name}" |
258 | 201 |
|
259 | | - kwargs.update(_azure_kwargs(auth)) |
| 202 | + kwargs.update(azure_blob_filesystem_kwargs(auth)) |
260 | 203 | fs = self.fs_class(**kwargs) |
261 | 204 |
|
262 | 205 | try: |
|
0 commit comments