datachain 0.30.1__tar.gz → 0.30.2__tar.gz

{datachain-0.30.1 → datachain-0.30.2}/.github/workflows/tests.yml

@@ -78,6 +78,24 @@ jobs:
           fetch-depth: 0
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
 
+      - name: Setup PostgreSQL
+        if: runner.os != 'Windows'
+        uses: ikalnytskyi/action-setup-postgres@10ab8a56cc77b4823c2bfa57b1d4dd5605ef0481 # v7
+        with:
+          username: test
+          password: test
+          database: test_datachain
+          port: 5432
+          postgres-version: "17"
+        id: postgres
+
+      - name: Set PostgreSQL URI
+        if: runner.os != 'Windows'
+        run: |
+          FULL_URI="${{ steps.postgres.outputs.connection-uri }}"
+          echo "TEST_POSTGRES_URI=${FULL_URI%/*}" >> "$GITHUB_ENV"
+        shell: bash
+
       - name: Set up Python ${{ matrix.pyv }}
         uses: actions/setup-python@v5
         with:

{datachain-0.30.1 → datachain-0.30.2}/.gitignore

@@ -145,3 +145,7 @@ cython_debug/
 *.pt
 
 .DS_Store/
+
+# for local dev, e.g. LLM generated files, .env.test to override
+# test variables, local scripts to try, etc
+local/

