From 947718458294157534d76affe6159e8efeed16cb Mon Sep 17 00:00:00 2001
From: Josh Hayes <35790761+hayescode@users.noreply.github.com>
Date: Fri, 5 Sep 2025 05:52:09 -0500
Subject: [PATCH] databricks bug fixes (#3416)

---
 mem0/vector_stores/databricks.py       | 4 ++--
 tests/vector_stores/test_databricks.py | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mem0/vector_stores/databricks.py b/mem0/vector_stores/databricks.py
index 0a5b06f49..6b5660e74 100644
--- a/mem0/vector_stores/databricks.py
+++ b/mem0/vector_stores/databricks.py
@@ -465,7 +465,7 @@ class Databricks(VectorStoreBase):
 
             # Parse results
             result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
-            data_array = result_data.data_array if hasattr(result_data, "data_array") else []
+            data_array = result_data.data_array if getattr(result_data, "data_array", None) else []
 
             memory_results = []
             for row in data_array:
@@ -708,7 +708,7 @@ class Databricks(VectorStoreBase):
                         pass
                 memory_id = row_dict.get("memory_id") or row_dict.get("id")
                 memory_results.append(MemoryResult(id=memory_id, payload=payload))
-            return memory_results
+            return [memory_results]
         except Exception as e:
             logger.error(f"Failed to list memories: {e}")
             return []
diff --git a/tests/vector_stores/test_databricks.py b/tests/vector_stores/test_databricks.py
index 5b87def07..9b2d1c560 100644
--- a/tests/vector_stores/test_databricks.py
+++ b/tests/vector_stores/test_databricks.py
@@ -300,8 +300,9 @@ def test_list_memories(db_instance_delta, mock_workspace_client):
         result=SimpleNamespace(data_array=[row])
     )
     res = db_instance_delta.list(limit=1)
-    assert len(res) == 1
-    assert res[0].id == "id3"
+    assert isinstance(res, list)
+    assert len(res[0]) == 1
+    assert res[0][0].id == "id3"
 
 
 # ---------------------- Reset Tests ---------------------- #