datachain 0.19.2__py3-none-any.whl → 0.20.0__py3-none-any.whl

datachain/__init__.py

@@ -1,3 +1,4 @@
+from datachain.lib import namespaces, projects
 from datachain.lib.data_model import DataModel, DataType, is_chain_type
 from datachain.lib.dc import (
     C,
@@ -67,7 +68,9 @@ __all__ = [
     "is_chain_type",
     "listings",
     "metrics",
+    "namespaces",
     "param",
+    "projects",
     "read_csv",
     "read_database",
     "read_dataset",

datachain/catalog/catalog.py

@@ -41,6 +41,7 @@ from datachain.dataset import (
     DatasetStatus,
     StorageURI,
     create_dataset_uri,
+    parse_dataset_name,
     parse_dataset_uri,
 )
 from datachain.error import (
@@ -48,12 +49,14 @@ from datachain.error import (
     DatasetInvalidVersionError,
     DatasetNotFoundError,
     DatasetVersionNotFoundError,
+    ProjectNotFoundError,
     QueryScriptCancelError,
     QueryScriptRunError,
 )
 from datachain.lib.listing import get_listing
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
+from datachain.project import Project
 from datachain.sql.types import DateTime, SQLType
 from datachain.utils import DataChainDir
 
@@ -155,9 +158,9 @@ class DatasetRowsFetcher(NodesThreadPool):
         self,
         metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
-        remote_ds_name: str,
+        remote_ds: DatasetRecord,
         remote_ds_version: str,
-        local_ds_name: str,
+        local_ds: DatasetRecord,
         local_ds_version: str,
         schema: dict[str, Union[SQLType, type[SQLType]]],
         max_threads: int = PULL_DATASET_MAX_THREADS,
@@ -169,9 +172,9 @@ class DatasetRowsFetcher(NodesThreadPool):
         self._check_dependencies()
         self.metastore = metastore
         self.warehouse = warehouse
-        self.remote_ds_name = remote_ds_name
+        self.remote_ds = remote_ds
         self.remote_ds_version = remote_ds_version
-        self.local_ds_name = local_ds_name
+        self.local_ds = local_ds
         self.local_ds_version = local_ds_version
         self.schema = schema
         self.last_status_check: Optional[float] = None
@@ -207,7 +210,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         Checks are done every PULL_DATASET_CHECK_STATUS_INTERVAL seconds
         """
         export_status_response = self.studio_client.dataset_export_status(
-            self.remote_ds_name, self.remote_ds_version
+            self.remote_ds, self.remote_ds_version
         )
         if not export_status_response.ok:
             raise DataChainError(export_status_response.message)
@@ -254,9 +257,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         import pandas as pd
 
         # metastore and warehouse are not thread safe
-        with self.metastore.clone() as metastore, self.warehouse.clone() as warehouse:
-            local_ds = metastore.get_dataset(self.local_ds_name)
-
+        with self.warehouse.clone() as warehouse:
             urls = list(urls)
 
             for url in urls:
@@ -269,7 +270,7 @@ class DatasetRowsFetcher(NodesThreadPool):
                 df = self.fix_columns(df)
 
                 inserted = warehouse.insert_dataset_rows(
-                    df, local_ds, self.local_ds_version
+                    df, self.local_ds, self.local_ds_version
                 )
                 self.increase_counter(inserted)  # type: ignore [arg-type]
                 # sometimes progress bar doesn't get updated so manually updating it
@@ -675,7 +676,11 @@ class Catalog:
             listing: Optional[Listing]
             if src.startswith("ds://"):
                 ds_name, ds_version = parse_dataset_uri(src)
-                dataset = self.get_dataset(ds_name)
+                ds_namespace, ds_project, ds_name = parse_dataset_name(ds_name)
+                assert ds_namespace
+                assert ds_project
+                project = self.metastore.get_project(ds_project, ds_namespace)
+                dataset = self.get_dataset(ds_name, project)
                 if not ds_version:
                     ds_version = dataset.latest_version
                 dataset_sources = self.warehouse.get_dataset_sources(
@@ -695,7 +700,11 @@ class Catalog:
                         dataset_name=dataset_name,
                     )
                     rows = DatasetQuery(
-                        name=dataset.name, version=ds_version, catalog=self
+                        name=dataset.name,
+                        namespace_name=dataset.project.namespace.name,
+                        project_name=dataset.project.name,
+                        version=ds_version,
+                        catalog=self,
                     ).to_db_records()
                     indexed_sources.append(
                         (
@@ -769,6 +778,7 @@ class Catalog:
     def create_dataset(
         self,
         name: str,
+        project: Optional[Project] = None,
         version: Optional[str] = None,
         *,
         columns: Sequence[Column],
@@ -788,6 +798,7 @@ class Catalog:
         If version is None, then next unused version is created.
         If version is given, then it must be an unused version.
         """
+        DatasetRecord.validate_name(name)
         assert [c.name for c in columns if c.name != "sys__id"], f"got {columns=}"
         if not listing and Client.is_data_source_uri(name):
             raise RuntimeError(
@@ -795,7 +806,7 @@ class Catalog:
             )
         default_version = DEFAULT_DATASET_VERSION
         try:
-            dataset = self.get_dataset(name)
+            dataset = self.get_dataset(name, project)
             default_version = dataset.next_version_patch
             if update_version == "major":
                 default_version = dataset.next_version_major
@@ -820,6 +831,7 @@ class Catalog:
             }
             dataset = self.metastore.create_dataset(
                 name,
+                project.id if project else None,
                 feature_schema=feature_schema,
                 query_script=query_script,
                 schema=schema,
@@ -892,7 +904,7 @@ class Catalog:
         )
 
         if create_rows_table:
-            table_name = self.warehouse.dataset_table_name(dataset.name, version)
+            table_name = self.warehouse.dataset_table_name(dataset, version)
             self.warehouse.create_dataset_rows_table(table_name, columns=columns)
             self.update_dataset_version_with_warehouse_info(dataset, version)
 
@@ -923,7 +935,13 @@ class Catalog:
 
         if not dataset_version.preview:
             values["preview"] = (
-                DatasetQuery(name=dataset.name, version=version, catalog=self)
+                DatasetQuery(
+                    name=dataset.name,
+                    namespace_name=dataset.project.namespace.name,
+                    project_name=dataset.project.name,
+                    version=version,
+                    catalog=self,
+                )
                 .limit(20)
                 .to_db_records()
             )
@@ -949,6 +967,7 @@ class Catalog:
             # updating name must result in updating dataset table names as well
             for version in [v.version for v in dataset.versions]:
                 self.warehouse.rename_dataset_table(
+                    dataset,
                     old_name,
                     new_name,
                     old_version=version,
@@ -986,6 +1005,7 @@ class Catalog:
         self,
         name: str,
         sources: list[str],
+        project: Optional[Project] = None,
         client_config=None,
         recursive=False,
     ) -> DatasetRecord:
@@ -994,6 +1014,8 @@ class Catalog:
 
         from datachain import read_dataset, read_storage
 
+        project = project or self.metastore.default_project
+
         chains = []
         for source in sources:
             if source.startswith(DATASET_PREFIX):
@@ -1006,10 +1028,11 @@ class Catalog:
         # create union of all dataset queries created from sources
         dc = reduce(lambda dc1, dc2: dc1.union(dc2), chains)
         try:
+            dc = dc.settings(project=project.name, namespace=project.namespace.name)
             dc.save(name)
         except Exception as e:  # noqa: BLE001
             try:
-                ds = self.get_dataset(name)
+                ds = self.get_dataset(name, project)
                 self.metastore.update_dataset_status(
                     ds,
                     DatasetStatus.FAILED,
@@ -1026,7 +1049,7 @@ class Catalog:
             except DatasetNotFoundError:
                 raise e from None
 
-        ds = self.get_dataset(name)
+        ds = self.get_dataset(name, project)
 
         self.update_dataset_version_with_warehouse_info(
             ds,
@@ -1034,49 +1057,67 @@ class Catalog:
             sources="\n".join(sources),
         )
 
-        return self.get_dataset(name)
+        return self.get_dataset(name, project)
 
-    def get_dataset(self, name: str) -> DatasetRecord:
-        return self.metastore.get_dataset(name)
+    def get_dataset(
+        self, name: str, project: Optional[Project] = None
+    ) -> DatasetRecord:
+        from datachain.lib.listing import is_listing_dataset
+
+        if is_listing_dataset(name):
+            project = self.metastore.listing_project
+        return self.metastore.get_dataset(name, project.id if project else None)
 
     def get_dataset_with_remote_fallback(
-        self, name: str, version: Optional[str] = None
+        self,
+        name: str,
+        namespace_name: str,
+        project_name: str,
+        version: Optional[str] = None,
     ) -> DatasetRecord:
         try:
-            ds = self.get_dataset(name)
+            project = self.metastore.get_project(project_name, namespace_name)
+            ds = self.get_dataset(name, project)
             if version and not ds.has_version(version):
                 raise DatasetVersionNotFoundError(
                     f"Dataset {name} does not have version {version}"
                 )
             return ds
 
-        except (DatasetNotFoundError, DatasetVersionNotFoundError):
+        except (
+            ProjectNotFoundError,
+            DatasetNotFoundError,
+            DatasetVersionNotFoundError,
+        ):
             print("Dataset not found in local catalog, trying to get from studio")
-
-            remote_ds_uri = f"{DATASET_PREFIX}{name}"
-            if version:
-                remote_ds_uri += f"@v{version}"
+            remote_ds_uri = create_dataset_uri(
+                name, namespace_name, project_name, version
+            )
 
             self.pull_dataset(
                 remote_ds_uri=remote_ds_uri,
                 local_ds_name=name,
                 local_ds_version=version,
             )
-            return self.get_dataset(name)
+            return self.get_dataset(
+                name, self.metastore.get_project(project_name, namespace_name)
+            )
 
     def get_dataset_with_version_uuid(self, uuid: str) -> DatasetRecord:
         """Returns dataset that contains version with specific uuid"""
         for dataset in self.ls_datasets():
             if dataset.has_version_with_uuid(uuid):
-                return self.get_dataset(dataset.name)
+                return self.get_dataset(dataset.name, dataset.project)
         raise DatasetNotFoundError(f"Dataset with version uuid {uuid} not found.")
 
-    def get_remote_dataset(self, name: str) -> DatasetRecord:
+    def get_remote_dataset(
+        self, namespace: str, project: str, name: str
+    ) -> DatasetRecord:
         from datachain.remote.studio import StudioClient
 
         studio_client = StudioClient()
 
-        info_response = studio_client.dataset_info(name)
+        info_response = studio_client.dataset_info(namespace, project, name)
         if not info_response.ok:
             raise DataChainError(info_response.message)
 
@@ -1085,9 +1126,9 @@ class Catalog:
         return DatasetRecord.from_dict(dataset_info)
 
     def get_dataset_dependencies(
-        self, name: str, version: str, indirect=False
+        self, name: str, version: str, project: Optional[Project] = None, indirect=False
     ) -> list[Optional[DatasetDependency]]:
-        dataset = self.get_dataset(name)
+        dataset = self.get_dataset(name, project)
 
         direct_dependencies = self.metastore.get_direct_dataset_dependencies(
             dataset, version
@@ -1101,9 +1142,10 @@ class Catalog:
                 # dependency has been removed
                 continue
             if d.is_dataset:
+                project = self.metastore.get_project(d.project, d.namespace)
                 # only datasets can have dependencies
                 d.dependencies = self.get_dataset_dependencies(
-                    d.name, d.version, indirect=indirect
+                    d.name, d.version, project, indirect=indirect
                 )
 
         return direct_dependencies
@@ -1113,9 +1155,12 @@ class Catalog:
         prefix: Optional[str] = None,
         include_listing: bool = False,
         studio: bool = False,
+        project: Optional[Project] = None,
     ) -> Iterator[DatasetListRecord]:
         from datachain.remote.studio import StudioClient
 
+        project_id = project.id if project else None
+
         if studio:
             client = StudioClient()
             response = client.ls_datasets(prefix=prefix)
@@ -1130,9 +1175,11 @@ class Catalog:
                 if not d.get("name", "").startswith(QUERY_DATASET_PREFIX)
             )
         elif prefix:
-            datasets = self.metastore.list_datasets_by_prefix(prefix)
+            datasets = self.metastore.list_datasets_by_prefix(
+                prefix, project_id=project_id
+            )
         else:
-            datasets = self.metastore.list_datasets()
+            datasets = self.metastore.list_datasets(project_id=project_id)
 
         for d in datasets:
             if not d.is_bucket_listing or include_listing:
@@ -1144,11 +1191,15 @@ class Catalog:
         include_listing: bool = False,
         with_job: bool = True,
         studio: bool = False,
+        project: Optional[Project] = None,
     ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", Optional["Job"]]]:
         """Iterate over all dataset versions with related jobs."""
         datasets = list(
             self.ls_datasets(
-                prefix=prefix, include_listing=include_listing, studio=studio
+                prefix=prefix,
+                include_listing=include_listing,
+                studio=studio,
+                project=project,
             )
         )
 
@@ -1184,6 +1235,7 @@ class Catalog:
             prefix=prefix,
             include_listing=True,
             with_job=False,
+            project=self.metastore.listing_project,
         )
 
         return [
@@ -1193,13 +1245,21 @@ class Catalog:
         ]
 
     def ls_dataset_rows(
-        self, name: str, version: str, offset=None, limit=None
+        self,
+        dataset: DatasetRecord,
+        version: str,
+        offset=None,
+        limit=None,
     ) -> list[dict]:
         from datachain.query.dataset import DatasetQuery
 
-        dataset = self.get_dataset(name)
-
-        q = DatasetQuery(name=dataset.name, version=version, catalog=self)
+        q = DatasetQuery(
+            name=dataset.name,
+            namespace_name=dataset.project.namespace.name,
+            project_name=dataset.project.name,
+            version=version,
+            catalog=self,
+        )
         if limit:
             q = q.limit(limit)
         if offset:
@@ -1232,35 +1292,29 @@ class Catalog:
         bucket_uri: str,
         name: str,
         version: str,
+        project: Optional[Project] = None,
         client_config=None,
     ) -> list[str]:
