vastdb 2.0.3__tar.gz → 2.0.5__tar.gz

{vastdb-2.0.3 → vastdb-2.0.5}/CHANGELOG.md

@@ -4,6 +4,19 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [2.0.5] (2025-01-08)
+[2.0.4]: https://github.com/vast-data/vastdb_sdk/compare/v2.0.4...v2.0.5
+
+### Added
+- end-user impersonation for adbc connection configurable via Session
+
+
+## [2.0.4] (2025-01-06)
+[2.0.3]: https://github.com/vast-data/vastdb_sdk/compare/v2.0.3...v2.0.4
+
+### Changed
+- Fix race condition in select (get data) mechanism
+
 ## [2.0.3] (2025-11-25)
 [2.0.2]: https://github.com/vast-data/vastdb_sdk/compare/v2.0.2...v2.0.3
 

{vastdb-2.0.3 → vastdb-2.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vastdb
-Version: 2.0.3
+Version: 2.0.5
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA

{vastdb-2.0.3 → vastdb-2.0.5}/setup.py

@@ -29,7 +29,7 @@ setup(
     name='vastdb',
     python_requires='>=3.10.0',
     description='VAST Data SDK',
-    version='2.0.3' + suffix,
+    version='2.0.5' + suffix,
     url='https://github.com/vast-data/vastdb_sdk',
     author='VAST DATA',
     author_email='hello@vastdata.com',

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/_adbc.py

@@ -17,6 +17,7 @@ log = logging.getLogger(__name__)
 
 
 TXID_OVERRIDE_PROPERTY: str = "vast.db.external_txid"
+END_USER_PROPERTY: str = "vast.db.end_user"
 VAST_DIST_ALIAS = "vast_pysdk_vector_dist"
 DEFAULT_ADBC_DRIVER_CACHE_DIR: str = "~/.vast/adbc_drivers_cache"
 DEFAULT_ADBC_DRIVER_CACHE_BY_URL_DIR: str = f"{DEFAULT_ADBC_DRIVER_CACHE_DIR}/by_url"
@@ -83,9 +84,18 @@ class AdbcDriver:
 
 
 def _get_adbc_connection(
-    adbc_driver_path: str, endpoint: str, access_key: str, secret_key: str, txid: int
+    adbc_driver_path: str,
+    endpoint: str,
+    access_key: str,
+    secret_key: str,
+    txid: int,
+    end_user: Optional[str],
 ) -> Connection:
     """Get an adbc connection in transaction."""
+    conn_kwargs = {TXID_OVERRIDE_PROPERTY: str(txid)}
+    if end_user is not None:
+        conn_kwargs[END_USER_PROPERTY] = end_user
+
     return connect(
         driver=adbc_driver_path,
         db_kwargs={
@@ -93,7 +103,7 @@ def _get_adbc_connection(
             "vast.db.access_key": access_key,
             "vast.db.secret_key": secret_key,
         },
-        conn_kwargs={TXID_OVERRIDE_PROPERTY: str(txid)},
+        conn_kwargs=conn_kwargs,
     )
 
 
@@ -157,9 +167,10 @@ class AdbcConnection:
         access_key: str,
         secret_key: str,
         txid: int,
+        end_user: Optional[str] = None,
     ):
         self._adbc_conn = _get_adbc_connection(
-            adbc_driver.local_path, endpoint, access_key, secret_key, txid
+            adbc_driver.local_path, endpoint, access_key, secret_key, txid, end_user
         )
 
         self._cursor = self._adbc_conn.cursor()

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/_internal.py

@@ -1117,7 +1117,7 @@ class VastdbApi:
         if not bucket.startswith(bucket_name):
             raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
         schemas_length = lists.SchemasLength()
-        count = int(res_headers['tabular-list-count']) if 'tabular-list-count' in res_headers else schemas_length
+        count = int(res_headers['tabular-list-count']) if count_only else schemas_length
         for i in range(schemas_length):
             schema_obj = lists.Schemas(i)
             name = schema_obj.Name().decode()
@@ -1427,7 +1427,7 @@ class VastdbApi:
         is_truncated = res_headers['tabular-is-truncated'] == 'true'
         lists = list_tables.GetRootAs(res.content)
         tables_length = lists.TablesLength()
-        count = int(res_headers['tabular-list-count']) if 'tabular-list-count' in res_headers else tables_length
+        count = int(res_headers['tabular-list-count']) if count_only else tables_length
         return lists, next_key, is_truncated, count
 
     def _list_tables_internal(self, bucket, schema, parse_properties, txid=0, client_tags=[], max_keys=1000, next_key=0, name_prefix="",
@@ -1574,8 +1574,8 @@ class VastdbApi:
         res_headers = res.headers
         next_key = int(res_headers['tabular-next-key'])
         is_truncated = res_headers['tabular-is-truncated'] == 'true'
-        count = int(res_headers['tabular-list-count'])
         columns = [] if count_only else pa.ipc.open_stream(res.content).schema
+        count = int(res_headers['tabular-list-count']) if count_only else len(columns)
 
         return columns, next_key, is_truncated, count
 
@@ -2044,7 +2044,6 @@ class VastdbApi:
         res_headers = res.headers
         next_key = int(res_headers['tabular-next-key'])
         is_truncated = res_headers['tabular-is-truncated'] == 'true'
-        count = int(res_headers['tabular-list-count'])
         lists = list_projections.GetRootAs(res.content)
         bucket_name = lists.BucketName().decode()
         schema_name = lists.SchemaName().decode()
@@ -2052,6 +2051,7 @@ class VastdbApi:
         if not bucket.startswith(bucket_name):  # ignore snapshot name
             raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
         projections_length = lists.ProjectionsLength()
+        count = int(res_headers['tabular-list-count']) if count_only else projections_length
         for i in range(projections_length):
             projections.append(_parse_table_info(lists.Projections(i), lambda x: x))
 
@@ -2092,10 +2092,10 @@ class VastdbApi:
         res_headers = res.headers
         next_key = int(res_headers['tabular-next-key'])
         is_truncated = res_headers['tabular-is-truncated'] == 'true'
-        count = int(res_headers['tabular-list-count'])
         columns = [] if count_only else [[f.name, f.type, f.metadata] for f in
                                          pa.ipc.open_stream(res.content).schema]
 
+        count = int(res_headers['tabular-list-count']) if count_only else len(columns)
         return columns, next_key, is_truncated, count
 
 

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/schema.py

@@ -89,6 +89,7 @@ class Schema:
         See https://support.vastdata.com/s/article/UUID-48d0a8cf-5786-5ef3-3fa3-9c64e63a0967 for more details.
 
         Args:
+        ----
             table_name: Name of the table to create
             columns: PyArrow schema defining the table columns
             fail_if_exists: Whether to fail if the table already exists
@@ -97,6 +98,7 @@ class Schema:
             vector_index: Optional vector index.
 
         Returns:
+        -------
             The created table
 
         """

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/session.py

@@ -25,7 +25,8 @@ class Session:
                  ssl_verify=True,
                  timeout=None,
                  backoff_config: Optional["BackoffConfig"] = None,
-                 adbc_driver: Optional[AdbcDriver] = None):
+                 adbc_driver: Optional[AdbcDriver] = None,
+                 end_user: Optional[str] = None):
         """Connect to a VAST Database endpoint, using specified credentials."""
         from . import _internal, features
 
@@ -49,6 +50,7 @@ class Session:
             backoff_config=backoff_config)
         self.features = features.Features(self.api.vast_version)
         self.adbc_driver: Optional[AdbcDriver] = adbc_driver
+        self._end_user = end_user
 
     def __repr__(self):
         """Don't show the secret key."""
@@ -63,4 +65,4 @@ class Session:
                 tx.bucket("bucket").create_schema("schema")
         """
         from . import transaction
-        return transaction.Transaction(self, _adbc_driver=self.adbc_driver)
+        return transaction.Transaction(self, _adbc_driver=self.adbc_driver, _end_user=self._end_user)

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/table.py

@@ -626,13 +626,12 @@ class TableInTransaction(ITable):
                             log.debug(
                                 "one worker thread finished, remaining: %d", tasks_running)
 
-                    # all host threads ended - wait for all futures to complete
-                    propagate_first_exception(futures, block=True)
                 finally:
                     stop_event.set()
                     while tasks_running > 0:
                         if record_batches_queue.get() is None:
                             tasks_running -= 1
+                    propagate_first_exception(futures, block=True)
 
         return pa.RecordBatchReader.from_batches(query_data_request.response_schema, batches_iterator())
 

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/tests/test_adbc_integration.py

@@ -1,6 +1,11 @@
+from typing import Optional
+from unittest.mock import MagicMock, patch
+
 import pyarrow as pa
 import pytest
 
+from vastdb._adbc import END_USER_PROPERTY, AdbcDriver
+from vastdb.session import Session
 from vastdb.table_metadata import TableRef
 from vastdb.transaction import NoAdbcConnectionError
 
@@ -87,3 +92,38 @@ def test_two_simulatnious_txs_with_adbc(session_factory, clean_bucket_name: str)
         # expecting adbc execute to "see" table if it shares the transaction with the pysdk
         tx.adbc_conn.cursor.execute(f"SELECT * FROM {table.ref.query_engine_full_path}")
         assert tx.adbc_conn.cursor.fetchall() == [(1,), (2,), (3,), (4,), (5,)]
+
+
+@pytest.mark.parametrize("end_user", [("mock-end-user",), (None,)])
+def test_end_user_passed_to_adbc_connect(end_user: Optional[str]):
+    mock_driver = MagicMock(spec=AdbcDriver)
+    mock_driver.local_path = "/mock/driver/path"
+
+    with (
+        patch("vastdb._adbc.connect") as mock_connect,
+        patch("vastdb._internal.VastdbApi") as MockVastdbApi,
+    ):
+        mock_api_instance = MockVastdbApi.return_value
+        mock_api_instance.begin_transaction.return_value.headers = {
+            "tabular-txid": "12345"
+        }
+        # A version that supports everything needed.
+        mock_api_instance.vast_version = (5, 4, 0, 0)
+
+        session = Session(
+            access="test_access",
+            secret="test_secret",
+            endpoint="http://localhost:9090",
+            adbc_driver=mock_driver,
+            end_user=end_user,
+        )
+
+        with session.transaction():
+            # The ADBC connection is established when the transaction starts
+            pass
+
+        mock_connect.assert_called_once()
+        call_kwargs = mock_connect.call_args.kwargs
+        conn_kwargs = call_kwargs.get("conn_kwargs", {})
+        assert (end_user is None) ^ (END_USER_PROPERTY in conn_kwargs)
+        assert (end_user is None) ^ (conn_kwargs.get(END_USER_PROPERTY) == end_user)

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/tests/test_vector_search.py

@@ -163,6 +163,7 @@ def test_ibis_to_query_engine_predicates(ibis_predicate: IbisPredicate, expected
     assert _ibis_to_qe_predicates(ibis_predicate) == expected
 
 
+@pytest.mark.skip(reason="see https://vastdata.atlassian.net/browse/ORION-307908")
 def test_with_predicates_get_vector_index_properties_from_server(
     session_factory: SessionFactory,
     clean_bucket_name: str

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb/transaction.py

@@ -54,6 +54,7 @@ class Transaction:
     txid: Optional[int] = None
     _adbc_driver: Optional[AdbcDriver] = None
     _adbc_conn: Optional[AdbcConnection] = None
+    _end_user: Optional[str] = None
 
     def __enter__(self):
         """Create a transaction and store its ID."""
@@ -67,6 +68,7 @@ class Transaction:
                 self._rpc.access,
                 self._rpc.secret,
                 self.txid,
+                self._end_user
             )
 
         log.debug("opened txid=%016x", self.txid)

{vastdb-2.0.3 → vastdb-2.0.5}/vastdb.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vastdb
-Version: 2.0.3
+Version: 2.0.5
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA