vastdb 2.0.3__py3-none-any.whl → 2.0.5__py3-none-any.whl

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/_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/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/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/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/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/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/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.dist-info → vastdb-2.0.5.dist-info}/METADATA

@@ -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.dist-info → vastdb-2.0.5.dist-info}/RECORD

@@ -1,18 +1,18 @@
 vastdb/__init__.py,sha256=uf-AXdzsD4nPxFP7WxkcAXGG0whv8BHLrrXCJtsPGaQ,436
-vastdb/_adbc.py,sha256=jJphuI9vFv9TNw6hbV2XzDRjOboxBKcgyf-7voPbuco,5735
+vastdb/_adbc.py,sha256=aUjn0hC92U5Jh2jQYeMFzTHYnVP086dVHlMoIoxhvh0,5985
 vastdb/_ibis_support.py,sha256=sJieOMvDWpsciPKh1mJzS56jxLtCRVlvK41hW84vexM,866
-vastdb/_internal.py,sha256=gZZAV3p4ya9bUUkVghgEXUbzR3HukkAu757mtlexpzw,112715
+vastdb/_internal.py,sha256=7c6tySuEacwZ5jqdW-HSAtbIIk8ep-_jA4uGrEU9gwc,112767
 vastdb/_table_interface.py,sha256=dRcXzC7j1SSpBQ4jCOPF_Zgg66LSt0tBeILM8W_A71c,4130
 vastdb/bucket.py,sha256=ulkTI_7xw5FYVzcrTFC7G2ijmTTVSmvJZUdgzycGHR0,2588
 vastdb/config.py,sha256=OehnsWrjzv0-SUouEXmkrKBugiWyhXOn4XiSLV3s9yk,2342
 vastdb/conftest.py,sha256=vDNqhYMsF630_ueXofnT2zhuSEpjcc4hxcXAjPqlxDI,5525
 vastdb/errors.py,sha256=GVmiwyrfHrag8GL55ivcF5Kwc1Ro4wvvaJAcuIWvmO0,5764
 vastdb/features.py,sha256=ivYbvhiGA858B00vhs_CNzlVV9QDUe53yW6V3J5EoxM,1874
-vastdb/schema.py,sha256=apU2NATGg_Cv8Fnq8a1_g6Cc2-Fake2vV0NzsB5Xe8Q,7753
-vastdb/session.py,sha256=UEtYmrcOn5t80PkPUFaD7pGYkDKDWc-0LbNQ3BTsk58,2420
-vastdb/table.py,sha256=PXKZenXdtCXWiZ48D_2XLrSq2t1D256uJcAAGwXr4qE,46321
+vastdb/schema.py,sha256=7k6Su8xjXzeLZeqNzsXIk06q9sIWMG9PBBlwy1aEAG0,7782
+vastdb/session.py,sha256=oUiTx8x_5XOs6EnNvFjon9dSVTk89XhOZ-SfLR8sITQ,2529
+vastdb/table.py,sha256=X0uOMlxA9R2ER1-uEyB1lBv_nn-c_kM_jjaM2ai2IrQ,46241
 vastdb/table_metadata.py,sha256=PSwx1A25AmcCs4ig-Sx_xyXr7xqyGaNIk2LMnIvrwmM,7358
-vastdb/transaction.py,sha256=qdCdPgqT1CH73ut2mGPg0VZL90jyeShXD_QR-fGBWvw,4667
+vastdb/transaction.py,sha256=JaTvUw86HZ-qFiqjR1fmHndF5bSgMtVWj549MtGsZ5Q,4734
 vastdb/util.py,sha256=8CUnVRsJukC3uNHNoB5D0qPf0FxS8OSdVB84nNoLJKc,6290
 vastdb/bench/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/bench/test_perf.py,sha256=0kbCxK8U9vYO0zCMUYcZHzEICaaII3I0-6FeR5-CNtM,4040
