dsi-workflow 1.2.4__tar.gz → 1.2.6__tar.gz

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsi-workflow
-Version: 1.2.4
+Version: 1.2.6
 Summary: A Data Science Infrastructure Project
 Author-email: Jesus Pulido <pulido@lanl.gov>, James Ahrens <ahrens@lanl.gov>, Divya Banesh <dbanesh@lanl.gov>, Hugh Greenberg <hng@lanl.gov>, Pascal Grosset <pascalgrosset@lanl.gov>, Vedant Iyer <iyer@lanl.gov>, Benjamin Sims <bsims@lanl.gov>, Terece Turton <tlturton@lanl.gov>
 License-Expression: BSD-3-Clause

dsi_workflow-1.2.6/dsi/_version.py

@@ -0,0 +1 @@
+__version__ = "1.2.6"

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/dsi/backends/duckdb.py

@@ -407,7 +407,7 @@ class DuckDB(Filesystem):
         """
         data = None
         command = query.strip().split(None, 1)[0].lower()
-        if command in {"select", "pragma"}:
+        if command in {"select", "pragma", "with"}:
             try:
                 data = self.cur.execute(query).fetch_df()
                 if isVerbose:

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/dsi/backends/ndp.py

@@ -106,7 +106,7 @@ class NDP(Webserver):
         if self.api_key:
             self.headers["Authorization"] = self.api_key
 
-        # skip data retrieval if only checking connection to oceans11
+        # skip data retrieval if only checking connection to ndp
         if kwargs.get("only_validate", False):
             return
 

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/dsi/backends/osti.py

@@ -120,7 +120,7 @@ class OSTI(Webserver):
         if self.api_key:
             self.headers["Authorization"] = self.api_key
 
-        # skip data retrieval if only checking connection to oceans11
+        # skip data retrieval if only checking connection to osti
         if kwargs.get("only_validate", False):
             return
 

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/dsi/backends/rcsbpdb.py

@@ -230,7 +230,7 @@ class RCSBPDB(Webserver):
         self.params = params or {}
         self.validate_error_msg = None
 
-        # skip data retrieval if only checking connection to oceans11
+        # skip data retrieval if only checking connection to rcsbpdb
         if kwargs.get("only_validate", False):
             return
 
@@ -257,7 +257,7 @@ class RCSBPDB(Webserver):
 
         if self.validate_on_init and not self.validate_connection():
             self._loaded = False
-            raise ConnectionError(self.validate_error_msg or "RCSBPDB connection validation failed.")
+            raise ConnectionError(self.validate_error_msg or "Validating RCSBPDB connection failed.")
 
         if self.auto_load:
             try:
@@ -1081,13 +1081,19 @@ class RCSBPDB(Webserver):
     
     def get_table(self, table_name: str, dict_return=False):
         resolved = self._resolve_table_name(table_name)
+
+        if resolved not in self.schemas:
+            raise ValueError(
+                f"Table '{table_name}' not found. Available tables: {list(self.schemas.keys())}"
+            )
+
         table = self.tables.get(resolved, OrderedDict())
 
         if dict_return:
             return table
 
         if not table:
-            return pd.DataFrame(columns=self.get_schema(resolved))
+            return pd.DataFrame(columns=self.schemas.get(resolved, []))
 
         return pd.DataFrame(table)
 
@@ -1202,10 +1208,56 @@ class RCSBPDB(Webserver):
         return self.tables
 
     def get_schema(self, table_name: str | None = None):
-        if table_name is None:
-            return self.schemas
-        resolved = self._resolve_table_name(table_name)
-        return self.schemas.get(resolved, [])
+        """
+        Return SQL-style CREATE TABLE schema text.
+
+        Parameters
+        ----------
+        table_name : str or None
+            If provided, return schema text for one table.
+            If None, return schema text for all loaded/known tables.
+        """
+        table_names = [self._resolve_table_name(table_name)] if table_name else list(self.schemas.keys())
+        schema_lines = []
+
+        for name in table_names:
+            if name not in self.schemas:
+                raise ValueError(
+                    f"Table '{table_name}' not found. Available tables: {list(self.schemas.keys())}"
+                )
+
+            table = self.tables.get(name, OrderedDict())
+            cols = []
+
+            for col_name in self.schemas[name]:
+                dtype = "TEXT"
+
+                values = table.get(col_name, []) if isinstance(table, dict) else []
+
+                for value in values:
+                    if value is None:
+                        continue
+
+                    if isinstance(value, bool):
+                        dtype = "BOOLEAN"
+                    elif isinstance(value, int):
+                        dtype = "INTEGER"
+                    elif isinstance(value, float):
+                        dtype = "REAL"
+
+                    break
+
+                cols.append(f"    {col_name} {dtype}")
+
+            create_stmt = (
+                f"CREATE TABLE {name} (\n"
+                + ",\n".join(cols)
+                + "\n);"
+            )
+
+            schema_lines.append(create_stmt)
+
+        return "\n\n".join(schema_lines)
 
     def num_tables(self):
         table_count = len(self.tables)
@@ -1353,16 +1405,25 @@ class RCSBPDB(Webserver):
 
         return matches
 
-    def find_cell(self, query_object, **kwargs):
+    def find_cell(self, query_object, row=True, **kwargs):
         """
         Finds all cells that match the given query_object.
 