{datachain-0.30.1 → datachain-0.30.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.30.1
+Version: 0.30.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -81,8 +81,10 @@ Provides-Extra: video
 Requires-Dist: ffmpeg-python; extra == "video"
 Requires-Dist: imageio[ffmpeg,pyav]>=2.37.0; extra == "video"
 Requires-Dist: opencv-python; extra == "video"
+Provides-Extra: postgres
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "postgres"
 Provides-Extra: tests
-Requires-Dist: datachain[audio,hf,remote,torch,vector,video]; extra == "tests"
+Requires-Dist: datachain[audio,hf,postgres,remote,torch,vector,video]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
@@ -90,6 +92,7 @@ Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
 Requires-Dist: pytest-servers[all]>=0.5.9; extra == "tests"
 Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
 Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
+Requires-Dist: pytest-env>=1.1.0; extra == "tests"
 Requires-Dist: virtualenv; extra == "tests"
 Requires-Dist: dulwich; extra == "tests"
 Requires-Dist: hypothesis; extra == "tests"

{datachain-0.30.1 → datachain-0.30.2}/pyproject.toml

@@ -96,8 +96,11 @@ video = [
   "imageio[ffmpeg,pyav]>=2.37.0",
   "opencv-python"
 ]
+postgres = [
+  "psycopg2-binary>=2.9.0"
+]
 tests = [
-  "datachain[torch,audio,remote,vector,hf,video]",
+  "datachain[torch,audio,remote,vector,hf,video,postgres]",
   "pytest>=8,<9",
   "pytest-sugar>=0.9.6",
   "pytest-cov>=4.1.0",
@@ -105,6 +108,7 @@ tests = [
   "pytest-servers[all]>=0.5.9",
   "pytest-benchmark[histogram]",
   "pytest-xdist>=3.3.1",
+  "pytest-env>=1.1.0",
   "virtualenv",
   "dulwich",
   "hypothesis",
@@ -150,13 +154,16 @@ namespaces = false
 
 [tool.pytest.ini_options]
 addopts = "-rfEs -m 'not examples' --benchmark-skip"
+env_override_existing_values = true
+env_files = "local/.env.test"
 markers = [
   "e2e: End-to-end tests",
   "examples: All examples",
   "computer_vision: Computer vision examples",
   "get_started: Get started examples",
   "llm_and_nlp: LLM and NLP examples",
-  "multimodal: Multimodal examples"
+  "multimodal: Multimodal examples",
+  "incremental_processing: Delta and retry examples"
 ]
 filterwarnings = [
   "error::pandas.errors.PerformanceWarning",

{datachain-0.30.1 → datachain-0.30.2}/src/datachain/lib/dc/database.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 import sqlalchemy
 
 from datachain.query.schema import ColumnMeta
+from datachain.utils import batched
 
 DEFAULT_DATABASE_BATCH_SIZE = 10_000
 
@@ -74,6 +75,7 @@ def to_database(
     *,
     batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
     on_conflict: Optional[str] = None,
+    conflict_columns: Optional[list[str]] = None,
     column_mapping: Optional[dict[str, Optional[str]]] = None,
 ) -> None:
     """
@@ -82,8 +84,6 @@ def to_database(
     This is the core implementation that handles the actual database operations.
     For user-facing documentation, see DataChain.to_database() method.
     """
-    from datachain.utils import batched
-
     if on_conflict and on_conflict not in ("ignore", "update"):
         raise ValueError(
             f"on_conflict must be 'ignore' or 'update', got: {on_conflict}"
@@ -105,19 +105,26 @@ def to_database(
         metadata = sqlalchemy.MetaData()
         table = sqlalchemy.Table(table_name, metadata, *columns)
 
-        # Check if table already exists to determine if we should clean up on error.
-        inspector = sqlalchemy.inspect(conn)
-        assert inspector  # to satisfy mypy
-        table_existed_before = table_name in inspector.get_table_names()
-
+        table_existed_before = False
         try:
-            table.create(conn, checkfirst=True)
-            rows_iter = chain._leaf_values()
-            for batch in batched(rows_iter, batch_rows):
-                _process_batch(
-                    conn, table, batch, on_conflict, column_indices_and_names
-                )
-            conn.commit()
+            with conn.begin():
+                # Check if table exists to determine if we should clean up on error.
+                inspector = sqlalchemy.inspect(conn)
+                assert inspector  # to satisfy mypy
+                table_existed_before = table_name in inspector.get_table_names()
+
+                table.create(conn, checkfirst=True)
+
+                rows_iter = chain._leaf_values()
+                for batch in batched(rows_iter, batch_rows):
+                    _process_batch(
+                        conn,
+                        table,
+                        batch,
+                        on_conflict,
+                        conflict_columns,
+                        column_indices_and_names,
+                    )
         except Exception:
             if not table_existed_before:
                 try:
@@ -183,7 +190,9 @@ def _prepare_columns(all_columns, column_mapping):
     return column_indices_and_names, columns
 
 
-def _process_batch(conn, table, batch, on_conflict, column_indices_and_names):
+def _process_batch(
+    conn, table, batch, on_conflict, conflict_columns, column_indices_and_names
+):
     """Process a batch of rows with conflict resolution."""
 
     def prepare_row(row_values):
@@ -217,7 +226,19 @@ def _process_batch(conn, table, batch, on_conflict, column_indices_and_names):
             update_values = {
                 col.name: insert_stmt.excluded[col.name] for col in table.columns
             }
-            insert_stmt = insert_stmt.on_conflict_do_update(set_=update_values)
+            if conn.engine.name == "postgresql":
+                if not conflict_columns:
+                    raise ValueError(
+                        "conflict_columns parameter is required when "
+                        "on_conflict='update' with PostgreSQL. Specify the column "
+                        "names that form a unique constraint."
+                    )
+
+                insert_stmt = insert_stmt.on_conflict_do_update(
+                    index_elements=conflict_columns, set_=update_values
+                )
+            else:
+                insert_stmt = insert_stmt.on_conflict_do_update(set_=update_values)
     elif on_conflict:
         import warnings
 

{datachain-0.30.1 → datachain-0.30.2}/src/datachain/lib/dc/datachain.py

@@ -2296,6 +2296,7 @@ class DataChain:
         *,
         batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
         on_conflict: Optional[str] = None,
+        conflict_columns: Optional[list[str]] = None,
         column_mapping: Optional[dict[str, Optional[str]]] = None,
     ) -> None:
         """Save chain to a database table using a given database connection.
@@ -2319,6 +2320,9 @@ class DataChain:
                   (default)
                 - "ignore": Skip duplicate rows silently
                 - "update": Update existing rows with new values
+            conflict_columns: List of column names that form a unique constraint
+                for conflict resolution. Required when on_conflict='update' and
+                using PostgreSQL.
             column_mapping: Optional mapping to rename or skip columns:
                 - Dict mapping DataChain column names to database column names
                 - Set values to None to skip columns entirely, or use `defaultdict` to
@@ -2377,6 +2381,7 @@ class DataChain:
             connection,
             batch_rows=batch_rows,
             on_conflict=on_conflict,
+            conflict_columns=conflict_columns,
             column_mapping=column_mapping,
         )
 

{datachain-0.30.1 → datachain-0.30.2}/src/datachain/sql/__init__.py

@@ -1,6 +1,8 @@
 from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import column
 
+# Import PostgreSQL dialect registration (registers PostgreSQL type converter)
+from . import postgresql_dialect  # noqa: F401
 from .default import setup as default_setup
 from .selectable import select, values
 

datachain-0.30.2/src/datachain/sql/postgresql_dialect.py

@@ -0,0 +1,9 @@
+"""
+PostgreSQL dialect registration for DataChain.
+"""
+
+from datachain.sql.postgresql_types import PostgreSQLTypeConverter
+from datachain.sql.types import register_backend_types
+
+# Register PostgreSQL type converter
+register_backend_types("postgresql", PostgreSQLTypeConverter())

datachain-0.30.2/src/datachain/sql/postgresql_types.py

@@ -0,0 +1,21 @@
+"""
+PostgreSQL-specific type converter for DataChain.
+
+Handles PostgreSQL-specific type mappings that differ from the default dialect.
+"""
+
+from sqlalchemy.dialects import postgresql
+
+from datachain.sql.types import TypeConverter
+
+
+class PostgreSQLTypeConverter(TypeConverter):
+    """PostgreSQL-specific type converter."""
+
+    def datetime(self):
+        """PostgreSQL uses TIMESTAMP WITH TIME ZONE to preserve timezone information."""
+        return postgresql.TIMESTAMP(timezone=True)
+
+    def json(self):
+        """PostgreSQL uses JSONB for better performance and query capabilities."""
+        return postgresql.JSONB()

datachain-0.30.2/src/datachain/sql/sqlite/__init__.py

@@ -0,0 +1,11 @@
+from .base import (
+    create_user_defined_sql_functions,
+    setup,
+    sqlite_dialect,
+)
+
+__all__ = [
+    "create_user_defined_sql_functions",
+    "setup",
+    "sqlite_dialect",
+]

{datachain-0.30.1 → datachain-0.30.2}/src/datachain/sql/sqlite/base.py

@@ -304,7 +304,11 @@ def register_user_defined_sql_functions() -> None:
 
 
 def adapt_datetime(val: datetime) -> str:
-    if not (val.tzinfo is timezone.utc or val.tzname() == "UTC"):
+    is_utc_check = val.tzinfo is timezone.utc
+    tzname_check = val.tzname() == "UTC"
+    combined_check = is_utc_check or tzname_check
+
+    if not combined_check:
         try:
             val = val.astimezone(timezone.utc)
         except (OverflowError, ValueError, OSError):
@@ -314,6 +318,7 @@ def adapt_datetime(val: datetime) -> str:
                 val = datetime.min.replace(tzinfo=timezone.utc)
             else:
                 raise
+
     return val.replace(tzinfo=None).isoformat(" ")
 
 

{datachain-0.30.1 → datachain-0.30.2}/src/datachain/sql/types.py

@@ -58,9 +58,14 @@ def converter(dialect) -> "TypeConverter":
     try:
         return registry[name]
     except KeyError:
-        raise ValueError(
-            f"No type converter registered for dialect: {dialect.name!r}"
-        ) from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No type converter registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def read_converter(dialect) -> "TypeReadConverter":
@@ -68,9 +73,14 @@ def read_converter(dialect) -> "TypeReadConverter":
     try:
         return read_converter_registry[name]
     except KeyError:
-        raise ValueError(
-            f"No read type converter registered for dialect: {dialect.name!r}"
-        ) from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return read_converter_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No read type converter registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def type_defaults(dialect) -> "TypeDefaults":
@@ -78,7 +88,14 @@ def type_defaults(dialect) -> "TypeDefaults":
     try:
         return type_defaults_registry[name]
     except KeyError:
-        raise ValueError(f"No type defaults registered for dialect: {name!r}") from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return type_defaults_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No type defaults registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def db_defaults(dialect) -> "DBDefaults":
@@ -86,7 +103,14 @@ def db_defaults(dialect) -> "DBDefaults":
     try:
         return db_defaults_registry[name]
     except KeyError:
-        raise ValueError(f"No DB defaults registered for dialect: {name!r}") from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return db_defaults_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No DB defaults registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 class SQLType(TypeDecorator):

{datachain-0.30.1 → datachain-0.30.2}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.30.1
+Version: 0.30.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -81,8 +81,10 @@ Provides-Extra: video
 Requires-Dist: ffmpeg-python; extra == "video"
 Requires-Dist: imageio[ffmpeg,pyav]>=2.37.0; extra == "video"
 Requires-Dist: opencv-python; extra == "video"
+Provides-Extra: postgres
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "postgres"
 Provides-Extra: tests
-Requires-Dist: datachain[audio,hf,remote,torch,vector,video]; extra == "tests"
+Requires-Dist: datachain[audio,hf,postgres,remote,torch,vector,video]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
@@ -90,6 +92,7 @@ Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
 Requires-Dist: pytest-servers[all]>=0.5.9; extra == "tests"
 Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
 Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
+Requires-Dist: pytest-env>=1.1.0; extra == "tests"
 Requires-Dist: virtualenv; extra == "tests"
 Requires-Dist: dulwich; extra == "tests"
 Requires-Dist: hypothesis; extra == "tests"

{datachain-0.30.1 → datachain-0.30.2}/src/datachain.egg-info/SOURCES.txt

@@ -231,6 +231,8 @@ src/datachain/query/utils.py
 src/datachain/remote/__init__.py
 src/datachain/remote/studio.py
 src/datachain/sql/__init__.py
+src/datachain/sql/postgresql_dialect.py
+src/datachain/sql/postgresql_types.py
 src/datachain/sql/selectable.py
 src/datachain/sql/types.py
 src/datachain/sql/utils.py

{datachain-0.30.1 → datachain-0.30.2}/src/datachain.egg-info/requires.txt

@@ -76,12 +76,15 @@ fsspec>=2024.12.0
 [hf:sys_platform == "linux" or sys_platform == "darwin"]
 datasets[audio]>=4.0.0
 
+[postgres]
+psycopg2-binary>=2.9.0
+
 [remote]
 lz4
 requests>=2.22.0
 
 [tests]
-datachain[audio,hf,remote,torch,vector,video]
+datachain[audio,hf,postgres,remote,torch,vector,video]
 pytest<9,>=8
 pytest-sugar>=0.9.6
 pytest-cov>=4.1.0
@@ -89,6 +92,7 @@ pytest-mock>=3.12.0
 pytest-servers[all]>=0.5.9
 pytest-benchmark[histogram]
 pytest-xdist>=3.3.1
+pytest-env>=1.1.0
 virtualenv
 dulwich
 hypothesis

{datachain-0.30.1 → datachain-0.30.2}/tests/examples/test_examples.py

@@ -93,6 +93,7 @@ def test_multimodal(example):
 
 
 @pytest.mark.examples
+@pytest.mark.incremental_processing
 @pytest.mark.parametrize("example", incremental_processing_examples)
 def test_incremental_processing_examples(example):
     smoke_test(example)