datamasque-python 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

datamasque/client/__init__.py

@@ -30,7 +30,7 @@ from datamasque.client.models.connection import (
     ConnectionId,
     DatabaseConnectionConfig,
     DatabaseType,
-    DatabricksDeltaS3ConnectionConfig,
+    DatabricksConnectionConfig,
     DynamoConnectionConfig,
     FileConnectionConfig,
     MongoConnectionConfig,
@@ -129,7 +129,7 @@ __all__ = [
     "DataMasqueUserError",
     "DatabaseConnectionConfig",
     "DatabaseType",
-    "DatabricksDeltaS3ConnectionConfig",
+    "DatabricksConnectionConfig",
     "DiscoveryMatch",
     "DynamoConnectionConfig",
     "FailedToStartError",

datamasque/client/base.py

@@ -1,7 +1,10 @@
 import logging
+import platform
+import sys
 import warnings
 from contextlib import contextmanager
 from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, version
 from io import BufferedIOBase, BytesIO, TextIOBase
 from pathlib import Path
 from typing import Any, Callable, Iterator, Optional, Type, TypeVar, Union
@@ -24,6 +27,44 @@ logger = logging.getLogger(__name__)
 FileOrContent = Union[str, bytes, TextIOBase, BufferedIOBase, Path]
 _T = TypeVar("_T", bound=BaseModel)
 
+
+def _build_user_agent() -> str:
+    """
+    Identify ourselves to the DataMasque server in access logs and audit trails.
+
+    Default `python-requests/x.y.z` is anonymous; this surfaces the SDK name +
+    version, Python interpreter, and OS so operators can correlate API traffic
+    with a specific SDK release (e.g. when triaging a bug report).
+    """
+
+    try:
+        sdk_version = version("datamasque-python")
+    except PackageNotFoundError:
+        # Source checkouts without installed metadata.
+        sdk_version = "dev"
+    py = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+    return f"datamasque-python/{sdk_version} (Python/{py}; {platform.system()}/{platform.release()})"
+
+
+USER_AGENT = _build_user_agent()
+
+
+def _build_session(verify_ssl: bool) -> requests.Session:
+    """
+    Build a configured `requests.Session` for one client's lifetime.
+
+    Centralises the `User-Agent` and `verify` defaults so every call site
+    inherits them automatically — keeping the per-call code free of
+    boilerplate and removing the risk of forgetting either flag on a new
+    endpoint.
+    """
+
+    session = requests.Session()
+    session.headers["User-Agent"] = USER_AGENT
+    session.verify = verify_ssl
+    return session
+
+
 # Substrings (case-insensitive) that mark a key whose value should be redacted
 # before logging on an error path, so that passwords, API tokens, and similar secrets don't
 # end up in user-visible logs when a request fails.
@@ -71,6 +112,15 @@ class BaseClient:
 
     Holds the connection config, cached auth token, and the core `make_request` dispatcher
     used by all per-feature mixins that compose `DataMasqueClient`.
+
+    Uses a single `requests.Session` for the lifetime of the client so that
+    per-host TCP / TLS connections are pooled across calls (paginated list
+    endpoints and tight polling loops benefit most). Session-wide defaults
+    (`User-Agent`, `verify`) are set once on construction; per-call headers
+    like `Authorization` are merged at request time.
+
+    `requests.Session` is not thread-safe; do not share a client between
+    threads. Construct one per worker.
     """
 
     token: str = ""
@@ -86,6 +136,7 @@ class BaseClient:
         self.password = connection_config.password
         self.verify_ssl = connection_config.verify_ssl
         self.token_source = connection_config.token_source
+        self._session = _build_session(self.verify_ssl)
 
     @contextmanager
     def _maybe_suppress_insecure_warning(self) -> Iterator[None]:
@@ -186,28 +237,32 @@ class BaseClient:
         url = urljoin(self.base_url, path)
 
         def send() -> Response:
-            headers: Optional[dict] = {"Authorization": self.token} if requires_authorization else None
+            headers = {"Authorization": self.token} if requires_authorization else None
             try:
                 with self._maybe_suppress_insecure_warning():
                     if files:
                         files_payload = {f.field_name: (f.filename, f.content, f.content_type or "") for f in files}
-                        return requests.request(
+                        return self._session.request(
                             method,
                             url,
                             data=data,
                             params=params,
                             headers=headers,
                             files=files_payload,
-                            verify=self.verify_ssl,
                         )
-                    return requests.request(
-                        method, url, json=data, params=params, headers=headers, verify=self.verify_ssl
-                    )
+                    return self._session.request(method, url, json=data, params=params, headers=headers)
             except requests.RequestException as e:
                 raise DataMasqueTransportError(f"Failed to reach DataMasque server at {url}: {e}") from e
 
         response = send()
-        if response.status_code == 401:
+        if response.status_code == 401 and requires_authorization:
+            # Token-expiry recovery: re-auth and replay. Only meaningful when the
+            # caller actually sent a token; on `requires_authorization=False`
+            # calls a 401 means the server itself is rejecting anonymous access
+            # (e.g. admin-install on an already-configured instance), and
+            # re-authing with whatever creds the client happens to hold would
+            # both misdiagnose the failure and emit a misleading
+            # "credentials are incorrect" error to the user.
             logger.debug("Re-authenticating")
             self.authenticate()
             # Reset file pointers so the retry doesn't send empty files

datamasque/client/ifm.py

@@ -17,7 +17,7 @@ import requests
 from pydantic import BaseModel
 from requests import Response
 
-from datamasque.client.base import suppress_insecure_warning_if_needed
+from datamasque.client.base import _build_session, suppress_insecure_warning_if_needed
 from datamasque.client.exceptions import (
     DataMasqueApiError,
     DataMasqueNotReadyError,
@@ -82,6 +82,9 @@ class DataMasqueIfmClient:
         self.password = connection_config.password
         self.verify_ssl = connection_config.verify_ssl
         self.token_source = connection_config.token_source
+        # One session for both admin-server (JWT login/refresh) and IFM (data plane)
+        # traffic -- different hosts, but a single session handles per-host pooling.
+        self._session = _build_session(self.verify_ssl)
 
     def authenticate(self) -> None:
         """Obtain an access (and refresh) token from the admin server, or via `token_source`."""
@@ -95,10 +98,9 @@ class DataMasqueIfmClient:
         login_url = urljoin(self.admin_server_base_url, "/api/auth/jwt/login/")
         try:
             with self._maybe_suppress_insecure_warning():
-                response = requests.post(
+                response = self._session.post(
                     login_url,
                     json={"username": self.username, "password": self.password},
-                    verify=self.verify_ssl,
                 )
         except requests.RequestException as e:
             raise DataMasqueTransportError(f"Failed to reach admin server at {login_url}: {e}") from e
@@ -122,10 +124,9 @@ class DataMasqueIfmClient:
         refresh_url = urljoin(self.admin_server_base_url, "/api/auth/jwt/refresh/")
         try:
             with self._maybe_suppress_insecure_warning():
-                response = requests.post(
+                response = self._session.post(
                     refresh_url,
                     json={"refresh": self.refresh_token},
-                    verify=self.verify_ssl,
                 )
         except requests.RequestException as e:
             raise DataMasqueTransportError(f"Failed to reach admin server at {refresh_url}: {e}") from e
@@ -187,13 +188,12 @@ class DataMasqueIfmClient:
         def send() -> Response:
             try:
                 with self._maybe_suppress_insecure_warning():
-                    return requests.request(
+                    return self._session.request(
                         method,
                         url,
                         json=json_body,
                         params=params,
                         headers={"Authorization": f"Bearer {self.access_token}"},
-                        verify=self.verify_ssl,
                     )
             except requests.RequestException as e:
                 raise DataMasqueTransportError(f"Failed to reach IFM server at {url}: {e}") from e

datamasque/client/models/connection.py

@@ -45,6 +45,8 @@ class DatabaseType(Enum):
     snowflake = "snowflake"
     mongodb = "mongodb"
     databricks_lakebase = "databricks_lakebase"
+    databricks = "databricks"
+    informix = "informix"
 
 
 class SnowflakeStageLocation(str, Enum):
@@ -280,6 +282,8 @@ class DatabaseConnectionConfig(ConnectionConfig):
             raise ValueError("For Snowflake, use the SnowflakeConnectionConfig class instead")
         if self.database_type is DatabaseType.mongodb:
             raise ValueError("For MongoDB, use the MongoConnectionConfig class instead")
+        if self.database_type is DatabaseType.databricks:
+            raise ValueError("For Databricks SQL Warehouse, use the DatabricksConnectionConfig class instead")
         return self
 
     mask_type: Literal["database"] = "database"
@@ -391,19 +395,36 @@ class MountedShareConnectionConfig(FileConnectionConfig):
     type: Literal["mounted_share_connection"] = "mounted_share_connection"
 
 
-class DatabricksDeltaS3ConnectionConfig(FileConnectionConfig):
-    """Connection configuration for Databricks Delta tables stored in S3."""
+class DatabricksConnectionConfig(ConnectionConfig):
+    """Connection configuration for a Databricks SQL Warehouse."""
 
-    type: Literal["databricks_delta_s3_connection"] = "databricks_delta_s3_connection"
-    bucket: str = ""
-    iam_role_arn: Optional[str] = None
+    server_hostname: str
+    http_path: str
+    access_token: Optional[str] = None
+    catalog: str
+    db_schema: Optional[str] = Field(default=None, alias="schema")
+    is_read_only: bool = False
+    version: str = "1.0"
+
+    mask_type: Literal["database"] = "database"
+    db_type: Literal["databricks"] = "databricks"
+
+    @property
+    def database_type(self) -> DatabaseType:
+        return DatabaseType.databricks
+
+    @model_validator(mode="before")
+    @classmethod
+    def _strip_encrypted_token(cls, data: dict) -> dict:
+        if isinstance(data, dict):
+            data.pop("access_token_encrypted", None)
+        return data
 
 
 FILE_TYPE_MAP: dict[str, type[FileConnectionConfig]] = {
     "s3_connection": S3ConnectionConfig,
     "azure_blob_connection": AzureConnectionConfig,
     "mounted_share_connection": MountedShareConnectionConfig,
-    "databricks_delta_s3_connection": DatabricksDeltaS3ConnectionConfig,
 }
 
 DB_TYPE_MAP: dict[str, type[ConnectionConfig]] = {
@@ -411,6 +432,7 @@ DB_TYPE_MAP: dict[str, type[ConnectionConfig]] = {
     DatabaseType.mongodb.value: MongoConnectionConfig,
     DatabaseType.snowflake.value: SnowflakeConnectionConfig,
     DatabaseType.mssql_linked.value: MssqlLinkedServerConnectionConfig,
+    DatabaseType.databricks.value: DatabricksConnectionConfig,
     # others use the default `DatabaseConnectionConfig`
 }
 

{datamasque_python-1.0.2.dist-info → datamasque_python-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datamasque-python
-Version: 1.0.2
+Version: 1.0.4
 Summary: Official Python client for the DataMasque data-masking API.
 Project-URL: Homepage, https://datamasque.com/
 Project-URL: Documentation, https://datamasque-python.readthedocs.io/

{datamasque_python-1.0.2.dist-info → datamasque_python-1.0.4.dist-info}/RECORD

@@ -1,11 +1,11 @@
-datamasque/client/__init__.py,sha256=kEWW_L7t0J2gsBC9zwLUN3p8I2p26RqGHXo02aI3WRI,5486
-datamasque/client/base.py,sha256=he-ObABs_rpXV_KSMmrtJ73D-9lb0giAK029lN5qB3Y,11817
+datamasque/client/__init__.py,sha256=rH7hQyA_nFVkv8eV_C1Ds_40v9xzCK__uXIJtWMLfuQ,5472
+datamasque/client/base.py,sha256=vP1LuYbqq6U8NAEJWnwNrdt_Fa2wjQTJXHf0gSFO79c,14115
 datamasque/client/connections.py,sha256=EFinx8fJRme0mTxuWY3d29UnmUFbsQhMaUQT0Ma2PK4,2885
 datamasque/client/discovery.py,sha256=uA8h6vRqsxSAzSAV9bebJwU47XINlaVy7V1nTYDiaCM,12634
 datamasque/client/dmclient.py,sha256=OPYMzc57gUHPs6iL_J2DYp06MfOaabpTlISmnNCpqS4,1553
 datamasque/client/exceptions.py,sha256=F9FYCxP-ERXkVD1L3yh_rWqcW3IsPL-c-Ic4qMYoGnw,2542
 datamasque/client/files.py,sha256=5Gzel4aLby8T7ncOIV6wtgVbFEk0eaEy7wxohh9u0zE,3468
-datamasque/client/ifm.py,sha256=K4n9q-8f2wsmxnKlE-jkn9zl40iC9RGAw2tRB0mAtyM,11737
+datamasque/client/ifm.py,sha256=uIMxpLIPvDiDO1m4bxezixNIUFIFY0MXWRMQNvTbpTA,11858
 datamasque/client/license.py,sha256=pluYaSU168OC6_laB9bM9H3Vuuxs8wa9gujCJvtoJCk,1392
 datamasque/client/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datamasque/client/ruleset_libraries.py,sha256=tyN--cndzG0gwFnHY9fDtObTLzGsK0wrV5lxxjRP72g,6868
@@ -14,7 +14,7 @@ datamasque/client/runs.py,sha256=ZPSkkuyqMiwy7dLbWZ1PAEKaesKLeNRX7xEJGfzieVg,742
 datamasque/client/settings.py,sha256=Ui8AyR2XdoW8MZ9FIrGn2jm8DLzUGWIL8i2DOTvu7hc,2898
 datamasque/client/users.py,sha256=VCUo2CJyOw4-aO_3mp_w_0BcSoa6-h1HcReMcAMyumw,3701
 datamasque/client/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datamasque/client/models/connection.py,sha256=mQtWLAKcExbOkMymKmadprdWLbXl0V9aFZ4Zla2v0jc,15091
+datamasque/client/models/connection.py,sha256=DQkn1kVTu3ODNNBjd-nR1JZhck5kalWhHJBuDWGNzU0,15778
 datamasque/client/models/data_selection.py,sha256=406yyUZ5NmLBSql2lYM1gsTWY5GWmnutmF-DjznAoLc,1904
 datamasque/client/models/discovery.py,sha256=BawKusPuhyt0gRRnWKYf-ZKF7V54GxkhgYIrEQ6XkOs,7283
 datamasque/client/models/dm_instance.py,sha256=yjjpHZJTFhJp3lAinTEffgKrxnadrJys1EuhO77wQcA,1561
@@ -27,7 +27,7 @@ datamasque/client/models/ruleset_library.py,sha256=gIqb6yn4-f9i2OydOLk1cd0zId9nG
 datamasque/client/models/runs.py,sha256=oVYo9jp9s5LjWts0LKsYpX3HUBmrVwuuzDqFtQmTjvo,5673
 datamasque/client/models/status.py,sha256=rjH6YSwAHoOUaUCADwvMhuKd0ygT0c2w2Ek5SV8PWD8,1984
 datamasque/client/models/user.py,sha256=UGAUzgJkf78m24_zFXXoA99zdut48BXkX_ivV8yq1Vc,2043
-datamasque_python-1.0.2.dist-info/METADATA,sha256=-D8IeXp8l_ZzTAp2gEjNBppqtyaXrz-j1SWOqTUJscI,4187
-datamasque_python-1.0.2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-datamasque_python-1.0.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-datamasque_python-1.0.2.dist-info/RECORD,,
+datamasque_python-1.0.4.dist-info/METADATA,sha256=eIH7wSpr6f9J6twjwzCbrKHZ7Sn3jFs-DLXo_F5Lzus,4187
+datamasque_python-1.0.4.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+datamasque_python-1.0.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+datamasque_python-1.0.4.dist-info/RECORD,,

{datamasque_python-1.0.2.dist-info → datamasque_python-1.0.4.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.29.0
+Generator: hatchling 1.30.1
 Root-Is-Purelib: true
 Tag: py3-none-any