-        dataset = self.get_dataset(name)
+        dataset = self.get_dataset(name, project)
 
         return self.warehouse.export_dataset_table(
             bucket_uri, dataset, version, client_config
         )
 
-    def dataset_table_export_file_names(self, name: str, version: str) -> list[str]:
-        dataset = self.get_dataset(name)
+    def dataset_table_export_file_names(
+        self, name: str, version: str, project: Optional[Project] = None
+    ) -> list[str]:
+        dataset = self.get_dataset(name, project)
         return self.warehouse.dataset_table_export_file_names(dataset, version)
 
     def remove_dataset(
         self,
         name: str,
+        project: Optional[Project] = None,
         version: Optional[str] = None,
         force: Optional[bool] = False,
-        studio: Optional[bool] = False,
     ):
-        from datachain.remote.studio import StudioClient
-
-        if studio:
-            client = StudioClient()
-            response = client.rm_dataset(name, version=version, force=force)
-            if not response.ok:
-                raise DataChainError(response.message)
-            return
-
-        dataset = self.get_dataset(name)
+        dataset = self.get_dataset(name, project)
         if not version and not force:
             raise ValueError(f"Missing dataset version from input for dataset {name}")
         if version and not dataset.has_version(version):
@@ -1282,19 +1336,21 @@ class Catalog:
     def edit_dataset(
         self,
         name: str,
+        project: Optional[Project] = None,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
     ) -> DatasetRecord:
         update_data = {}
         if new_name:
