unstructured/test_unstructured_ingest/unit/test_enhanced_dataclass.py
Roman Isecke f193d3d43b
feat: improve sensitive data handling by fsspec connectors (#2194)
### Description
Building off of PR
https://github.com/Unstructured-IO/unstructured/pull/2179, updating
fsspec based connectors to use better authentication field handling.
This PR adds in the following changes:

* Update the base classes to inherit from the enhanced json mixin
* Add in a new access config dataclass that should be used as a nest
dataclass in the connector configs
* Update the code extracting configs out of the cli options dictionary
to support the nested access config if it exists on the parent config
* Update all fsspec connectors with explicit access configs given what
each one's SDKs support
* Update the json mixin and enhanced field to support a name override
when serializing/deserializing from json/dicts. This allows a different
name to be used for the CLI option than what the name of the field is on
the dataclass.
* Update all the writes to use class-based approach and share the same
structure of the runner classes
* Above update allowed for better code to be used in the base source and
destination CLI commands
* Add in utility code around paring a flat dictionary (coming from the
click based options) into dataclass-based configs with potentially
nested dataclasses.

**Slightly unrelated changes:**
* session handle removed from pinecone connector as this was breaking
the serialization of the write config and didn't have any benefit as a
connection was never being shared, the index used simply makes a new
http call each time it's invoked.
* Dedicated write configs were created for all destination connectors to
better support serialization
* Refactor of Elasticsearch connector included, with update to ingest
test to use auth

**TODOs**
* Left a `#TODO` in the code but the way session handler is implemented
right now, it breaks serialization since it adds a generic variable
based on the library being used for a connector (i.e.
`googleapiclient.discovery.Resource`) which is not serializable. This
will need to be updated to omit that from serialization but still
support the current workflow.

---------

Co-authored-by: ryannikolaidis <1208590+ryannikolaidis@users.noreply.github.com>
Co-authored-by: rbiseck3 <rbiseck3@users.noreply.github.com>
2023-12-05 20:55:19 +00:00

61 lines
1.8 KiB
Python

import json
from dataclasses import Field, dataclass, fields
import pytest
from unstructured.ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
from unstructured.ingest.enhanced_dataclass.dataclasses import EnhancedField
@dataclass
class AuthData(EnhancedDataClassJsonMixin):
username: str
password: str = enhanced_field(sensitive=True)
date: int = enhanced_field(overload_name="time")
auth = AuthData(username="my name", password="top secret", date=3)
def test_enhanced_field():
fs = fields(AuthData)
for f in fs:
if f.name == "username":
assert isinstance(f, Field)
assert hasattr(f, "sensitive") is False
else:
assert isinstance(f, EnhancedField)
if f.name == "password":
assert f.sensitive is True
else:
assert not f.sensitive
@pytest.mark.parametrize(
("apply_name_overload", "expected_dict"),
[
(True, {"username": "my name", "password": "THIS IS REDACTED", "time": 3}),
(False, {"username": "my name", "password": "THIS IS REDACTED", "date": 3}),
],
)
def test_to_json(apply_name_overload: bool, expected_dict: dict):
j = auth.to_json(
redact_sensitive=True,
redacted_text="THIS IS REDACTED",
apply_name_overload=apply_name_overload,
)
expected = json.dumps(expected_dict)
assert j == expected
@pytest.mark.parametrize(
("apply_name_overload", "expected_dict"),
[
(True, {"username": "my name", "password": "***REDACTED***", "time": 3}),
(False, {"username": "my name", "password": "***REDACTED***", "date": 3}),
],
)
def test_to_dict(apply_name_overload: bool, expected_dict: dict):
d = auth.to_dict(redact_sensitive=True, apply_name_overload=apply_name_overload)
assert d == expected_dict