vastdb 0.0.5.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

vastdb/tests/test_imports.py

@@ -0,0 +1,125 @@
+import pytest
+
+from tempfile import NamedTemporaryFile
+import logging
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from vastdb.errors import InvalidArgument, ImportFilesError
+from vastdb import util
+
+
+log = logging.getLogger(__name__)
+
+
+def test_parallel_imports(session, clean_bucket_name, s3):
+    num_rows = 1000
+    num_files = 53
+    ds = {'num': [i for i in range(num_rows)]}
+    files = []
+    table = pa.Table.from_pydict(ds)
+    with NamedTemporaryFile() as f:
+        pq.write_table(table, f.name)
+        s3.put_object(Bucket=clean_bucket_name, Key='prq0', Body=f)
+        files.append(f'/{clean_bucket_name}/prq0')
+
+    for i in range(1, num_files):
+        copy_source = {
+            'Bucket': clean_bucket_name,
+            'Key': 'prq0'
+        }
+        s3.copy(copy_source, clean_bucket_name, f'prq{i}')
+        files.append(f'/{clean_bucket_name}/prq{i}')
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = s.create_table('t1', pa.schema([('num', pa.int64())]))
+        log.info("Starting import of %d files", num_files)
+        t.import_files(files)
+        arrow_table = pa.Table.from_batches(t.select(columns=['num']))
+        assert arrow_table.num_rows == num_rows * num_files
+        arrow_table = pa.Table.from_batches(t.select(columns=['num'], predicate=t['num'] == 100))
+        assert arrow_table.num_rows == num_files
+
+
+def test_create_table_from_files(session, clean_bucket_name, s3):
+    datasets = [
+        {'num': [0],
+         'varch': ['z']},
+        {'num': [1, 2, 3, 4, 5],
+         'varch': ['a', 'b', 'c', 'd', 'e']},
+        {'num': [1, 2, 3, 4, 5],
+         'bool': [True, False, None, None, False],
+         'varch': ['a', 'b', 'c', 'd', 'e']},
+        {'num': [1, 2],
+         'bool': [True, True]},
+        {'varch': ['a', 'b', 'c'],
+         'mismatch': [1, 2, 3]}
+    ]
+    for i, ds in enumerate(datasets):
+        table = pa.Table.from_pydict(ds)
+        with NamedTemporaryFile() as f:
+            pq.write_table(table, f.name)
+            s3.put_object(Bucket=clean_bucket_name, Key=f'prq{i}', Body=f)
+
+    same_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(2)]
+    contained_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(4)]
+    different_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(5)]
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = util.create_table_from_files(s, 't1', contained_schema_files)
+        assert len(t.arrow_schema) == 3
+        assert t.arrow_schema == pa.schema([('num', pa.int64()), ('bool', pa.bool_()), ('varch', pa.string())])
+
+        with pytest.raises(InvalidArgument):
+            util.create_table_from_files(s, 't2', different_schema_files)
+
+        with pytest.raises(InvalidArgument):
+            util.create_table_from_files(s, 't2', contained_schema_files, schema_merge_func=util.strict_schema_merge)
+
+        util.create_table_from_files(s, 't2', different_schema_files, schema_merge_func=util.union_schema_merge)
+        util.create_table_from_files(s, 't3', same_schema_files, schema_merge_func=util.strict_schema_merge)
+
+
+def test_import_name_mismatch_error(session, clean_bucket_name, s3):
+    ds = {'varch': ['a', 'b', 'c'],
+          'invalid_column_name': [1, 2, 3]}
+    prq_name = 'name_mismatch.parquet'
+    table = pa.Table.from_pydict(ds)
+    with NamedTemporaryFile() as f:
+        pq.write_table(table, f.name)
+        s3.put_object(Bucket=clean_bucket_name, Key=prq_name, Body=f)
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = s.create_table('t1', pa.schema([('varch', pa.string()), ('num', pa.int64())]))
+        with pytest.raises(ImportFilesError) as exc:
+            t.import_files([f'/{clean_bucket_name}/{prq_name}'])
+        assert exc.value.error_dict['object_name'] == prq_name
+        assert exc.value.error_dict['res'] == 'TabularMismatchColumnName'
+        assert 'invalid_column_name' in exc.value.error_dict['err_msg']
+
+
+def test_import_type_mismatch_error(session, clean_bucket_name, s3):
+    ds = {'varch': ['a', 'b', 'c'],
+          'num_type_mismatch': [1, 2, 3]}
+    prq_name = 'type_mismatch.parquet'
+    table = pa.Table.from_pydict(ds)
+    with NamedTemporaryFile() as f:
+        pq.write_table(table, f.name)
+        s3.put_object(Bucket=clean_bucket_name, Key=prq_name, Body=f)
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = s.create_table('t1', pa.schema([('varch', pa.string()), ('num_type_mismatch', pa.bool_())]))
+        with pytest.raises(ImportFilesError) as exc:
+            t.import_files([f'/{clean_bucket_name}/{prq_name}'])
+        assert exc.value.error_dict['object_name'] == prq_name
+        assert exc.value.error_dict['res'] == 'TabularMismatchColumnType'
+        assert 'num_type_mismatch' in exc.value.error_dict['err_msg']

vastdb/tests/test_projections.py

