datachain 0.7.5__py3-none-any.whl → 0.7.7__py3-none-any.whl

datachain/lib/pytorch.py

@@ -70,20 +70,17 @@ class PytorchDataset(IterableDataset):
         # For compatibility with multiprocessing,
         # we can only store params in __init__(), as Catalog isn't picklable
         # see https://github.com/iterative/dvcx/issues/954
-        self._idgen_params = catalog.id_generator.clone_params()
         self._ms_params = catalog.metastore.clone_params()
         self._wh_params = catalog.warehouse.clone_params()
         self._catalog_params = catalog.get_init_params()
         self.catalog: Optional[Catalog] = None
 
     def _get_catalog(self) -> "Catalog":
-        idgen_cls, idgen_args, idgen_kwargs = self._idgen_params
-        idgen = idgen_cls(*idgen_args, **idgen_kwargs)
         ms_cls, ms_args, ms_kwargs = self._ms_params
         ms = ms_cls(*ms_args, **ms_kwargs)
         wh_cls, wh_args, wh_kwargs = self._wh_params
         wh = wh_cls(*wh_args, **wh_kwargs)
-        return Catalog(idgen, ms, wh, **self._catalog_params)
+        return Catalog(ms, wh, **self._catalog_params)
 
     def __iter__(self) -> Iterator[Any]:
         if self.catalog is None:

datachain/query/dataset.py

