vastdb 1.4.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

vastdb/bucket.py

@@ -36,7 +36,7 @@ class Bucket:
         """Get a specific schema (a container of tables) under this bucket."""
         return self._root_schema.schema(name=name, fail_if_missing=fail_if_missing)
 
-    def schemas(self, batch_size=None):
+    def schemas(self, batch_size: Optional[int] = None) -> Iterable["Schema"]:
         """List bucket's schemas."""
         return self._root_schema.schemas(batch_size=batch_size)
 

vastdb/conftest.py

@@ -1,20 +1,42 @@
 import os
 import sqlite3
 from pathlib import Path
+from typing import Iterable
 
 import boto3
 import pytest
 
 import vastdb
 import vastdb.errors
+from vastdb.schema import Schema
+from vastdb.session import Session
+
+AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID"
+AWS_SECRET_ACCESS_KEY = "AWS_SECRET_ACCESS_KEY"
+
+
+def get_aws_cred_from_system(cred_name):
+    orion_dir = Path(os.path.expanduser(os.environ.get("ORION_DIR", "~/orion")))
+    orion_file_path = orion_dir / 'data' / cred_name
+
+    env_value = os.environ.get(cred_name)
+
+    if env_value is not None:
+        return env_value
+
+    if not orion_file_path.exists():
+        return None
+
+    with open(orion_file_path, "r") as f:
+        return f.read().strip()
 
 
 def pytest_addoption(parser):
     parser.addoption("--tabular-bucket-name", help="Name of the S3 bucket with Tabular enabled", default="vastdb")
     parser.addoption("--tabular-access-key", help="Access key with Tabular permissions (AWS_ACCESS_KEY_ID)",
-                     default=os.environ.get("AWS_ACCESS_KEY_ID", None))
+                     default=get_aws_cred_from_system(AWS_ACCESS_KEY_ID))
     parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions (AWS_SECRET_ACCESS_KEY)",
-                     default=os.environ.get("AWS_SECRET_ACCESS_KEY", None))
+                     default=get_aws_cred_from_system(AWS_SECRET_ACCESS_KEY))
     parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint", default=[], action="append")
     parser.addoption("--data-path", help="Data files location", default=None)
     parser.addoption("--crater-path", help="Save benchmark results in a dedicated location", default=None)