+            DatasetRecord.validate_name(new_name)
             update_data["name"] = new_name
         if description is not None:
             update_data["description"] = description
         if attrs is not None:
             update_data["attrs"] = attrs  # type: ignore[assignment]
 
-        dataset = self.get_dataset(name)
+        dataset = self.get_dataset(name, project)
         return self.update_dataset(dataset, **update_data)
 
     def ls(
@@ -1351,7 +1407,29 @@ class Catalog:
         except Exception as e:
             raise DataChainError("Error when parsing dataset uri") from e
 
-        remote_ds = self.get_remote_dataset(remote_ds_name)
+        remote_namespace, remote_project, remote_ds_name = parse_dataset_name(
+            remote_ds_name
+        )
+        if not remote_namespace or not remote_project:
+            raise DataChainError(
+                f"Invalid fully qualified dataset name {remote_ds_name}, namespace"
+                f" or project missing"
+            )
+
+        if local_ds_name:
+            local_namespace, local_project, local_ds_name = parse_dataset_name(
+                local_ds_name
+            )
+            if local_namespace and local_namespace != remote_namespace:
+                raise DataChainError(
+                    "Local namespace must be the same to remote namespace"
+                )
+            if local_project and local_project != remote_project:
+                raise DataChainError("Local project must be the same to remote project")
+
+        remote_ds = self.get_remote_dataset(
+            remote_namespace, remote_project, remote_ds_name
+        )
 
         try:
             # if version is not specified in uri, take the latest one
@@ -1359,7 +1437,12 @@ class Catalog:
                 version = remote_ds.latest_version
                 print(f"Version not specified, pulling the latest one (v{version})")
                 # updating dataset uri with latest version
-                remote_ds_uri = create_dataset_uri(remote_ds_name, version)
+                remote_ds_uri = create_dataset_uri(
+                    remote_ds.name,
+                    remote_ds.project.namespace.name,
+                    remote_ds.project.name,
+                    version,
+                )
             remote_ds_version = remote_ds.get_version(version)
         except (DatasetVersionNotFoundError, StopIteration) as exc:
             raise DataChainError(
@@ -1368,7 +1451,13 @@ class Catalog:
 
         local_ds_name = local_ds_name or remote_ds.name
         local_ds_version = local_ds_version or remote_ds_version.version
-        local_ds_uri = create_dataset_uri(local_ds_name, local_ds_version)
+
+        local_ds_uri = create_dataset_uri(
+            local_ds_name,
+            remote_ds.project.namespace.name,
+            remote_ds.project.name,
+            local_ds_version,
+        )
 
         try:
             # try to find existing dataset with the same uuid to avoid pulling again
@@ -1377,7 +1466,10 @@ class Catalog:
                 remote_ds_version.uuid
             )
             existing_ds_uri = create_dataset_uri(
-                existing_ds.name, existing_ds_version.version
+                existing_ds.name,
+                existing_ds.project.namespace.name,
+                existing_ds.project.name,
+                existing_ds_version.version,
             )
             if existing_ds_uri == remote_ds_uri:
                 print(f"Local copy of dataset {remote_ds_uri} already present")
@@ -1391,8 +1483,26 @@ class Catalog:
         except DatasetNotFoundError:
             pass
 
+        # Create namespace and project if doesn't exist
+        print(
+            f"Creating namespace {remote_ds.project.namespace.name} and project"
+            f" {remote_ds.project.name}"
+        )
+
+        namespace = self.metastore.create_namespace(
+            remote_ds.project.namespace.name,
+            description=remote_ds.project.namespace.description,
+            uuid=remote_ds.project.namespace.uuid,
+        )
+        project = self.metastore.create_project(
+            remote_ds.project.name,
+            namespace.name,
+            description=remote_ds.project.description,
+            uuid=remote_ds.project.uuid,
+        )
+
         try:
-            local_dataset = self.get_dataset(local_ds_name)
+            local_dataset = self.get_dataset(local_ds_name, project=project)
             if local_dataset and local_dataset.has_version(local_ds_version):
                 raise DataChainError(
                     f"Local dataset {local_ds_uri} already exists with different uuid,"
@@ -1414,6 +1524,7 @@ class Catalog:
 
         local_ds = self.create_dataset(
             local_ds_name,
+            project,
             local_ds_version,
             query_script=remote_ds_version.query_script,
             create_rows=True,
@@ -1426,7 +1537,7 @@ class Catalog:
         # asking remote to export dataset rows table to s3 and to return signed
         # urls of exported parts, which are in parquet format
         export_response = studio_client.export_dataset_table(
-            remote_ds_name, remote_ds_version.version
+            remote_ds, remote_ds_version.version
         )
         if not export_response.ok:
             raise DataChainError(export_response.message)
@@ -1457,9 +1568,9 @@ class Catalog:
                 rows_fetcher = DatasetRowsFetcher(
                     metastore,
                     warehouse,
-                    remote_ds_name,
+                    remote_ds,
                     remote_ds_version.version,
-                    local_ds_name,
+                    local_ds,
                     local_ds_version,
                     schema,
                     progress_bar=dataset_save_progress_bar,
@@ -1469,7 +1580,7 @@ class Catalog:
                         iter(batch(signed_urls)), dataset_save_progress_bar
                     )
                 except:
-                    self.remove_dataset(local_ds_name, local_ds_version)
+                    self.remove_dataset(local_ds_name, project, local_ds_version)
                     raise
 
         local_ds = self.metastore.update_dataset_status(
@@ -1526,7 +1637,11 @@ class Catalog:
             )
 
         self.create_dataset_from_sources(
-            output, sources, client_config=client_config, recursive=recursive
+            output,
+            sources,
+            self.metastore.default_project,
+            client_config=client_config,
+            recursive=recursive,
         )
 
     def query(

datachain/cli/__init__.py

@@ -152,9 +152,6 @@ def handle_dataset_command(args, catalog):
             new_name=args.new_name,
             description=args.description,
             attrs=args.attrs,
-            studio=args.studio,
-            local=args.local,
-            all=args.all,
             team=args.team,
         ),
         "ls": lambda: list_datasets(
@@ -172,8 +169,6 @@ def handle_dataset_command(args, catalog):
             version=args.version,
             force=args.force,
             studio=args.studio,
-            local=args.local,
-            all=args.all,
             team=args.team,
         ),
         "remove": lambda: rm_dataset(
@@ -182,8 +177,6 @@ def handle_dataset_command(args, catalog):
             version=args.version,
             force=args.force,
             studio=args.studio,
-            local=args.local,
-            all=args.all,
             team=args.team,
         ),
     }