pixeltable 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

pixeltable/dataframe.py

@@ -88,12 +88,12 @@ class DataFrameResultSet:
     def __iter__(self) -> Iterator[dict[str, Any]]:
         return (self._row_to_dict(i) for i in range(len(self)))
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, DataFrameResultSet):
             return False
         return self.to_pandas().equals(other.to_pandas())
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(self.to_pandas())
 
 
@@ -571,7 +571,7 @@ class DataFrame:
             expr = exprs.Expr.from_object(raw_expr)
             if expr is None:
                 raise excs.Error(f'Invalid expression: {raw_expr}')
-            if expr.col_type.is_invalid_type():
+            if expr.col_type.is_invalid_type() and not (isinstance(expr, exprs.Literal) and expr.val is None):
                 raise excs.Error(f'Invalid type: {raw_expr}')
             if not expr.is_bound_by(self._from_clause.tbls):
                 raise excs.Error(
@@ -624,6 +624,8 @@ class DataFrame:
 
             >>> df = person.where(t.age > 30)
         """
+        if self.where_clause is not None:
+            raise excs.Error('Where clause already specified')
         if not isinstance(pred, exprs.Expr):
             raise excs.Error(f'Where() requires a Pixeltable expression, but instead got {type(pred)}')
         if not pred.col_type.is_bool_type():

pixeltable/env.py

@@ -30,6 +30,7 @@ from tqdm import TqdmWarning
 from pixeltable import exceptions as excs
 from pixeltable.config import Config
 from pixeltable.utils.console_output import ConsoleLogger, ConsoleMessageFilter, ConsoleOutputHandler, map_level
+from pixeltable.utils.dbms import CockroachDbms, Dbms, PostgresqlDbms
 from pixeltable.utils.http_server import make_server
 
 if TYPE_CHECKING:
@@ -43,7 +44,9 @@ T = TypeVar('T')
 
 class Env:
     """
-    Store for runtime globals.
+    Store runtime globals for both local and non-local environments.
+    For a local environment, Pixeltable uses an embedded PostgreSQL server that runs locally in a separate process.
+    For a non-local environment, Pixeltable uses a connection string to the externally managed database.
     """
 
     _instance: Optional[Env] = None
@@ -58,7 +61,7 @@ class Env:
     _sa_engine: Optional[sql.engine.base.Engine]
     _pgdata_dir: Optional[Path]
     _db_name: Optional[str]
-    _db_server: Optional[pixeltable_pgserver.PostgresServer]
+    _db_server: Optional[pixeltable_pgserver.PostgresServer]  # set only when running in local environment
     _db_url: Optional[str]
     _default_time_zone: Optional[ZoneInfo]
 
@@ -81,6 +84,7 @@ class Env:
     _resource_pool_info: dict[str, Any]
     _current_conn: Optional[sql.Connection]
     _current_session: Optional[sql.orm.Session]
+    _dbms: Optional[Dbms]
 
     @classmethod
     def get(cls) -> Env:
@@ -98,7 +102,7 @@ class Env:
         cls._instance = env
         cls.__initializing = False
 
-    def __init__(self):
+    def __init__(self) -> None:
         assert self._instance is None, 'Env is a singleton; use Env.get() to access the instance'
 
         self._media_dir = None  # computed media files
@@ -112,7 +116,6 @@ class Env:
         self._db_server = None
         self._db_url = None
         self._default_time_zone = None
-
         self.__optional_packages = {}
         self._spacy_nlp = None
         self._httpd = None
@@ -136,6 +139,7 @@ class Env:
         self._resource_pool_info = {}
         self._current_conn = None
         self._current_session = None
+        self._dbms = None
 
     @property
     def db_url(self) -> str:
@@ -147,6 +151,18 @@ class Env:
         assert self._http_address is not None
         return self._http_address
 
+    @property
+    def user(self) -> Optional[str]:
+        return Config.get().get_string_value('user')
+
+    @user.setter
+    def user(self, user: Optional[str]) -> None:
+        if user is None:
+            if 'PIXELTABLE_USER' in os.environ:
+                del os.environ['PIXELTABLE_USER']
+        else:
+            os.environ['PIXELTABLE_USER'] = user
+
     @property
     def default_time_zone(self) -> Optional[ZoneInfo]:
         return self._default_time_zone
@@ -170,9 +186,19 @@ class Env:
         assert self._current_session is not None
         return self._current_session
 
+    @property
+    def dbms(self) -> Optional[Dbms]:
+        assert self._dbms is not None
+        return self._dbms
+
     def in_xact(self) -> bool:
         return self._current_conn is not None
 
+    @property
+    def is_local(self) -> bool:
+        assert self._db_url is not None  # is_local should be called only after db initialization
+        return self._db_server is not None
+
     @contextmanager
     def begin_xact(self) -> Iterator[sql.Connection]:
         """Return a context manager that yields a connection to the database. Idempotent."""
@@ -346,16 +372,13 @@ class Env:
 
         self.clear_tmp_dir()
 
-        self._db_name = os.environ.get('PIXELTABLE_DB', 'pixeltable')
-        self._pgdata_dir = Path(os.environ.get('PIXELTABLE_PGDATA', str(Config.get().home / 'pgdata')))
+        # configure pixeltable database
+        self._init_db(config)
 
-        # cleanup_mode=None will leave the postgres process running after Python exits
-        # cleanup_mode='stop' will terminate the postgres process when Python exits
-        # On Windows, we need cleanup_mode='stop' because child processes are killed automatically when the parent
-        # process (such as Terminal or VSCode) exits, potentially leaving it in an unusable state.
-        cleanup_mode = 'stop' if platform.system() == 'Windows' else None
-        self._db_server = pixeltable_pgserver.get_server(self._pgdata_dir, cleanup_mode=cleanup_mode)
-        self._db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
+        if reinit_db and not self.is_local:
+            raise excs.Error(
+                'Reinitializing pixeltable database is not supported when running in non-local environment'
+            )
 
         tz_name = config.get_string_value('time_zone')
         if tz_name is not None:
@@ -382,11 +405,8 @@ class Env:
         # Create the SQLAlchemy engine. This will also set the default time zone.
         self._create_engine(time_zone_name=tz_name, echo=echo)
 
-        if create_db:
-            from pixeltable import metadata
-
-            metadata.schema.base_metadata.create_all(self._sa_engine)
-            metadata.create_system_info(self._sa_engine)
+        # Create catalog tables and system metadata
+        self._init_metadata()
 
         self.console_logger.info(f'Connected to Pixeltable database at: {self.db_url}')
 
@@ -394,12 +414,65 @@ class Env:
         self._set_up_runtime()
         self.log_to_stdout(False)
 
+    def _init_db(self, config: Config) -> None:
+        """
+        Initialize the pixeltable database along with its associated DBMS.
+        """
+        db_connect_str = config.get_string_value('DB_CONNECT_STR')
+        if db_connect_str is not None:
+            try:
+                db_url = sql.make_url(db_connect_str)
+            except sql.exc.ArgumentError as e:
+                error = f'Invalid db connection string {db_connect_str}: {e}'
+                self._logger.error(error)
+                raise excs.Error(error) from e
+            self._db_url = db_url.render_as_string(hide_password=False)
+            self._db_name = db_url.database  # use the dbname given in connect string
+            dialect = db_url.get_dialect().name
+            if dialect == 'cockroachdb':
+                self._dbms = CockroachDbms(db_url)
+            else:
+                raise excs.Error(f'Unsupported DBMS {dialect}')
+            # Check if database exists
+            if not self._store_db_exists():
+                error = f'Database {self._db_name!r} does not exist'
+                self._logger.error(error)
+                raise excs.Error(error)
+            self._logger.info(f'Using database at: {self.db_url}')
+        else:
+            self._db_name = os.environ.get('PIXELTABLE_DB', 'pixeltable')
+            self._pgdata_dir = Path(os.environ.get('PIXELTABLE_PGDATA', str(Config.get().home / 'pgdata')))
+            # cleanup_mode=None will leave the postgres process running after Python exits
+            # cleanup_mode='stop' will terminate the postgres process when Python exits
+            # On Windows, we need cleanup_mode='stop' because child processes are killed automatically when the parent
+            # process (such as Terminal or VSCode) exits, potentially leaving it in an unusable state.
+            cleanup_mode = 'stop' if platform.system() == 'Windows' else None
+            self._db_server = pixeltable_pgserver.get_server(self._pgdata_dir, cleanup_mode=cleanup_mode)
+            self._db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
+            self._dbms = PostgresqlDbms(sql.make_url(self._db_url))
+        assert self._dbms is not None
+        assert self._db_url is not None
+        assert self._db_name is not None
+
+    def _init_metadata(self) -> None:
+        """
+        Create pixeltable metadata tables and system metadata.
+        This is an idempotent operation.
+        """
+        assert self._sa_engine is not None
+        from pixeltable import metadata
+
+        metadata.schema.base_metadata.create_all(self._sa_engine, checkfirst=True)
+        metadata.create_system_info(self._sa_engine)
+
     def _create_engine(self, time_zone_name: Optional[str], echo: bool = False) -> None:
         connect_args = {} if time_zone_name is None else {'options': f'-c timezone={time_zone_name}'}
         self._sa_engine = sql.create_engine(
-            self.db_url, echo=echo, isolation_level='REPEATABLE READ', connect_args=connect_args
+            self.db_url, echo=echo, isolation_level=self._dbms.transaction_isolation_level, connect_args=connect_args
         )
+
         self._logger.info(f'Created SQLAlchemy engine at: {self.db_url}')
+
         with self.engine.begin() as conn:
             tz_name = conn.execute(sql.text('SHOW TIME ZONE')).scalar()
             assert isinstance(tz_name, str)
@@ -409,8 +482,7 @@ class Env:
     def _store_db_exists(self) -> bool:
         assert self._db_name is not None
         # don't try to connect to self.db_name, it may not exist
-        db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
-        engine = sql.create_engine(db_url, future=True)
+        engine = sql.create_engine(self._dbms.default_system_db_url(), future=True)
         try:
             with engine.begin() as conn:
                 stmt = f"SELECT COUNT(*) FROM pg_database WHERE datname = '{self._db_name}'"
@@ -423,23 +495,17 @@ class Env:
     def _create_store_db(self) -> None:
         assert self._db_name is not None
         # create the db
-        pg_db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
-        engine = sql.create_engine(pg_db_url, future=True, isolation_level='AUTOCOMMIT')
+        engine = sql.create_engine(self._dbms.default_system_db_url(), future=True, isolation_level='AUTOCOMMIT')
         preparer = engine.dialect.identifier_preparer
         try:
             with engine.begin() as conn:
-                # use C collation to get standard C/Python-style sorting
-                stmt = (
-                    f'CREATE DATABASE {preparer.quote(self._db_name)} '
-                    "ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
-                )
+                stmt = self._dbms.create_db_stmt(preparer.quote(self._db_name))
                 conn.execute(sql.text(stmt))
         finally:
             engine.dispose()
 
         # enable pgvector
-        store_db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
-        engine = sql.create_engine(store_db_url, future=True, isolation_level='AUTOCOMMIT')
+        engine = sql.create_engine(self.db_url, future=True, isolation_level='AUTOCOMMIT')
         try:
             with engine.begin() as conn:
                 conn.execute(sql.text('CREATE EXTENSION vector'))
@@ -448,21 +514,21 @@ class Env:
 
     def _drop_store_db(self) -> None:
         assert self._db_name is not None
-        db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
-        engine = sql.create_engine(db_url, future=True, isolation_level='AUTOCOMMIT')
+        engine = sql.create_engine(self._dbms.default_system_db_url(), future=True, isolation_level='AUTOCOMMIT')
         preparer = engine.dialect.identifier_preparer
         try:
             with engine.begin() as conn:
                 # terminate active connections
-                stmt = f"""
-                    SELECT pg_terminate_backend(pg_stat_activity.pid)
-                    FROM pg_stat_activity
-                    WHERE pg_stat_activity.datname = '{self._db_name}'
-                    AND pid <> pg_backend_pid()
-                """
-                conn.execute(sql.text(stmt))
+                if self._db_server is not None:
+                    stmt = f"""
+                        SELECT pg_terminate_backend(pg_stat_activity.pid)
+                        FROM pg_stat_activity
+                        WHERE pg_stat_activity.datname = '{self._db_name}'
+                        AND pid <> pg_backend_pid()
+                    """
+                    conn.execute(sql.text(stmt))
                 # drop db
-                stmt = f'DROP DATABASE {preparer.quote(self._db_name)}'
+                stmt = self._dbms.drop_db_stmt(preparer.quote(self._db_name))
                 conn.execute(sql.text(stmt))
         finally:
             engine.dispose()
@@ -524,7 +590,7 @@ class Env:
         port = self._httpd.server_address[1]
         self._http_address = f'http://127.0.0.1:{port}'
 
-        def run_server():
+        def run_server() -> None:
             logging.log(logging.INFO, f'running web server at {self._http_address}')
             self._httpd.serve_forever()
 
@@ -779,7 +845,7 @@ class RateLimitsInfo:
     def reset(self) -> None:
         self.resource_limits.clear()
 
-    def record(self, **kwargs) -> None:
+    def record(self, **kwargs: Any) -> None:
         now = datetime.datetime.now(tz=datetime.timezone.utc)
         if len(self.resource_limits) == 0:
             self.resource_limits = {k: RateLimitInfo(k, now, *v) for k, v in kwargs.items() if v is not None}

pixeltable/exec/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401
+
 from .aggregation_node import AggregationNode
 from .cache_prefetch_node import CachePrefetchNode
 from .component_iteration_node import ComponentIterationNode

pixeltable/exec/aggregation_node.py

@@ -2,11 +2,9 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Any, AsyncIterator, Iterable, Iterator, Optional, cast
+from typing import Any, AsyncIterator, Iterable, Optional, cast
 
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exceptions as excs, exprs
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -52,20 +50,20 @@ class AggregationNode(ExecNode):
         for fn_call in self.agg_fn_calls:
             try:
                 fn_call.reset_agg()
-            except Exception as e:
+            except Exception as exc:
                 _, _, exc_tb = sys.exc_info()
                 expr_msg = f'init() function of the aggregate {fn_call}'
-                raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, [], row_num)
+                raise excs.ExprEvalError(fn_call, expr_msg, exc, exc_tb, [], row_num) from exc
 
     def _update_agg_state(self, row: exprs.DataRow, row_num: int) -> None:
         for fn_call in self.agg_fn_calls:
             try:
                 fn_call.update(row)
-            except Exception as e:
+            except Exception as exc:
                 _, _, exc_tb = sys.exc_info()
                 expr_msg = f'update() function of the aggregate {fn_call}'
                 input_vals = [row[d.slot_idx] for d in fn_call.dependencies()]
-                raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, input_vals, row_num)
+                raise excs.ExprEvalError(fn_call, expr_msg, exc, exc_tb, input_vals, row_num) from exc
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         prev_row: Optional[exprs.DataRow] = None

pixeltable/exec/cache_prefetch_node.py

@@ -12,10 +12,7 @@ from pathlib import Path
 from typing import Any, AsyncIterator, Iterator, Optional
 from uuid import UUID
 
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
-from pixeltable import catalog
+from pixeltable import catalog, env, exceptions as excs, exprs
 from pixeltable.utils.filecache import FileCache
 
 from .data_row_batch import DataRowBatch
@@ -234,7 +231,7 @@ class CachePrefetchNode(ExecNode):
         assert len(parsed.scheme) > 1 and parsed.scheme != 'file'
         # preserve the file extension, if there is one
         extension = ''
-        if parsed.path != '':
+        if parsed.path:
             p = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
             extension = p.suffix
         tmp_path = env.Env.get().create_tmp_path(extension=extension)
@@ -253,12 +250,12 @@ class CachePrefetchNode(ExecNode):
                         }
                         self.boto_client = get_client(**config)
                 self.boto_client.download_file(parsed.netloc, parsed.path.lstrip('/'), str(tmp_path))
-            elif parsed.scheme == 'http' or parsed.scheme == 'https':
+            elif parsed.scheme in ('http', 'https'):
                 with urllib.request.urlopen(url) as resp, open(tmp_path, 'wb') as f:
                     data = resp.read()
                     f.write(data)
             else:
-                assert False, f'Unsupported URL scheme: {parsed.scheme}'
+                raise AssertionError(f'Unsupported URL scheme: {parsed.scheme}')
             _logger.debug(f'Downloaded {url} to {tmp_path}')
             return tmp_path, None
         except Exception as e:

pixeltable/exec/component_iteration_node.py

@@ -1,8 +1,6 @@
 from typing import AsyncIterator
 
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exceptions as excs, exprs
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode

pixeltable/exec/data_row_batch.py

@@ -3,8 +3,7 @@ from __future__ import annotations
 import logging
 from typing import Iterator, Optional
 
-import pixeltable.catalog as catalog
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exprs
 from pixeltable.utils.media_store import MediaStore
 
 _logger = logging.getLogger('pixeltable')

pixeltable/exec/exec_context.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import sqlalchemy as sql
 
-import pixeltable.exprs as exprs
+from pixeltable import exprs
 
 
 class ExecContext:

pixeltable/exec/exec_node.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import abc
 import asyncio
 import logging
-import sys
 from typing import AsyncIterator, Iterable, Iterator, Optional, TypeVar
 
 from pixeltable import exprs
@@ -70,7 +69,7 @@ class ExecNode(abc.ABC):
 
             nest_asyncio.apply()
             loop = running_loop
-            _logger.debug(f'Patched running loop')
+            _logger.debug('Patched running loop')
         except RuntimeError:
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)

pixeltable/exec/expr_eval/__init__.py

@@ -1 +1,3 @@
+# ruff: noqa: F401
+
 from .expr_eval_node import ExprEvalNode