dbt-polars 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- dbt_polars-0.0.1/PKG-INFO +87 -0
- dbt_polars-0.0.1/README.md +55 -0
- dbt_polars-0.0.1/dbt/__init__.py +3 -0
- dbt_polars-0.0.1/dbt/adapters/polars/__init__.py +84 -0
- dbt_polars-0.0.1/dbt/adapters/polars/__version__.py +1 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/__init__.py +24 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/azureCatalog.py +291 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/baseCatalog.py +160 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/formats/__init__.py +33 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/formats/delta.py +222 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/formats/file.py +193 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/icebergCatalog.py +486 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/iceberg_write_utils.py +169 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/localCatalog.py +138 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/s3Catalog.py +229 -0
- dbt_polars-0.0.1/dbt/adapters/polars/catalogs/storageCatalog.py +292 -0
- dbt_polars-0.0.1/dbt/adapters/polars/connections.py +157 -0
- dbt_polars-0.0.1/dbt/adapters/polars/impl.py +1326 -0
- dbt_polars-0.0.1/dbt/adapters/polars/relation.py +42 -0
- dbt_polars-0.0.1/dbt/adapters/polars/sql_rewrite.py +158 -0
- dbt_polars-0.0.1/dbt/include/polars/__init__.py +3 -0
- dbt_polars-0.0.1/dbt_polars.egg-info/PKG-INFO +87 -0
- dbt_polars-0.0.1/dbt_polars.egg-info/SOURCES.txt +26 -0
- dbt_polars-0.0.1/dbt_polars.egg-info/dependency_links.txt +1 -0
- dbt_polars-0.0.1/dbt_polars.egg-info/requires.txt +15 -0
- dbt_polars-0.0.1/dbt_polars.egg-info/top_level.txt +1 -0
- dbt_polars-0.0.1/pyproject.toml +102 -0
- dbt_polars-0.0.1/setup.cfg +4 -0
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dbt-polars
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: A dbt adapter for running models with Polars
|
|
5
|
+
Author-email: Jonas Crevecoeur <jonas.crevecoeur@dataminded.com>
|
|
6
|
+
Maintainer-email: Jonas Crevecoeur <jonas.crevecoeur@dataminded.com>
|
|
7
|
+
License: Apache-2.0
|
|
8
|
+
Keywords: dbt,polars,adapter,data
|
|
9
|
+
Classifier: Development Status :: 3 - Alpha
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Classifier: Topic :: Database
|
|
18
|
+
Requires-Python: <3.14,>=3.10
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
Requires-Dist: dbt-adapters>=1.0.0
|
|
21
|
+
Requires-Dist: dbt-common>=1.0.0
|
|
22
|
+
Requires-Dist: dbt-core>=1.11.2
|
|
23
|
+
Requires-Dist: polars[deltalake,pyarrow]!=1.43.0,>=1.0.0
|
|
24
|
+
Requires-Dist: sqlglot>=1.0.0
|
|
25
|
+
Provides-Extra: iceberg
|
|
26
|
+
Requires-Dist: pyiceberg[adlfs,sql-sqlite]<=0.11.1,>=0.11.0; extra == "iceberg"
|
|
27
|
+
Provides-Extra: s3
|
|
28
|
+
Requires-Dist: boto3>=1.20.0; extra == "s3"
|
|
29
|
+
Provides-Extra: azure
|
|
30
|
+
Requires-Dist: azure-storage-file-datalake>=12.0.0; extra == "azure"
|
|
31
|
+
Requires-Dist: azure-identity>=1.0.0; extra == "azure"
|
|
32
|
+
|
|
33
|
+
# dbt-polars
|
|
34
|
+
|
|
35
|
+
A [dbt](https://www.getdbt.com/) adapter that runs models locally using [Polars](https://pola.rs/). Instead of connecting to a database, dbt-polars reads and writes files — on your local filesystem, Azure Blob Storage, AWS S3, or an Iceberg catalog.
|
|
36
|
+
|
|
37
|
+
## Installation
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
pip install dbt-polars
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
Extra dependencies are required for cloud and Iceberg backends:
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install 'dbt-polars[azure]' # Azure Blob Storage
|
|
47
|
+
pip install 'dbt-polars[s3]' # AWS S3
|
|
48
|
+
pip install 'dbt-polars[iceberg]' # Apache Iceberg (SQLite, REST, Databricks, …)
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
## Catalogs
|
|
52
|
+
|
|
53
|
+
dbt-polars uses **catalogs** to define where data is stored. Each catalog maps to a storage backend. You can configure multiple catalogs in a single profile and reference them from your models.
|
|
54
|
+
|
|
55
|
+
| Catalog | Backend | Status |
|
|
56
|
+
|---|---|---|
|
|
57
|
+
| `local` | Local filesystem | Usable |
|
|
58
|
+
| `iceberg` with SQLite | Iceberg + SQLite metastore | Usable |
|
|
59
|
+
| `iceberg` with REST + Databricks | Iceberg REST API (Unity Catalog) | Experimental |
|
|
60
|
+
| `iceberg` with other backends | Any pyiceberg-supported backend | Experimental |
|
|
61
|
+
| `azure` | Azure Blob Storage / ADLS Gen2 | Experimental |
|
|
62
|
+
|
|
63
|
+
## Quick start
|
|
64
|
+
|
|
65
|
+
A minimal `profiles.yml` for local development:
|
|
66
|
+
|
|
67
|
+
```yaml
|
|
68
|
+
my_project:
|
|
69
|
+
target: dev
|
|
70
|
+
outputs:
|
|
71
|
+
dev:
|
|
72
|
+
type: polars
|
|
73
|
+
schema: dev
|
|
74
|
+
catalogs:
|
|
75
|
+
- name: my_catalog
|
|
76
|
+
type: local
|
|
77
|
+
root: ./data
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
## Documentation
|
|
81
|
+
|
|
82
|
+
- [Catalog overview and profile reference](docs/index.md)
|
|
83
|
+
- [Local catalog](docs/catalogs/local.md)
|
|
84
|
+
- [Azure catalog](docs/catalogs/azure.md)
|
|
85
|
+
- [Iceberg catalog](docs/catalogs/iceberg.md)
|
|
86
|
+
- [Python models and tests](docs/python-models.md)
|
|
87
|
+
- [Roadmap and known limitations](docs/roadmap.md)
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
# dbt-polars
|
|
2
|
+
|
|
3
|
+
A [dbt](https://www.getdbt.com/) adapter that runs models locally using [Polars](https://pola.rs/). Instead of connecting to a database, dbt-polars reads and writes files — on your local filesystem, Azure Blob Storage, AWS S3, or an Iceberg catalog.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install dbt-polars
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
Extra dependencies are required for cloud and Iceberg backends:
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
pip install 'dbt-polars[azure]' # Azure Blob Storage
|
|
15
|
+
pip install 'dbt-polars[s3]' # AWS S3
|
|
16
|
+
pip install 'dbt-polars[iceberg]' # Apache Iceberg (SQLite, REST, Databricks, …)
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
## Catalogs
|
|
20
|
+
|
|
21
|
+
dbt-polars uses **catalogs** to define where data is stored. Each catalog maps to a storage backend. You can configure multiple catalogs in a single profile and reference them from your models.
|
|
22
|
+
|
|
23
|
+
| Catalog | Backend | Status |
|
|
24
|
+
|---|---|---|
|
|
25
|
+
| `local` | Local filesystem | Usable |
|
|
26
|
+
| `iceberg` with SQLite | Iceberg + SQLite metastore | Usable |
|
|
27
|
+
| `iceberg` with REST + Databricks | Iceberg REST API (Unity Catalog) | Experimental |
|
|
28
|
+
| `iceberg` with other backends | Any pyiceberg-supported backend | Experimental |
|
|
29
|
+
| `azure` | Azure Blob Storage / ADLS Gen2 | Experimental |
|
|
30
|
+
|
|
31
|
+
## Quick start
|
|
32
|
+
|
|
33
|
+
A minimal `profiles.yml` for local development:
|
|
34
|
+
|
|
35
|
+
```yaml
|
|
36
|
+
my_project:
|
|
37
|
+
target: dev
|
|
38
|
+
outputs:
|
|
39
|
+
dev:
|
|
40
|
+
type: polars
|
|
41
|
+
schema: dev
|
|
42
|
+
catalogs:
|
|
43
|
+
- name: my_catalog
|
|
44
|
+
type: local
|
|
45
|
+
root: ./data
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
## Documentation
|
|
49
|
+
|
|
50
|
+
- [Catalog overview and profile reference](docs/index.md)
|
|
51
|
+
- [Local catalog](docs/catalogs/local.md)
|
|
52
|
+
- [Azure catalog](docs/catalogs/azure.md)
|
|
53
|
+
- [Iceberg catalog](docs/catalogs/iceberg.md)
|
|
54
|
+
- [Python models and tests](docs/python-models.md)
|
|
55
|
+
- [Roadmap and known limitations](docs/roadmap.md)
|
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
from dbt.adapters.base.plugin import AdapterPlugin
|
|
2
|
+
from dbt.exceptions import ModelConfigError
|
|
3
|
+
from dbt_common.dataclass_schema import ValidationError
|
|
4
|
+
from dbt_common.exceptions import DbtRuntimeError
|
|
5
|
+
|
|
6
|
+
from dbt.adapters.polars.connections import PolarsCredentials
|
|
7
|
+
from dbt.adapters.polars.impl import PolarsAdapter
|
|
8
|
+
from dbt.include.polars import PACKAGE_PATH
|
|
9
|
+
|
|
10
|
+
# Two dbt-core monkeypatches to support Python singular tests (dbt-core 1.11.x):
|
|
11
|
+
#
|
|
12
|
+
# 1. get_file_types_for_project: adds ".py" to ParseFileType.SingularTest extensions
|
|
13
|
+
# so Python test files are discovered by the parser.
|
|
14
|
+
#
|
|
15
|
+
# 2. SingularTestParser.render_update: runs parse_python_model for .py test nodes so
|
|
16
|
+
# dbt.ref() calls are statically scanned and added to depends_on. Without this, dbt's
|
|
17
|
+
# runtime ref() validation rejects any dbt.ref() call in a .py test. Requires a
|
|
18
|
+
# "def model" entry point (satisfied by _TEST_ENTRYPOINT_WRAPPER below).
|
|
19
|
+
_TEST_ENTRYPOINT_WRAPPER = "\n\n\ndef model(dbt, pl):\n return test(dbt, pl)\n"
|
|
20
|
+
|
|
21
|
+
try:
|
|
22
|
+
import dbt.parser.read_files as _read_files
|
|
23
|
+
from dbt.node_types import ModelLanguage
|
|
24
|
+
from dbt.parser.models import ModelParser, verify_python_model_code
|
|
25
|
+
from dbt.parser.singular_test import SingularTestParser
|
|
26
|
+
|
|
27
|
+
_original_get_file_types_for_project = _read_files.get_file_types_for_project
|
|
28
|
+
|
|
29
|
+
def _get_file_types_for_project_with_py_tests(project):
|
|
30
|
+
file_types = _original_get_file_types_for_project(project)
|
|
31
|
+
try:
|
|
32
|
+
file_types[_read_files.ParseFileType.SingularTest]["extensions"].append(
|
|
33
|
+
".py"
|
|
34
|
+
)
|
|
35
|
+
except (KeyError, AttributeError, TypeError):
|
|
36
|
+
pass
|
|
37
|
+
return file_types
|
|
38
|
+
|
|
39
|
+
_read_files.get_file_types_for_project = _get_file_types_for_project_with_py_tests
|
|
40
|
+
|
|
41
|
+
_original_singular_test_render_update = SingularTestParser.render_update
|
|
42
|
+
|
|
43
|
+
def _singular_test_render_update(
|
|
44
|
+
self, node, config, validate_config_call_dict=False
|
|
45
|
+
):
|
|
46
|
+
if node.language == ModelLanguage.python:
|
|
47
|
+
if "def test(" not in node.raw_code:
|
|
48
|
+
raise DbtRuntimeError(
|
|
49
|
+
f"Python singular test '{node.name}' must define "
|
|
50
|
+
+ "a function named 'test(dbt, pl)'"
|
|
51
|
+
)
|
|
52
|
+
if not node.raw_code.endswith(_TEST_ENTRYPOINT_WRAPPER):
|
|
53
|
+
node.raw_code += _TEST_ENTRYPOINT_WRAPPER
|
|
54
|
+
try:
|
|
55
|
+
verify_python_model_code(node)
|
|
56
|
+
context = self._context_for(node, config)
|
|
57
|
+
ModelParser.parse_python_model(self, node, config, context)
|
|
58
|
+
self.update_parsed_node_config(
|
|
59
|
+
node, config, context=context, validate_config_call_dict=True
|
|
60
|
+
)
|
|
61
|
+
except ValidationError as exc:
|
|
62
|
+
raise ModelConfigError(exc, node=node) from exc
|
|
63
|
+
except Exception:
|
|
64
|
+
return _original_singular_test_render_update(
|
|
65
|
+
self, node, config, validate_config_call_dict
|
|
66
|
+
)
|
|
67
|
+
return
|
|
68
|
+
return _original_singular_test_render_update(
|
|
69
|
+
self, node, config, validate_config_call_dict
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
SingularTestParser.render_update = _singular_test_render_update # type: ignore[method-assign]
|
|
73
|
+
except (ImportError, AttributeError):
|
|
74
|
+
from dbt.adapters.events.logging import AdapterLogger
|
|
75
|
+
|
|
76
|
+
AdapterLogger("polars").warning(
|
|
77
|
+
"Python singular tests are not supported with this version of dbt-core."
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
Plugin = AdapterPlugin(
|
|
81
|
+
adapter=PolarsAdapter, # type: ignore[arg-type]
|
|
82
|
+
credentials=PolarsCredentials,
|
|
83
|
+
include_path=PACKAGE_PATH,
|
|
84
|
+
)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
version = "1.7.0"
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
from dbt.adapters.polars.catalogs.azureCatalog import (
|
|
2
|
+
AzureBlobStorageCatalog,
|
|
3
|
+
AzureBlobStorageCatalogConfig,
|
|
4
|
+
)
|
|
5
|
+
from dbt.adapters.polars.catalogs.baseCatalog import BaseCatalog, CatalogConfig
|
|
6
|
+
from dbt.adapters.polars.catalogs.icebergCatalog import (
|
|
7
|
+
IcebergCatalog,
|
|
8
|
+
IcebergCatalogConfig,
|
|
9
|
+
)
|
|
10
|
+
from dbt.adapters.polars.catalogs.localCatalog import LocalCatalog, LocalCatalogConfig
|
|
11
|
+
from dbt.adapters.polars.catalogs.s3Catalog import AWSS3Catalog, S3CatalogConfig
|
|
12
|
+
|
|
13
|
+
CATALOG_CONFIG_REGISTRY: dict[str, type[CatalogConfig]] = {
|
|
14
|
+
"local": LocalCatalogConfig,
|
|
15
|
+
"iceberg": IcebergCatalogConfig,
|
|
16
|
+
"azure": AzureBlobStorageCatalogConfig,
|
|
17
|
+
"s3": S3CatalogConfig,
|
|
18
|
+
}
|
|
19
|
+
CATALOG_REGISTRY: dict[str, type[BaseCatalog]] = {
|
|
20
|
+
"local": LocalCatalog,
|
|
21
|
+
"iceberg": IcebergCatalog,
|
|
22
|
+
"azure": AzureBlobStorageCatalog,
|
|
23
|
+
"s3": AWSS3Catalog,
|
|
24
|
+
}
|
|
@@ -0,0 +1,291 @@
|
|
|
1
|
+
from typing import TYPE_CHECKING
|
|
2
|
+
|
|
3
|
+
from dbt.adapters.contracts.relation import RelationType
|
|
4
|
+
from dbt.adapters.events.logging import AdapterLogger
|
|
5
|
+
from dbt_common.exceptions import DbtRuntimeError
|
|
6
|
+
from deltalake import DeltaTable
|
|
7
|
+
|
|
8
|
+
from dbt.adapters.polars.catalogs.baseCatalog import CatalogConfig
|
|
9
|
+
from dbt.adapters.polars.catalogs.formats import FILE_FORMATS
|
|
10
|
+
from dbt.adapters.polars.catalogs.storageCatalog import StorageCatalog
|
|
11
|
+
from dbt.adapters.polars.relation import PolarsRelation
|
|
12
|
+
|
|
13
|
+
if TYPE_CHECKING:
|
|
14
|
+
from azure.core.credentials import AccessToken
|
|
15
|
+
|
|
16
|
+
logger = AdapterLogger("polars")
|
|
17
|
+
|
|
18
|
+
_FILE_FORMATS = frozenset(FILE_FORMATS)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class AzureBlobStorageCatalogConfig(CatalogConfig):
|
|
22
|
+
def __init__(
|
|
23
|
+
self,
|
|
24
|
+
name: str,
|
|
25
|
+
type: str,
|
|
26
|
+
account_name: str,
|
|
27
|
+
container: str = "",
|
|
28
|
+
prefix: str = "",
|
|
29
|
+
credentials: dict | None = None,
|
|
30
|
+
) -> None:
|
|
31
|
+
self.name = name
|
|
32
|
+
self.type = type
|
|
33
|
+
self.account_name = account_name
|
|
34
|
+
self.container = container
|
|
35
|
+
self.prefix = prefix.strip("/")
|
|
36
|
+
self.credentials = credentials or {}
|
|
37
|
+
|
|
38
|
+
def unique_field(self) -> str:
|
|
39
|
+
return f"{self.account_name}/{self.container}/{self.prefix}"
|
|
40
|
+
|
|
41
|
+
def connection_keys(self) -> tuple[str, ...]:
|
|
42
|
+
return ("name", "account_name", "container", "prefix")
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class _StaticTokenCredential:
|
|
46
|
+
def __init__(self, token: str) -> None:
|
|
47
|
+
self._token = token
|
|
48
|
+
|
|
49
|
+
def get_token(self, *_args, **_kwargs):
|
|
50
|
+
import time
|
|
51
|
+
|
|
52
|
+
from azure.core.credentials import AccessToken
|
|
53
|
+
|
|
54
|
+
return AccessToken(self._token, int(time.time()) + 3600)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class AzureBlobStorageCatalog(StorageCatalog):
|
|
58
|
+
config: AzureBlobStorageCatalogConfig
|
|
59
|
+
|
|
60
|
+
def __init__(
|
|
61
|
+
self, config: AzureBlobStorageCatalogConfig, project_root: str
|
|
62
|
+
) -> None:
|
|
63
|
+
try:
|
|
64
|
+
import azure.identity # noqa: F401
|
|
65
|
+
import azure.storage.filedatalake # noqa: F401
|
|
66
|
+
except ImportError as exc:
|
|
67
|
+
raise DbtRuntimeError(
|
|
68
|
+
"The azure extra is required for AzureBlobStorageCatalog. "
|
|
69
|
+
"Install it with: pip install 'dbt-polars[azure]'"
|
|
70
|
+
) from exc
|
|
71
|
+
super().__init__(config, project_root)
|
|
72
|
+
self._cached_token: AccessToken | None = None
|
|
73
|
+
self._service_client = None # DataLakeServiceClient | None
|
|
74
|
+
self._dac = None # DefaultAzureCredential | None
|
|
75
|
+
|
|
76
|
+
# ── URI and credential helpers ────────────────────────────────────────────
|
|
77
|
+
|
|
78
|
+
def _blob_prefix(self, *parts: str) -> str:
|
|
79
|
+
"""Join the catalog prefix with additional path parts."""
|
|
80
|
+
components = [self.config.prefix] if self.config.prefix else []
|
|
81
|
+
components.extend(p.strip("/") for p in parts if p)
|
|
82
|
+
return "/".join(components)
|
|
83
|
+
|
|
84
|
+
def _object_path(self, schema: str, *parts: str) -> str:
|
|
85
|
+
return self._blob_prefix(schema, *parts)
|
|
86
|
+
|
|
87
|
+
def _dfs_url(self) -> str:
|
|
88
|
+
return f"https://{self.config.account_name}.dfs.core.windows.net"
|
|
89
|
+
|
|
90
|
+
def _get_uri(self, relation: PolarsRelation) -> str:
|
|
91
|
+
schema = relation.schema or ""
|
|
92
|
+
identifier = relation.identifier or ""
|
|
93
|
+
container = self.config.container
|
|
94
|
+
path = self._object_path(schema, identifier)
|
|
95
|
+
if relation.file_format != "delta":
|
|
96
|
+
path = f"{path}.{relation.file_format}"
|
|
97
|
+
return f"az://{container}/{path}"
|
|
98
|
+
|
|
99
|
+
_DELEGATED_CRED_KEYS = ("bearer_token", "account_key", "sas_token")
|
|
100
|
+
|
|
101
|
+
def _get_storage_options(self, _uri: str) -> dict[str, str]:
|
|
102
|
+
base: dict[str, str] = {
|
|
103
|
+
"account_name": self.config.account_name,
|
|
104
|
+
"timeout": "120s",
|
|
105
|
+
}
|
|
106
|
+
creds = self.config.credentials
|
|
107
|
+
|
|
108
|
+
override = creds.get("storage_options")
|
|
109
|
+
if override is not None:
|
|
110
|
+
return {**base, **override}
|
|
111
|
+
|
|
112
|
+
direct = {k: str(v) for k, v in creds.items() if k in self._DELEGATED_CRED_KEYS}
|
|
113
|
+
if direct:
|
|
114
|
+
return {**base, **direct}
|
|
115
|
+
|
|
116
|
+
base["bearer_token"] = self._default_chain_token()
|
|
117
|
+
return base
|
|
118
|
+
|
|
119
|
+
def _get_dac(self):
|
|
120
|
+
if self._dac is None:
|
|
121
|
+
from azure.identity import DefaultAzureCredential
|
|
122
|
+
|
|
123
|
+
dac_kwargs = {
|
|
124
|
+
k: v
|
|
125
|
+
for k, v in self.config.credentials.items()
|
|
126
|
+
if k != "storage_options"
|
|
127
|
+
}
|
|
128
|
+
self._dac = DefaultAzureCredential(**dac_kwargs)
|
|
129
|
+
return self._dac
|
|
130
|
+
|
|
131
|
+
def _default_chain_token(self) -> str:
|
|
132
|
+
import time
|
|
133
|
+
|
|
134
|
+
if (
|
|
135
|
+
self._cached_token is None
|
|
136
|
+
or self._cached_token.expires_on < time.time() + 300
|
|
137
|
+
):
|
|
138
|
+
self._cached_token = self._get_dac().get_token(
|
|
139
|
+
"https://storage.azure.com/.default"
|
|
140
|
+
)
|
|
141
|
+
return self._cached_token.token
|
|
142
|
+
|
|
143
|
+
def _get_service_client(self):
|
|
144
|
+
if self._service_client is not None:
|
|
145
|
+
return self._service_client
|
|
146
|
+
|
|
147
|
+
from azure.storage.filedatalake import DataLakeServiceClient
|
|
148
|
+
|
|
149
|
+
url = self._dfs_url()
|
|
150
|
+
creds = self.config.credentials
|
|
151
|
+
if "bearer_token" in creds:
|
|
152
|
+
service = DataLakeServiceClient(
|
|
153
|
+
url, credential=_StaticTokenCredential(creds["bearer_token"])
|
|
154
|
+
)
|
|
155
|
+
elif "account_key" in creds:
|
|
156
|
+
service = DataLakeServiceClient(url, credential=creds["account_key"])
|
|
157
|
+
elif "sas_token" in creds:
|
|
158
|
+
service = DataLakeServiceClient(f"{url}?{creds['sas_token']}")
|
|
159
|
+
else:
|
|
160
|
+
service = DataLakeServiceClient(url, credential=self._get_dac())
|
|
161
|
+
self._service_client = service
|
|
162
|
+
return self._service_client
|
|
163
|
+
|
|
164
|
+
def _get_file_system_client(self, file_system: str | None = None):
|
|
165
|
+
name = file_system or self.config.container
|
|
166
|
+
return self._get_service_client().get_file_system_client(name)
|
|
167
|
+
|
|
168
|
+
def _remove_directory(self, file_system: str, path: str) -> None:
|
|
169
|
+
from azure.core.exceptions import ResourceNotFoundError
|
|
170
|
+
|
|
171
|
+
try:
|
|
172
|
+
self._get_file_system_client(file_system).get_directory_client(
|
|
173
|
+
path
|
|
174
|
+
).delete_directory()
|
|
175
|
+
except ResourceNotFoundError:
|
|
176
|
+
pass
|
|
177
|
+
|
|
178
|
+
# ── schema management ─────────────────────────────────────────────────────
|
|
179
|
+
|
|
180
|
+
def create_schema(self, relation: PolarsRelation) -> None:
|
|
181
|
+
schema = relation.schema or ""
|
|
182
|
+
path = self._object_path(schema)
|
|
183
|
+
logger.debug(f"create_schema: creating directory {path}")
|
|
184
|
+
self._get_file_system_client().create_directory(path)
|
|
185
|
+
|
|
186
|
+
def drop_schema(self, relation: PolarsRelation) -> None:
|
|
187
|
+
schema = relation.schema or ""
|
|
188
|
+
path = self._object_path(schema)
|
|
189
|
+
logger.debug(f"Dropping schema {relation.catalog}/{schema}")
|
|
190
|
+
self._remove_directory(self.config.container, path)
|
|
191
|
+
|
|
192
|
+
def list_schemas(self) -> list[str]:
|
|
193
|
+
from azure.core.exceptions import ResourceNotFoundError
|
|
194
|
+
|
|
195
|
+
base = self.config.prefix or ""
|
|
196
|
+
try:
|
|
197
|
+
return [
|
|
198
|
+
p.name.split("/")[-1]
|
|
199
|
+
for p in self._get_file_system_client().get_paths(
|
|
200
|
+
path=base, recursive=False
|
|
201
|
+
)
|
|
202
|
+
if p.is_directory
|
|
203
|
+
]
|
|
204
|
+
except ResourceNotFoundError:
|
|
205
|
+
return []
|
|
206
|
+
|
|
207
|
+
def table_exists(self, relation: PolarsRelation) -> bool:
|
|
208
|
+
uri = self._get_uri(relation)
|
|
209
|
+
opts = self._get_storage_options(uri)
|
|
210
|
+
if relation.file_format == "delta":
|
|
211
|
+
return DeltaTable.is_deltatable(uri, storage_options=opts)
|
|
212
|
+
file_path = uri.removeprefix(f"az://{self.config.container}/")
|
|
213
|
+
return (
|
|
214
|
+
self._get_file_system_client(self.config.container)
|
|
215
|
+
.get_file_client(file_path)
|
|
216
|
+
.exists()
|
|
217
|
+
)
|
|
218
|
+
|
|
219
|
+
def drop_relation(self, relation: PolarsRelation) -> None:
|
|
220
|
+
schema = relation.schema or ""
|
|
221
|
+
identifier = relation.identifier or ""
|
|
222
|
+
file_system = self.config.container
|
|
223
|
+
logger.debug(
|
|
224
|
+
f"Dropping table if exists {relation.catalog}/{schema}/{identifier}"
|
|
225
|
+
)
|
|
226
|
+
if relation.file_format == "delta":
|
|
227
|
+
self._remove_directory(file_system, self._object_path(schema, identifier))
|
|
228
|
+
else:
|
|
229
|
+
from azure.core.exceptions import ResourceNotFoundError
|
|
230
|
+
|
|
231
|
+
file_path = (
|
|
232
|
+
f"{self._object_path(schema, identifier)}.{relation.file_format}"
|
|
233
|
+
)
|
|
234
|
+
try:
|
|
235
|
+
self._get_file_system_client(file_system).get_file_client(
|
|
236
|
+
file_path
|
|
237
|
+
).delete_file()
|
|
238
|
+
except ResourceNotFoundError:
|
|
239
|
+
pass
|
|
240
|
+
|
|
241
|
+
def list_relations_without_caching(
|
|
242
|
+
self, schema_relation: PolarsRelation
|
|
243
|
+
) -> list[PolarsRelation]:
|
|
244
|
+
from azure.core.exceptions import ResourceNotFoundError
|
|
245
|
+
|
|
246
|
+
schema = schema_relation.schema or ""
|
|
247
|
+
file_system = self.config.container
|
|
248
|
+
prefix = self._object_path(schema)
|
|
249
|
+
relations: list[PolarsRelation] = []
|
|
250
|
+
|
|
251
|
+
try:
|
|
252
|
+
paths = list(
|
|
253
|
+
self._get_file_system_client(file_system).get_paths(
|
|
254
|
+
path=prefix, recursive=False
|
|
255
|
+
)
|
|
256
|
+
)
|
|
257
|
+
except ResourceNotFoundError:
|
|
258
|
+
return []
|
|
259
|
+
|
|
260
|
+
for path in paths:
|
|
261
|
+
name = path.name.split("/")[-1]
|
|
262
|
+
if path.is_directory:
|
|
263
|
+
uri = f"az://{file_system}/{path.name}"
|
|
264
|
+
if DeltaTable.is_deltatable(
|
|
265
|
+
uri, storage_options=self._get_storage_options(uri)
|
|
266
|
+
):
|
|
267
|
+
relations.append(
|
|
268
|
+
schema_relation.create(
|
|
269
|
+
database=schema_relation.database,
|
|
270
|
+
schema=schema,
|
|
271
|
+
identifier=name,
|
|
272
|
+
type=RelationType.Table,
|
|
273
|
+
catalog=schema_relation.catalog,
|
|
274
|
+
file_format="delta",
|
|
275
|
+
)
|
|
276
|
+
)
|
|
277
|
+
else:
|
|
278
|
+
stem, _, ext = name.rpartition(".")
|
|
279
|
+
if ext in _FILE_FORMATS:
|
|
280
|
+
relations.append(
|
|
281
|
+
schema_relation.create(
|
|
282
|
+
database=schema_relation.database,
|
|
283
|
+
schema=schema,
|
|
284
|
+
identifier=stem,
|
|
285
|
+
type=RelationType.Table,
|
|
286
|
+
catalog=schema_relation.catalog,
|
|
287
|
+
file_format=ext,
|
|
288
|
+
)
|
|
289
|
+
)
|
|
290
|
+
|
|
291
|
+
return relations
|