@@ -0,0 +1,41 @@
+import pyarrow as pa
+import logging
+
+log = logging.getLogger(__name__)
+
+def test_basic_projections(session, clean_bucket_name):
+    with session.transaction() as tx:
+        s = tx.bucket(clean_bucket_name).create_schema('s1')
+        columns = pa.schema([
+            ('a', pa.int8()),
+            ('b', pa.int16()),
+            ('c', pa.string()),
+            ('d', pa.int16()),
+            ('e', pa.int64()),
+            ('s', pa.struct([('x', pa.int8()), ('y', pa.int16())]))
+        ])
+
+        assert s.tables() == []
+        t = s.create_table('t1', columns)
+        assert s.tables() == [t]
+
+        sorted_columns = ['a']
+        unsorted_columns = ['b']
+        p1 = t.create_projection('p1', sorted_columns, unsorted_columns)
+
+        sorted_columns = ['b']
+        unsorted_columns = ['c', 'd']
+        p2 = t.create_projection('p2', sorted_columns, unsorted_columns)
+
+        projs = t.projections()
+        assert projs == [t.projection('p1'), t.projection('p2')]
+        p1 = t.projection('p1')
+        assert p1.name == 'p1'
+        p2 = t.projection('p2')
+        assert p2.name == 'p2'
+
+        p1.rename('p_new')
+        p2.drop()
+        projs = t.projections()
+        assert len(projs) == 1
+        assert projs[0].name == 'p_new'

vastdb/tests/test_sanity.py

@@ -1,26 +1,45 @@
-import logging
-
-import threading
 from http.server import HTTPServer, BaseHTTPRequestHandler
-from vastdb import api
 from itertools import cycle
+import logging
+import threading
+import contextlib
+
+import pytest
+import requests
+
+import vastdb
+
 
 log = logging.getLogger(__name__)
 
-def test_hello_world(rpc):
-    with rpc.transaction() as tx:
+
+def test_hello_world(session):
+    with session.transaction() as tx:
         assert tx.txid is not None
 
+
+def test_bad_credentials(session):
+    bad_session = vastdb.connect(access='BAD', secret='BAD', endpoint=session.api.url)
+    with pytest.raises(vastdb.errors.Forbidden):
+        with bad_session.transaction():
+            pass
+
+
+def test_bad_endpoint(session):
+    with pytest.raises(requests.exceptions.ConnectionError):
+        vastdb.connect(access='BAD', secret='BAD', endpoint='http://invalid-host-name-for-tests:12345')
+
+
 def test_version_extraction():
     # A list of version and expected version parsed by API
     TEST_CASES = [
-            (None, None), # vast server without version in header
-            ("5", None),                                    # major only is not supported
-            ("5.2", "5.2"),                                 # major.minor
-            ("5.2.0", "5.2.0"),                             # major.minor.patch
-            ("5.2.0.0", "5.2.0.0"),                         # major.minor.patch.protocol
-            ("5.2.0.0 some other things", "5.2.0.0"),       # Test forward comptibility 1
-            ("5.2.0.0.20 some other things", "5.2.0.0"),    # Test forward comptibility 2
+            (None, None),                                   # vast server without version in header
+            ("5", None),                                    # major
+            ("5.2", None),                                  # major.minor
+            ("5.2.0", None),                                # major.minor.patch
+            ("5.2.0.10", "5.2.0.10"),                       # major.minor.patch.protocol
+            ("5.2.0.10 some other things", None),           # suffix
+            ("5.2.0.10.20", None),                          # extra version
     ]
 
     # Mock OPTIONS handle that cycles through the test cases response
@@ -55,9 +74,10 @@ def test_version_extraction():
     server_thread.start()
 
     try:
-        for test_case in TEST_CASES:
-            tester = api.VastdbApi(endpoint=f"http://localhost:{httpd.server_port}", access_key="abc", secret_key="abc")
-            assert tester.vast_version == test_case[1]
+        for _, expected in TEST_CASES:
+            with (pytest.raises(NotImplementedError) if expected is None else contextlib.nullcontext()):
+                s = vastdb.connect(endpoint=f"http://localhost:{httpd.server_port}", access="abc", secret="abc")
+                assert s.api.vast_version == expected
     finally:
         # make sure we shut the server down no matter what
         httpd.shutdown()

vastdb/tests/test_schemas.py

@@ -1,8 +1,8 @@
 import pytest
 
 
-def test_schemas(rpc, clean_bucket_name):
-    with rpc.transaction() as tx:
+def test_schemas(session, clean_bucket_name):
+    with session.transaction() as tx:
         b = tx.bucket(clean_bucket_name)
         assert b.schemas() == []
 
@@ -19,8 +19,8 @@ def test_schemas(rpc, clean_bucket_name):
         assert b.schemas() == []
 
 
-def test_commits_and_rollbacks(rpc, clean_bucket_name):
-    with rpc.transaction() as tx:
+def test_commits_and_rollbacks(session, clean_bucket_name):
+    with session.transaction() as tx:
         b = tx.bucket(clean_bucket_name)
         assert b.schemas() == []
         b.create_schema("s3")
@@ -28,12 +28,18 @@ def test_commits_and_rollbacks(rpc, clean_bucket_name):
         # implicit commit
 
     with pytest.raises(ZeroDivisionError):
-        with rpc.transaction() as tx:
+        with session.transaction() as tx:
             b = tx.bucket(clean_bucket_name)
             b.schema("s3").drop()
             assert b.schemas() == []
             1/0  # rollback schema dropping
 
-    with rpc.transaction() as tx:
+    with session.transaction() as tx:
         b = tx.bucket(clean_bucket_name)
         assert b.schemas() != []
+
+def test_list_snapshots(session, clean_bucket_name):
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.snapshots()
+        assert s == []