@@ -44,7 +44,7 @@ vastdb/bench/perf_bench/query/query_pyarrow.py,sha256=Dj5YPUvb4dAj7RskHfJcPijJnM
 vastdb/bench/perf_bench/query/query_vastdb.py,sha256=SZYem_EmsaynEftAa_VFobjSJZDAcli9BckyRS3SFvg,2810
 vastdb/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/tests/metrics.py,sha256=ZCSeBYFSPMG3yI0JrAHs2CrY6wFjx_5GwRTYHVAwLKA,1026
-vastdb/tests/test_adbc_integration.py,sha256=7huOOe9ww6M6RABDo8nzv2uW1A6Xg7XnnhQ0-_e4Woc,3056
+vastdb/tests/test_adbc_integration.py,sha256=6Zt4plbKrr4lIQ2AxQN7GP4yh5vzOgYdSU0ql5g97SY,4521
 vastdb/tests/test_duckdb.py,sha256=XZxlw3SXb8tbd2BsJdr9F7jBBka4dW9PVtwvwqEcOXo,2057
 vastdb/tests/test_fixed_list.py,sha256=1KpsijmGvxHcM-pHJcm4zn2o_G8-ISbCFuoKSkXjmfk,14099
 vastdb/tests/test_imports.py,sha256=FhIHjYpttiQW_x3fMe8ENd9ynn7OyUtjxr2ca4z1F-s,21783
@@ -56,7 +56,7 @@ vastdb/tests/test_table_in_tx.py,sha256=-OFGaZMZQc78HOmR23xHIsH7d0VBneFG3tGdDlog
 vastdb/tests/test_tables.py,sha256=yU5fBeHentNe22Rc1e73mj33M0nl63Y_acOy6VTzF-4,54496
 vastdb/tests/test_util.py,sha256=n7gvT5Wg6b6bxgqkFXkYqvFd_W1GlUdVfmPv66XYXyA,1956
 vastdb/tests/test_vector_index.py,sha256=qgppSQG77PUgmdXzfprX1mxUkkqq7tbxCA3MRaYFFgQ,6460
-vastdb/tests/test_vector_search.py,sha256=z3oxzejIaJw8nPZEvrG74k-zhDGRvu9C5Dhkqv--Zpc,7020
+vastdb/tests/test_vector_search.py,sha256=50mnBFuw0PAApMLkYwDUMV5VxQorswzD_f_XdpaEXTo,7103
 vastdb/tests/util.py,sha256=cv8Qcvy92vpBbiiuqXr7pAfbh2K8ziBVZ_j1EQ4vT6o,5994
 vastdb/vast_flatbuf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/vast_flatbuf/org/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -219,8 +219,8 @@ vastdb/vast_flatbuf/tabular/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 vastdb/vast_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/vast_tests/test_ha.py,sha256=744P4G6VJ09RIkHhMQL4wlipCBJWQVMhyvUrSc4k1HQ,975
 vastdb/vast_tests/test_scale.py,sha256=5jGwOdZH6Tv5tPdZYPWoqcxOceI2jA5i2D1zNKZHER4,3958
-vastdb-2.0.3.dist-info/licenses/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
-vastdb-2.0.3.dist-info/METADATA,sha256=lGMKtf0KnNZhPAtFyXCXu6HA0jbsmIDMVw1II061vww,1721
-vastdb-2.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-vastdb-2.0.3.dist-info/top_level.txt,sha256=nnKAaZaQa8GFbYpWAexr_B9HrhonZbUlX6hL6AC--yA,7
-vastdb-2.0.3.dist-info/RECORD,,
+vastdb-2.0.5.dist-info/licenses/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
+vastdb-2.0.5.dist-info/METADATA,sha256=4wi1lvs2ENDLmNaGUK4cFq_AS_IpWtYBIQJdNvrV7Nc,1721
+vastdb-2.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+vastdb-2.0.5.dist-info/top_level.txt,sha256=nnKAaZaQa8GFbYpWAexr_B9HrhonZbUlX6hL6AC--yA,7
+vastdb-2.0.5.dist-info/RECORD,,