+        Parameters
+        ----------
+        query_object : Any
+            Value to search for.
+        row : bool, optional
+            If True, return full matching rows.
+            If False, return individual matching cells.
+
         Matching behavior:
         - Exact match for all data types
         - Case-insensitive partial match for strings
         - String representation match for complex objects
         """
         matches = []
+        seen_rows = set()
 
         is_str_query = isinstance(query_object, str)
         query_lower = query_object.lower() if is_str_query else None
@@ -1374,9 +1435,14 @@ class RCSBPDB(Webserver):
             cols = list(table_data.keys())
             df = pd.DataFrame(table_data)
 
-            for row_idx, row in df.iterrows():
+            if df.empty:
+                continue
+
+            for row_idx, row_data in df.iterrows():
+                row_matched = False
+
                 for col in cols:
-                    cell = row[col]
+                    cell = row_data[col]
                     match = False
 
                     if (
@@ -1389,21 +1455,42 @@ class RCSBPDB(Webserver):
                         )
                     ):
                         match = True
-
                     elif is_str_query and isinstance(cell, (dict, list, tuple)):
                         cell_str = str(cell).lower()
+
                         if query_lower in cell_str:
                             match = True
 
-                    if match:
-                        val = ValueObject()
-                        val.t_name = table_name
-                        val.c_name = cols
-                        val.row_num = row_idx
-                        val.value = row.tolist()
-                        val.type = "cell"
+                    if not match:
+                        continue
+
+                    if row:
+                        row_matched = True
+                        break
+
+                    val = ValueObject()
+                    val.t_name = table_name
+                    val.c_name = [col]
+                    val.row_num = int(row_idx)
+                    val.value = cell
+                    val.type = "cell"
+                    matches.append(val)
 
-                        matches.append(val)
+                if row and row_matched:
+                    key = (table_name, int(row_idx))
+
+                    if key in seen_rows:
+                        continue
+
+                    seen_rows.add(key)
+
+                    val = ValueObject()
+                    val.t_name = table_name
+                    val.c_name = cols
+                    val.row_num = int(row_idx)
+                    val.value = row_data.to_dict()
+                    val.type = "row"
+                    matches.append(val)
 
         return matches
 
@@ -1549,25 +1636,19 @@ class RCSBPDB(Webserver):
             print(f"  - num of rows: {df.shape[0]}")
         print()
 
-    def display(self, table_name=None, num_rows=25, display_cols=None, **kwargs):
+    def display(self, table_name, num_rows=25, display_cols=None, **kwargs):
         """
-        Print data from a table.
+        Return a preview DataFrame from a table.
 
         Parameters
         ----------
         table_name : str
             Name of the table to display.
         num_rows : int, optional