@@ -442,9 +442,6 @@ class UDFStep(Step, ABC):
                 udf_info = {
                     "udf_data": filtered_cloudpickle_dumps(self.udf),
                     "catalog_init": self.catalog.get_init_params(),
-                    "id_generator_clone_params": (
-                        self.catalog.id_generator.clone_params()
-                    ),
                     "metastore_clone_params": self.catalog.metastore.clone_params(),
                     "warehouse_clone_params": self.catalog.warehouse.clone_params(),
                     "table": udf_table,
@@ -969,8 +966,6 @@ class SQLGroupBy(SQLClause):
     def apply_sql_clause(self, query) -> Select:
         if not self.cols:
             raise ValueError("No columns to select")
-        if not self.group_by:
-            raise ValueError("No columns to group by")
 
         subquery = query.subquery()
 

datachain/query/dispatch.py

@@ -65,7 +65,6 @@ def udf_entrypoint() -> int:
     dispatch = UDFDispatcher(
         udf_info["udf_data"],
         udf_info["catalog_init"],
-        udf_info["id_generator_clone_params"],
         udf_info["metastore_clone_params"],
         udf_info["warehouse_clone_params"],
         udf_fields=udf_info["udf_fields"],
@@ -119,7 +118,6 @@ class UDFDispatcher:
         self,
         udf_data,
         catalog_init_params,
-        id_generator_clone_params,
         metastore_clone_params,
         warehouse_clone_params,
         udf_fields: "Sequence[str]",
@@ -129,11 +127,6 @@ class UDFDispatcher:
     ):
         self.udf_data = udf_data
         self.catalog_init_params = catalog_init_params
-        (
-            self.id_generator_class,
-            self.id_generator_args,
-            self.id_generator_kwargs,
-        ) = id_generator_clone_params
         (
             self.metastore_class,
             self.metastore_args,
@@ -155,18 +148,13 @@ class UDFDispatcher:
 
     def _create_worker(self) -> "UDFWorker":
         if not self.catalog:
-            id_generator = self.id_generator_class(
-                *self.id_generator_args, **self.id_generator_kwargs
-            )
             metastore = self.metastore_class(
                 *self.metastore_args, **self.metastore_kwargs
             )
             warehouse = self.warehouse_class(
                 *self.warehouse_args, **self.warehouse_kwargs
             )
-            self.catalog = Catalog(
-                id_generator, metastore, warehouse, **self.catalog_init_params
-            )
+            self.catalog = Catalog(metastore, warehouse, **self.catalog_init_params)
         self.udf = loads(self.udf_data)
         return UDFWorker(
             self.catalog,

datachain/query/session.py

@@ -85,7 +85,6 @@ class Session:
         if self.is_new_catalog:
             self.catalog.metastore.close_on_exit()
             self.catalog.warehouse.close_on_exit()
-            self.catalog.id_generator.close_on_exit()
 
         if Session.SESSION_CONTEXTS:
             Session.SESSION_CONTEXTS.pop()

datachain/remote/studio.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import logging
 import os
@@ -23,7 +24,8 @@ DatasetStatsData = Optional[DatasetStats]
 DatasetRowsData = Optional[Iterable[dict[str, Any]]]
 DatasetExportStatus = Optional[dict[str, Any]]
 DatasetExportSignedUrls = Optional[list[str]]
-
+FileUploadData = Optional[dict[str, Any]]
+JobData = Optional[dict[str, Any]]
 
 logger = logging.getLogger("datachain")
 
@@ -308,3 +310,33 @@ class StudioClient:
             "datachain/dataset-export-status",
             {"dataset_name": name, "dataset_version": version},
         )
+
+    def upload_file(self, file_name: str, content: bytes) -> Response[FileUploadData]:
+        data = {
+            "file_content": base64.b64encode(content).decode("utf-8"),
+            "file_name": file_name,
+        }
+        return self._send_request("datachain/upload-file", data)
+
+    def create_job(
+        self,
+        query: str,
+        query_type: str,
+        environment: Optional[str] = None,
+        workers: Optional[int] = None,
+        query_name: Optional[str] = None,
+        files: Optional[list[str]] = None,
+        python_version: Optional[str] = None,
+        requirements: Optional[str] = None,
+    ) -> Response[JobData]:
+        data = {
+            "query": query,
+            "query_type": query_type,
+            "environment": environment,
+            "workers": workers,
+            "query_name": query_name,
+            "files": files,
+            "python_version": python_version,
+            "requirements": requirements,
+        }
+        return self._send_request("datachain/job", data)

datachain/studio.py

@@ -34,6 +34,19 @@ def process_studio_cli_args(args: "Namespace"):
         print(tabulate(rows, headers="keys"))
         return 0
 
+    if args.cmd == "run":
+        return create_job(
+            args.query_file,
+            args.team,
+            args.env_file,
+            args.env,
+            args.workers,
+            args.files,
+            args.python_version,
+            args.req,
+            args.req_file,
+        )
+
     if args.cmd == "team":
         return set_team(args)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
@@ -168,3 +181,70 @@ def save_config(hostname, token):
         conf["studio"] = studio_conf
 
     return config.config_file()
+
+
+def create_job(
+    query_file: str,
+    team_name: Optional[str],
+    env_file: Optional[str] = None,
+    env: Optional[list[str]] = None,
+    workers: Optional[int] = None,
+    files: Optional[list[str]] = None,
+    python_version: Optional[str] = None,
+    req: Optional[list[str]] = None,
+    req_file: Optional[str] = None,
+):
+    query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
+    with open(query_file) as f:
+        query = f.read()
+
+    environment = "\n".join(env) if env else ""
+    if env_file:
+        with open(env_file) as f:
+            environment = f.read() + "\n" + environment
+
+    requirements = "\n".join(req) if req else ""
+    if req_file:
+        with open(req_file) as f:
+            requirements = f.read() + "\n" + requirements
+
+    client = StudioClient(team=team_name)
+    file_ids = upload_files(client, files) if files else []
+
+    response = client.create_job(
+        query=query,
+        query_type=query_type,
+        environment=environment,
+        workers=workers,
+        query_name=os.path.basename(query_file),
+        files=file_ids,
+        python_version=python_version,
+        requirements=requirements,
+    )
+    if not response.ok:
+        raise_remote_error(response.message)
+
+    if not response.data:
+        raise DataChainError("Failed to create job")
+
+    print(f"Job {response.data.get('job', {}).get('id')} created")
+    print("Open the job in Studio at", response.data.get("job", {}).get("url"))
+
+
+def upload_files(client: StudioClient, files: list[str]) -> list[str]:
+    file_ids = []
+    for file in files:
+        file_name = os.path.basename(file)
+        with open(file, "rb") as f:
+            file_content = f.read()
+        response = client.upload_file(file_name, file_content)
+        if not response.ok:
+            raise_remote_error(response.message)
+
+        if not response.data:
+            raise DataChainError(f"Failed to upload file {file_name}")
+
+        file_id = response.data.get("blob", {}).get("id")
+        if file_id:
+            file_ids.append(str(file_id))
+    return file_ids

{datachain-0.7.5.dist-info → datachain-0.7.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.7.5
+Version: 0.7.7
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.7.5.dist-info → datachain-0.7.7.dist-info}/RECORD

@@ -2,10 +2,10 @@ datachain/__init__.py,sha256=ofPJ6B-d-ybSDRrE7J6wqF_ZRAB2W9U8l-eeuBtqPLg,865
 datachain/__main__.py,sha256=hG3Y4ARGEqe1AWwNMd259rBlqtphx1Wk39YbueQ0yV8,91
 datachain/asyn.py,sha256=5aKrjnUxk0mtnZeFKNJd1DCE0MsnSoyJBZkr0y9H_a0,9313
 datachain/cache.py,sha256=s0YHN7qurmQv-eC265TjeureK84TebWWAnL07cxchZQ,2997
-datachain/cli.py,sha256=1hiBClE1kbRyx0DK3uX5KMVa0ktbsG6TsFSNvoT2xxs,39399
+datachain/cli.py,sha256=Ysm-6Kb-54FfkN35VJIe5vW7Kik8VGA3wcyCUnqPBHg,42245
 datachain/cli_utils.py,sha256=jrn9ejGXjybeO1ur3fjdSiAyCHZrX0qsLLbJzN9ErPM,2418
 datachain/config.py,sha256=g8qbNV0vW2VEKpX-dGZ9pAn0DAz6G2ZFcr7SAV3PoSM,4272
-datachain/dataset.py,sha256=4KciGp7XesQ8MsAzs8G7854ZoVcMdPE0ultVh-XSkIw,18597
+datachain/dataset.py,sha256=P-pDBgvPqJGDhq_I7fwCfb6hY8E8mIAO8Q0NT7SNlNE,19128
 datachain/error.py,sha256=bxAAL32lSeMgzsQDEHbGTGORj-mPzzpCRvWDPueJNN4,1092
 datachain/job.py,sha256=Jt4sNutMHJReaGsj3r3scueN5aESLGfhimAa8pUP7Is,1271
 datachain/listing.py,sha256=TgKg25ZWAP5enzKgw2_2GUPJVdnQUh6uySHB5SJrUY4,7773
@@ -14,13 +14,13 @@ datachain/nodes_fetcher.py,sha256=ILMzUW5o4_6lUOVrLDC9gJPCXfcgKnMG68plrc7dAOA,11
 datachain/nodes_thread_pool.py,sha256=uPo-xl8zG5m9YgODjPFBpbcqqHjI-dcxH87yAbj_qco,3192
 datachain/progress.py,sha256=5KotcvvzAUL_RF0GEj4JY0IB1lyImnmHxe89YkT1XO4,4330
 datachain/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/studio.py,sha256=w41vgVPrBfJ02XQOaDccLbh-1uSAfq9cAgOmkYUqExE,4845
+datachain/studio.py,sha256=MthVADn-jM2I5TlESOfbzFnKGZjpuk9bM8m2vqOK-C8,7227
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=-mSFowjIidJ4_sMXInvNHLn4rK_QnHuIlLuH1_lMGmI,13897
 datachain/catalog/__init__.py,sha256=g2iAAFx_gEIrqshXlhSEbrc8qDaEH11cjU40n3CHDz4,409
-datachain/catalog/catalog.py,sha256=l_HAxor5i_F03VvbmMuwhi4INhsmNrqubyydPhXWo2Y,57980
+datachain/catalog/catalog.py,sha256=s4fat0jjP3JPq0RGQ9zfzRkX1JavxxCrcB1tJKMgsks,57686
 datachain/catalog/datasource.py,sha256=D-VWIVDCM10A8sQavLhRXdYSCG7F4o4ifswEF80_NAQ,1412
-datachain/catalog/loader.py,sha256=-6VelNfXUdgUnwInVyA8g86Boxv2xqhTh9xNS-Zlwig,8242
+datachain/catalog/loader.py,sha256=HA_mBC7q_My8j2WnSvIjUGuJpl6SIdg5vvy_lagxJlA,5733
 datachain/client/__init__.py,sha256=T4wiYL9KIM0ZZ_UqIyzV8_ufzYlewmizlV4iymHNluE,86
 datachain/client/azure.py,sha256=ffxs26zm6KLAL1aUWJm-vtzuZP3LSNha7UDGXynMBKo,2234
 datachain/client/fileslice.py,sha256=bT7TYco1Qe3bqoc8aUkUZcPdPofJDHlryL5BsTn9xsY,3021
@@ -29,21 +29,20 @@ datachain/client/gcs.py,sha256=cnTIr5GS6dbYOEYfqehhyQu3dr6XNjPHSg5U3FkivUk,4124
 datachain/client/hf.py,sha256=XeVJVbiNViZCpn3sfb90Fr8SYO3BdLmfE3hOWMoqInE,951
 datachain/client/local.py,sha256=vwbgCwZ7IqY2voj2l7tLJjgov7Dp--fEUvUwUBsMbls,4457
 datachain/client/s3.py,sha256=CVHBUZ1Ic2Q3370nl-Bbe69phuWjFlrVv9dTJKBpRT0,6019
-datachain/data_storage/__init__.py,sha256=cEOJpyu1JDZtfUupYucCDNFI6e5Wmp_Oyzq6rZv32Y8,398
+datachain/data_storage/__init__.py,sha256=9Wit-oe5P46V7CJQTD0BJ5MhOa2Y9h3ddJ4VWTe-Lec,273
 datachain/data_storage/db_engine.py,sha256=81Ol1of9TTTzD97ORajCnP366Xz2mEJt6C-kTUCaru4,3406
-datachain/data_storage/id_generator.py,sha256=lCEoU0BM37Ai2aRpSbwo5oQT0GqZnSpYwwvizathRMQ,4292
 datachain/data_storage/job.py,sha256=w-7spowjkOa1P5fUVtJou3OltT0L48P0RYWZ9rSJ9-s,383
-datachain/data_storage/metastore.py,sha256=VPq-Dl8P-RbZQMzn6vB9aXBPKUWPTwP8ypkaVfE-7PU,37661
+datachain/data_storage/metastore.py,sha256=EzSsfR_l_84i1AewYygpdsJyzGqEmvXjpeohlYF7h4A,37435
 datachain/data_storage/schema.py,sha256=-QVlRvD0dfu-ZFUxylEoSnLJLnleMEjVlcAb2OGu-AY,9895
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
-datachain/data_storage/sqlite.py,sha256=nF-2B-n8YZh9cJlZv4XnbahAJDW6pvrp1h9L-140M7A,27538
-datachain/data_storage/warehouse.py,sha256=kFLhYEFkpsfl65Lr1c4t4HJt3nO1Ez_QQ76aQNN30fc,30966
-datachain/func/__init__.py,sha256=4VUt5BaLdBAl_BnAku0Jb8plqd7kDOiYrQTMG3pN0c4,794
+datachain/data_storage/sqlite.py,sha256=D_ZQ0PHmZzHO2dinv4naVJocUDIZUwV4WAz692C1cyk,22521
+datachain/data_storage/warehouse.py,sha256=tjIkU-5JywBR0apCyqTcwSyaRtGxhu2L7IVjrz-55uc,30802
+datachain/func/__init__.py,sha256=VAN7N2-eCHgidMCFI-fJTkCwdI1U_NIuCOgYc4sfYUQ,812
 datachain/func/aggregate.py,sha256=7_IPrIwb2XSs3zG4iOr1eTvzn6kNVe2mkzvNzjusDHk,10942
 datachain/func/array.py,sha256=zHDNWuWLA7HVa9FEvQeHhVi00_xqenyleTqcLwkXWBI,5477
 datachain/func/base.py,sha256=wA0sBQAVyN9LPxoo7Ox83peS0zUVnyuKxukwAcjGLfY,534
 datachain/func/conditional.py,sha256=mQroxsoExpBW84Zm5dAYP4OpBblWmzfnF2qJq9rba54,2223
-datachain/func/func.py,sha256=9wqdxxisoDL0w8qKGQmL6sNdgJeIOzotEUPlxu9t2IQ,12326
+datachain/func/func.py,sha256=GykhTvNbACFSwaSXsgVlDnqR48kpP_GNAxm3bcq1RYg,12560
 datachain/func/path.py,sha256=mqN_mfkwv44z2II7DMTp_fGGw95hmTCNls_TOFNpr4k,3155
 datachain/func/random.py,sha256=pENOLj9rSmWfGCnOsUIaCsVC5486zQb66qfQvXaz9Z4,452
 datachain/func/string.py,sha256=NQzaXXYu7yb72HPADy4WrFlcgvTS77L9x7-qvCKJtnk,4522
@@ -53,7 +52,7 @@ datachain/lib/arrow.py,sha256=b5efxAUaNNYVwtXVJqj07D3zf5KC-BPlLCxKEZbEG6w,9429
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=zS4lmXHVBXc9ntcyea2a1CRLXGSAN_0glXcF88CohgY,2685
 datachain/lib/dataset_info.py,sha256=IjdF1E0TQNOq9YyynfWiCFTeZpbyGfyJvxgJY4YN810,2493
-datachain/lib/dc.py,sha256=J7liATKQBJCkeHanVLr0s3d1t5wxiiiSJuSbuxKBbLg,89527
+datachain/lib/dc.py,sha256=t5y5tsYyU7uuk3gEPPhhcDSZ1tL1aHkKG2W54eHiUq8,89492
 datachain/lib/file.py,sha256=-XMkL6ED1sE7TMhWoMRTEuOXswZJw8X6AEmJDONFP74,15019
 datachain/lib/hf.py,sha256=a-zFpDmZIR4r8dlNNTjfpAKSnuJ9xyRXlgcdENiXt3E,5864
 datachain/lib/image.py,sha256=AMXYwQsmarZjRbPCZY3M1jDsM2WAB_b3cTY4uOIuXNU,2675
@@ -61,7 +60,7 @@ datachain/lib/listing.py,sha256=cVkCp7TRVpcZKSx-Bbk9t51bQI9Mw0o86W6ZPhAsuzM,3667
 datachain/lib/listing_info.py,sha256=9ua40Hw0aiQByUw3oAEeNzMavJYfW0Uhe8YdCTK-m_g,1110
 datachain/lib/meta_formats.py,sha256=anK2bDVbaeCCh0yvKUBaW2MVos3zRgdaSV8uSduzPcU,6680
 datachain/lib/model_store.py,sha256=DNIv8Y6Jtk1_idNLzIpsThOsdW2BMAudyUCbPUcgcxk,2515
-datachain/lib/pytorch.py,sha256=W-ARi2xH1f1DUkVfRuerW-YWYgSaJASmNCxtz2lrJGI,6072
+datachain/lib/pytorch.py,sha256=Nh6fUbQMLX8OpZvX4tw4bJjTCQpRKi0jSLgkJnLHdTM,5880
 datachain/lib/settings.py,sha256=ZELRCTLbi5vzRPiDX6cQ9LLg9TefJ_A05gIGni0lll8,2535
 datachain/lib/signal_schema.py,sha256=_uh19nCKhiD9ua8oIN1Q8R9iYv1BZAuqTJCLYVmyW8k,24557
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
@@ -88,15 +87,15 @@ datachain/model/ultralytics/pose.py,sha256=71KBTcoST2wcEtsyGXqLVpvUtqbp9gwZGA15p
 datachain/model/ultralytics/segment.py,sha256=Z1ab0tZRJubSYNH4KkFlzhYeGNTfAyC71KmkQcToHDQ,2760
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
 datachain/query/batch.py,sha256=5fEhORFe7li12SdYddaSK3LyqksMfCHhwN1_A6TfsA4,3485
-datachain/query/dataset.py,sha256=bQVG4WnJfBQpvnxouIdDlsJF2gB8V4lDp4Zu9JeZ-rc,54771
-datachain/query/dispatch.py,sha256=wjjTWw6sFQbB9SKRh78VbfvwSMgJXCfqJklS3-9KnCU,12025
+datachain/query/dataset.py,sha256=J6SbCLnFlZgCxRchc3tVk5tcC7xo1Hp616JGlEZXCDo,54547
+datachain/query/dispatch.py,sha256=fZ0TgGFRcsrYh1iXQoZVjkUl4Xetom9PSHoeDes3IRs,11606
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
 datachain/query/queue.py,sha256=waqM_KzavU8C-G95-4211Nd4GXna_u2747Chgwtgz2w,3839
 datachain/query/schema.py,sha256=b_KnVy6B26Ol4nYG0LqNNpeQ1QYPk95YRGUjXfdaQWs,6606
-datachain/query/session.py,sha256=50SOdLNCjqHHKI-L4xGXyzTVxzMWfANqKqjeYre-c2k,5959
+datachain/query/session.py,sha256=vvLIJ5b8eElovHLAWq_CZJXmN5t7C7iAZA7x9wPPOms,5905
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/remote/studio.py,sha256=z9DTDqfdWKT8MC23wRDTOHvI8hc_OySS1Ce3F617gjA,9906
+datachain/remote/studio.py,sha256=jp6NWo7OPUxqO8uYEHP0_XFlmj47rMxC80qKQ7rA3Xk,11024
 datachain/sql/__init__.py,sha256=6SQRdbljO3d2hx3EAVXEZrHQKv5jth0Jh98PogT59No,262
 datachain/sql/selectable.py,sha256=cTc60qVoAwqqss0Vop8Lt5Z-ROnM1XrQmL_GLjRxhXs,1765
 datachain/sql/types.py,sha256=ASSPkmM5EzdRindqj2O7WHLXq8VHAgFYedG8lYfGvVI,14045
@@ -117,9 +116,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ZgDcrNiKiPXZmKD591_1z9qRIXitu5zwAsoVPB7ykiU,2508
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.7.5.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.7.5.dist-info/METADATA,sha256=GHZwnQPiDUBHJMdyMf-t7LJkMPyAkygNewihcXwETSs,18006
-datachain-0.7.5.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-datachain-0.7.5.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.7.5.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.7.5.dist-info/RECORD,,
+datachain-0.7.7.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.7.7.dist-info/METADATA,sha256=laxYaz9f-PIJ30f3krSjRu45CjyfbnBM8Q4kddXa9dM,18006
+datachain-0.7.7.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+datachain-0.7.7.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.7.7.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.7.7.dist-info/RECORD,,

datachain/data_storage/id_generator.py

@@ -1,136 +0,0 @@
-import logging
-from abc import ABC, abstractmethod
-from collections.abc import Iterable
-from functools import cached_property
-from typing import TYPE_CHECKING, Optional
-
-from sqlalchemy import Column, Integer, Table, Text
-
-from datachain.data_storage.serializer import Serializable
-
-if TYPE_CHECKING:
-    from sqlalchemy.schema import SchemaItem
-
-    from datachain.data_storage.db_engine import DatabaseEngine
-
-
-logger = logging.getLogger("datachain")
-
-
-class AbstractIDGenerator(ABC, Serializable):
-    """
-    Abstract ID Generator class. This class is responsible for generating
-    unique IDs for each prefix (e.g. S3 bucket or dataset).
-    """
-
-    @abstractmethod
-    def clone(self) -> "AbstractIDGenerator":
-        """Clones AbstractIDGenerator implementation."""
-
-    def init(self) -> None:
-        """Initialize ID generator."""
-
-    def cleanup_for_tests(self):
-        """Cleanup for tests."""
-
-    def close(self) -> None:
-        """Closes any active database connections."""
-
-    def close_on_exit(self) -> None:
-        """Closes any active database or HTTP connections, called on Session exit or
-        for test cleanup only, as some ID Generator implementations may handle this
-        differently.
-        """
-        self.close()
-
-    @abstractmethod
-    def init_id(self, uri: str) -> None:
-        """Initializes the ID generator for the given URI with zero last_id."""
-
-    @abstractmethod
-    def get_next_ids(self, uri: str, count: int) -> range:
-        """Returns a range of IDs for the given URI."""
-
-    def get_next_id(self, uri: str) -> int:
-        """Returns the next ID for the given URI."""
-        return self.get_next_ids(uri, 1)[0]
-
-    def delete_uri(self, uri: str):
-        """Deletes the given URI."""
-        self.delete_uris([uri])
-
-    def delete_uris(self, uris: Iterable[str]):
-        """Deletes the given URIs."""
-
-
-class AbstractDBIDGenerator(AbstractIDGenerator):
-    """
-    Abstract ID Generator class, to be implemented by any Database Adapters
-    for a specific database system. This class is responsible for generating
-    unique IDs for each prefix (e.g. S3 bucket or dataset) and storing them
-    in a database. It is also responsible for initializing the database
-    and creating the necessary tables.
-    """
-
-    _db: "DatabaseEngine"
-    _table_prefix: Optional[str] = None
-    _skip_db_init: bool = False
-    _base_table_name = "id_generator"
-
-    def __init__(
-        self,
-        db: "DatabaseEngine",
-        table_prefix: Optional[str] = None,
-        skip_db_init: bool = False,
-    ):
-        self._db = db
-        self._table_prefix = table_prefix
-        self._skip_db_init = skip_db_init
-        if db and not skip_db_init:
-            self.init()
-
-    @abstractmethod
-    def clone(self) -> "AbstractDBIDGenerator":
-        """Clones AbstractIDGenerator implementation."""
-
-    def close(self) -> None:
-        """Closes any active database connections."""
-        self.db.close()
-
-    @property
-    def db(self) -> "DatabaseEngine":
-        return self._db
-
-    @property
-    def _columns(self) -> list["SchemaItem"]:
-        return [
-            Column("uri", Text, primary_key=True, nullable=False),
-            # This is the last id used (and starts at zero if no ids have been used)
-            Column("last_id", Integer, nullable=False),
-        ]
-
-    @cached_property
-    def _table(self) -> Table:
-        table_name = self._base_table_name
-        if self._table_prefix:
-            table_name = f"{self._table_prefix}_{table_name}"
-        return Table(table_name, self.db.metadata, *self._columns, extend_existing=True)
-
-    def init(self) -> None:
-        self.db.create_table(self._table, if_not_exists=True)
-
-    def cleanup_for_tests(self):
-        """Cleanup for tests."""
-        self.db.drop_table(self._table, if_exists=True)
-
-    @abstractmethod
-    def init_id(self, uri: str) -> None:
-        """Initializes the ID generator for the given URI with zero last_id."""
-
-    @abstractmethod
-    def get_next_ids(self, uri: str, count: int) -> range:
-        """Returns a range of IDs for the given URI."""
-
-    def delete_uris(self, uris: Iterable[str]):
-        """Deletes the given URIs from the database."""
-        self.db.execute(self._table.delete().where(self._table.c.uri.in_(uris)))