vastdb 0.0.5.2__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

@@ -0,0 +1,83 @@
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from itertools import cycle
+import logging
+import threading
+import contextlib
+
+import pytest
+import requests
+
+import vastdb
+
+
+log = logging.getLogger(__name__)
+
+
+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
+            ("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
+    class MockOptionsHandler(BaseHTTPRequestHandler):
+        versions_iterator = cycle(TEST_CASES)
+
+        def __init__(self, *args) -> None:
+            super().__init__(*args)
+
+        def do_OPTIONS(self):
+            self.send_response(204)
+            self.end_headers()
+
+        def version_string(self):
+            version = next(self.versions_iterator)[0]
+            return f"vast {version}" if version else "vast"
+
+        def log_message(self, format, *args):
+            log.debug(format,*args)
+
+    # start the server on localhost on some available port port
+    server_address =('localhost', 0)
+    httpd = HTTPServer(server_address, MockOptionsHandler)
+
+    def start_http_server_in_thread():
+        log.info(f"Mock HTTP server is running on port {httpd.server_port}")
+        httpd.serve_forever()
+        log.info("Mock HTTP server killed")
+
+    # start the server in a thread so we have the main thread to operate the API
+    server_thread = threading.Thread(target=start_http_server_in_thread)
+    server_thread.start()
+
+    try:
+        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

@@ -0,0 +1,45 @@
+import pytest
+
+
+def test_schemas(session, clean_bucket_name):
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        assert b.schemas() == []
+
+        s = b.create_schema('s1')
+        assert s.bucket == b
+        assert b.schemas() == [s]
+
+        s.rename('s2')
+        assert s.bucket == b
+        assert s.name == 's2'
+        assert b.schemas()[0].name == 's2'
+
+        s.drop()
+        assert b.schemas() == []
+
+
+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")
+        assert b.schemas() != []
+        # implicit commit
+
+    with pytest.raises(ZeroDivisionError):
+        with session.transaction() as tx:
+            b = tx.bucket(clean_bucket_name)
+            b.schema("s3").drop()
+            assert b.schemas() == []
+            1/0  # rollback schema dropping
+
+    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 == []