-            Number of rows to print.
+            Number of rows to return.
         display_cols : list[str], optional
             Specific columns to display.
         """
-        if table_name is None:
-            raise ValueError(
-                "display() requires a table_name. "
-                f"Available tables: {self.list(True)}"
-            )
-
         resolved = self._resolve_table_name(table_name)
 
         if resolved not in self.schemas:
@@ -1592,7 +1673,10 @@ class RCSBPDB(Webserver):
 
             df = df[display_cols]
 
-        print(df.head(num_rows).to_string(index=False))
+        if num_rows:
+            df = df.head(num_rows)
+
+        return df
 
 
     def summary(self, table_name=None, *args, **kwargs):

{dsi_workflow-1.2.4 → dsi_workflow-1.2.6}/dsi/backends/tests/test_rcsbpdb.py

@@ -257,9 +257,18 @@ def test_rcsbpdb_classify_usability():
 def test_rcsbpdb_initialization_no_auto_load(empty_backend):
     assert empty_backend._loaded is True
     assert empty_backend.list(True) == []
-    assert empty_backend.get_schema("datasets") == RCSBPDB.DATASET_SCHEMA
-    assert empty_backend.get_schema("resources") == RCSBPDB.RESOURCE_SCHEMA
-    assert empty_backend.get_schema("errors") == RCSBPDB.ERROR_SCHEMA
+
+    datasets_schema = empty_backend.get_schema("datasets")
+    resources_schema = empty_backend.get_schema("resources")
+    errors_schema = empty_backend.get_schema("errors")
+
+    assert isinstance(datasets_schema, str)
+    assert isinstance(resources_schema, str)
+    assert isinstance(errors_schema, str)
+
+    assert "CREATE TABLE datasets" in datasets_schema
+    assert "CREATE TABLE resources" in resources_schema
+    assert "CREATE TABLE errors" in errors_schema
 
 
 def test_rcsbpdb_initialization_with_identifiers(loaded_backend):
@@ -406,10 +415,26 @@ def test_rcsbpdb_get_schema(empty_backend):
     resource_schema = empty_backend.get_schema("resource")
     error_schema = empty_backend.get_schema("error")
 
-    assert isinstance(full_schema, dict)
-    assert dataset_schema == RCSBPDB.DATASET_SCHEMA
-    assert resource_schema == RCSBPDB.RESOURCE_SCHEMA
-    assert error_schema == RCSBPDB.ERROR_SCHEMA
+    assert isinstance(full_schema, str)
+    assert isinstance(dataset_schema, str)
+    assert isinstance(resource_schema, str)
+    assert isinstance(error_schema, str)
+
+    assert "CREATE TABLE datasets" in full_schema
+    assert "CREATE TABLE resources" in full_schema
+    assert "CREATE TABLE errors" in full_schema
+
+    assert "CREATE TABLE datasets" in dataset_schema
+    assert "dataset_id" in dataset_schema
+    assert "doi" in dataset_schema
+
+    assert "CREATE TABLE resources" in resource_schema
+    assert "resource_id" in resource_schema
+    assert "download_url" in resource_schema
+
+    assert "CREATE TABLE errors" in error_schema
+    assert "identifier" in error_schema
+    assert "status" in error_schema
 
 
 def test_rcsbpdb_table_name_resolution(empty_backend):
@@ -582,7 +607,7 @@ def test_rcsbpdb_find_searches_table_column_and_cell_values(loaded_backend):
 
     assert isinstance(matches, list)
     assert len(matches) >= 1
-    assert any(match.type == "cell" for match in matches)
+    assert any(match.type == "row" for match in matches)
     assert any(match.t_name == "datasets" for match in matches)
 
 
@@ -624,7 +649,7 @@ def test_rcsbpdb_find_cell_matches_value(loaded_backend):
 
     assert isinstance(matches, list)
     assert len(matches) >= 1
-    assert all(match.type == "cell" for match in matches)
+    assert all(match.type == "row" for match in matches)
 
     dataset_matches = [
         match for match in matches
@@ -632,8 +657,28 @@ def test_rcsbpdb_find_cell_matches_value(loaded_backend):
     ]
 
     assert len(dataset_matches) >= 1
-    assert any("1CBS" in match.value for match in dataset_matches)
+    assert any(
+        match.value["dataset_id"] == "1CBS"
+        for match in dataset_matches
+    )
+
+def test_rcsbpdb_find_cell_cell_type_support(loaded_backend):
+    matches = loaded_backend.find_cell("1CBS", row=False)
+
+    assert isinstance(matches, list)
+    assert len(matches) >= 1
+    assert all(match.type == "cell" for match in matches)
 
+    dataset_matches = [
+        match for match in matches
+        if match.t_name == "datasets"
+    ]
+
+    assert len(dataset_matches) >= 1
+    assert any(
+        match.c_name == ["dataset_id"] and match.value == "1CBS"
+        for match in dataset_matches
+    )
 
 def test_rcsbpdb_find_relation_condition_string(loaded_backend):
     matches = loaded_backend.find_relation("dataset_id = 1CBS")
@@ -842,11 +887,25 @@ def test_rcsbpdb_summary_single_table(loaded_backend):
     assert "dataset_id" in set(summary["column"])
 
 
-def test_rcsbpdb_display_returns_none(loaded_backend):
+def test_rcsbpdb_display_returns_dataframe(loaded_backend):
     result = loaded_backend.display("datasets", num_rows=1)
 
-    assert result is None
+    assert isinstance(result, pd.DataFrame)
+    assert len(result) == 1
 
+def test_rcsbpdb_display_cols(loaded_backend):
+    result = loaded_backend.display(
+        "datasets",
+        num_rows=2,
+        display_cols=["dataset_id", "title"],
+    )
+
+    assert isinstance(result, pd.DataFrame)
+    assert list(result.columns) == ["dataset_id", "title"]
+
+def test_rcsbpdb_display_requires_table_name(loaded_backend):
+    with pytest.raises(TypeError):
+        loaded_backend.display()
 
 def test_rcsbpdb_display_missing_table_raises_value_error(empty_backend):
     with pytest.raises(ValueError):