@@ -24,7 +46,7 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="session")
-def session_kwargs(request, tabular_endpoint_urls):
+def session_kwargs(request: pytest.FixtureRequest, tabular_endpoint_urls):
     return dict(
         access=request.config.getoption("--tabular-access-key"),
         secret=request.config.getoption("--tabular-secret-key"),
@@ -38,30 +60,30 @@ def session(session_kwargs):
 
 
 @pytest.fixture(scope="session")
-def num_workers(request):
+def num_workers(request: pytest.FixtureRequest):
     return int(request.config.getoption("--num-workers"))
 
 
 @pytest.fixture(scope="session")
-def test_bucket_name(request):
+def test_bucket_name(request: pytest.FixtureRequest):
     return request.config.getoption("--tabular-bucket-name")
 
 
 @pytest.fixture(scope="session")
-def tabular_endpoint_urls(request):
+def tabular_endpoint_urls(request: pytest.FixtureRequest):
     return request.config.getoption("--tabular-endpoint-url") or ["http://localhost:9090"]
 
 
-def iter_schemas(s):
-    """Recusively scan all schemas."""
-    children = s.schemas()
-    for c in children:
-        yield from iter_schemas(c)
-    yield s
+def iter_schemas(schema: Schema) -> Iterable[Schema]:
+    """Recursively scan all schemas."""
+    children = schema.schemas()
+    for child in children:
+        yield from iter_schemas(child)
+    yield schema
 
 
 @pytest.fixture(scope="function")
-def clean_bucket_name(request, test_bucket_name, session):
+def clean_bucket_name(request: pytest.FixtureRequest, test_bucket_name: str, session: Session) -> str:
     with session.transaction() as tx:
         b = tx.bucket(test_bucket_name)
         for top_schema in b.schemas():
@@ -69,7 +91,8 @@ def clean_bucket_name(request, test_bucket_name, session):
                 for t_name in s.tablenames():
                     try:
                         t = s.table(t_name)
-                        t.drop()
+                        if t is not None:
+                            t.drop()
                     except vastdb.errors.NotSupportedSchema:
                         # Use internal API to drop the table in case unsupported schema prevents creating a table
                         # object.
@@ -79,7 +102,7 @@ def clean_bucket_name(request, test_bucket_name, session):
 
 
 @pytest.fixture(scope="session")
-def s3(request, tabular_endpoint_urls):
+def s3(request: pytest.FixtureRequest, tabular_endpoint_urls):
     return boto3.client(
         's3',
         aws_access_key_id=request.config.getoption("--tabular-access-key"),
@@ -88,22 +111,22 @@ def s3(request, tabular_endpoint_urls):
 
 
 @pytest.fixture(scope="function")
-def parquets_path(request):
+def parquets_path(request: pytest.FixtureRequest):
     return Path(request.config.getoption("--data-path"))
 
 
 @pytest.fixture(scope="function")
-def crater_path(request):
+def crater_path(request: pytest.FixtureRequest):
     return request.config.getoption("--crater-path")
 
 
 @pytest.fixture(scope="function")
-def schema_name(request):
+def schema_name(request: pytest.FixtureRequest):
     return request.config.getoption("--schema-name")
 
 
 @pytest.fixture(scope="function")
-def table_name(request):
+def table_name(request: pytest.FixtureRequest):
     return request.config.getoption("--table-name")
 
 

vastdb/schema.py

@@ -10,7 +10,10 @@ from typing import TYPE_CHECKING, Iterable, List, Optional
 
 import pyarrow as pa
 
+from vastdb.table_metadata import TableMetadata, TableRef, TableType
+
 from . import bucket, errors, schema, table
+from ._ibis_support import validate_ibis_support_schema
 
 if TYPE_CHECKING:
     from .table import Table
@@ -91,7 +94,7 @@ class Schema:
         if use_external_row_ids_allocation:
             self.tx._rpc.features.check_external_row_ids_allocation()
 
-        table.Table.validate_ibis_support_schema(columns)
+        validate_ibis_support_schema(columns)
         self.tx._rpc.api.create_table(self.bucket.name, self.name, table_name, columns, txid=self.tx.txid,
                                       use_external_row_ids_allocation=use_external_row_ids_allocation,
                                       sorting_key=sorting_key)
@@ -117,7 +120,7 @@ class Schema:
         while True:
             _bucket_name, _schema_name, curr_tables, next_key, is_truncated, _ = \
                 self.tx._rpc.api.list_tables(
-                    bucket=self.bucket.name, schema=self.name, next_key=next_key, max_keys=page_size, txid=self.tx.txid,
+                    bucket=self.bucket.name, schema=self.name, next_key=next_key, max_keys=page_size, txid=self.tx.active_txid,
                     exact_match=exact_match, name_prefix=name_prefix, include_list_stats=exact_match)
             if not curr_tables:
                 break
@@ -152,4 +155,13 @@ class Schema:
 
 
 def _parse_table_info(table_info, schema: "schema.Schema"):
-    return table.Table(name=table_info.name, schema=schema, handle=int(table_info.handle), _imports_table=False, sorted_table=table_info.sorting_key_enabled)
+    ref = TableRef(bucket=schema.bucket.name,
+                   schema=schema.name,
+                   table=table_info.name)
+
+    table_type = TableType.Elysium if table_info.sorting_key_enabled else TableType.Regular
+    table_metadata = TableMetadata(ref, table_type=table_type)
+
+    return table.Table(handle=int(table_info.handle),
+                       metadata=table_metadata,
+                       tx=schema.tx)

vastdb/session.py

@@ -10,6 +10,8 @@ For more details see:
 import os
 from typing import TYPE_CHECKING, Optional
 
+from vastdb.transaction import Transaction
+
 if TYPE_CHECKING:
     from .config import BackoffConfig
 
@@ -45,7 +47,7 @@ class Session:
         """Don't show the secret key."""
         return f'{self.__class__.__name__}(endpoint={self.api.url}, access={self.api.access_key})'
 
-    def transaction(self):
+    def transaction(self) -> Transaction:
         """Create a non-initialized transaction object.
 
         It should be used as a context manager: