-
Notifications
You must be signed in to change notification settings - Fork 120
streaming ingestion support for PUT operation #643
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,34 @@ | ||
#!/usr/bin/env python3 | ||
""" | ||
Simple example of streaming PUT operations. | ||
|
||
This demonstrates the basic usage of streaming PUT with the __input_stream__ token. | ||
""" | ||
|
||
import io | ||
import os | ||
from databricks import sql | ||
|
||
with sql.connect( | ||
server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"), | ||
http_path=os.getenv("DATABRICKS_HTTP_PATH"), | ||
access_token=os.getenv("DATABRICKS_TOKEN"), | ||
) as connection: | ||
|
||
with connection.cursor() as cursor: | ||
# Create a simple data stream | ||
data = b"Hello, streaming world!" | ||
stream = io.BytesIO(data) | ||
|
||
# Get catalog, schema, and volume from environment variables | ||
catalog = os.getenv("DATABRICKS_CATALOG") | ||
schema = os.getenv("DATABRICKS_SCHEMA") | ||
volume = os.getenv("DATABRICKS_VOLUME") | ||
|
||
# Upload to Unity Catalog volume | ||
cursor.execute( | ||
f"PUT '__input_stream__' INTO '/Volumes/{catalog}/{schema}/{volume}/hello.txt' OVERWRITE", | ||
input_stream=stream | ||
) | ||
|
||
print("File uploaded successfully!") |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,65 @@ | ||
#!/usr/bin/env python3 | ||
""" | ||
E2E tests for streaming PUT operations. | ||
""" | ||
|
||
import io | ||
import logging | ||
import pytest | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
class PySQLStreamingPutTestSuiteMixin: | ||
"""Test suite for streaming PUT operations.""" | ||
|
||
def test_streaming_put_basic(self, catalog, schema): | ||
sreekanth-db marked this conversation as resolved.
Show resolved
Hide resolved
|
||
"""Test basic streaming PUT functionality.""" | ||
|
||
# Create test data | ||
test_data = b"Hello, streaming world! This is test data." | ||
filename = "streaming_put_test.txt" | ||
file_path = f"/Volumes/{catalog}/{schema}/e2etests/{filename}" | ||
|
||
try: | ||
with self.connection() as conn: | ||
with conn.cursor() as cursor: | ||
self._cleanup_test_file(file_path) | ||
|
||
with io.BytesIO(test_data) as stream: | ||
cursor.execute( | ||
f"PUT '__input_stream__' INTO '{file_path}'", | ||
input_stream=stream | ||
) | ||
|
||
# Verify file exists | ||
cursor.execute(f"LIST '/Volumes/{catalog}/{schema}/e2etests/'") | ||
files = cursor.fetchall() | ||
|
||
# Check if our file is in the list | ||
file_paths = [row[0] for row in files] | ||
assert file_path in file_paths, f"File {file_path} not found in {file_paths}" | ||
finally: | ||
self._cleanup_test_file(file_path) | ||
|
||
def test_streaming_put_missing_stream(self, catalog, schema): | ||
"""Test that missing stream raises appropriate error.""" | ||
|
||
with self.connection() as conn: | ||
with conn.cursor() as cursor: | ||
# Test without providing stream | ||
with pytest.raises(Exception): # Should fail | ||
cursor.execute( | ||
f"PUT '__input_stream__' INTO '/Volumes/{catalog}/{schema}/e2etests/test.txt'" | ||
# Note: No input_stream parameter | ||
) | ||
|
||
def _cleanup_test_file(self, file_path): | ||
"""Clean up a test file if it exists.""" | ||
try: | ||
with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn: | ||
with conn.cursor() as cursor: | ||
cursor.execute(f"REMOVE '{file_path}'") | ||
logger.info("Successfully cleaned up test file: %s", file_path) | ||
except Exception as e: | ||
logger.error("Cleanup failed for %s: %s", file_path, e) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,113 @@ | ||
import io | ||
from unittest.mock import patch, Mock, MagicMock | ||
|
||
import pytest | ||
|
||
sreekanth-db marked this conversation as resolved.
Show resolved
Hide resolved
|
||
import databricks.sql.client as client | ||
|
||
|
||
class TestStreamingPut: | ||
"""Unit tests for streaming PUT functionality.""" | ||
|
||
@pytest.fixture | ||
def cursor(self): | ||
return client.Cursor(connection=Mock(), backend=Mock()) | ||
|
||
def _setup_mock_staging_put_stream_response(self, mock_backend): | ||
"""Helper method to set up mock staging PUT stream response.""" | ||
mock_result_set = Mock() | ||
mock_result_set.is_staging_operation = True | ||
mock_backend.execute_command.return_value = mock_result_set | ||
|
||
mock_row = Mock() | ||
mock_row.operation = "PUT" | ||
mock_row.localFile = "__input_stream__" | ||
mock_row.presignedUrl = "https://example.com/upload" | ||
mock_row.headers = "{}" | ||
mock_result_set.fetchone.return_value = mock_row | ||
|
||
return mock_result_set | ||
|
||
def test_execute_with_valid_stream(self, cursor): | ||
"""Test execute method with valid input stream.""" | ||
|
||
# Mock the backend response | ||
self._setup_mock_staging_put_stream_response(cursor.backend) | ||
|
||
# Test with valid stream | ||
test_stream = io.BytesIO(b"test data") | ||
|
||
with patch.object(cursor, "_handle_staging_put_stream") as mock_handler: | ||
cursor.execute( | ||
"PUT '__input_stream__' INTO '/Volumes/test/cat/schema/vol/file.txt'", | ||
input_stream=test_stream, | ||
) | ||
|
||
# Verify staging handler was called | ||
mock_handler.assert_called_once() | ||
|
||
def test_execute_with_none_stream_for_staging_put(self, cursor): | ||
"""Test execute method rejects None stream for streaming PUT operations.""" | ||
|
||
# Mock staging operation response for None case | ||
self._setup_mock_staging_put_stream_response(cursor.backend) | ||
|
||
# None with __input_stream__ raises ProgrammingError | ||
with pytest.raises(client.ProgrammingError) as excinfo: | ||
cursor.execute( | ||
"PUT '__input_stream__' INTO '/Volumes/test/cat/schema/vol/file.txt'", | ||
input_stream=None, | ||
) | ||
error_msg = str(excinfo.value) | ||
assert "No input stream provided for streaming operation" in error_msg | ||
|
||
def test_handle_staging_put_stream_success(self, cursor): | ||
"""Test successful streaming PUT operation.""" | ||
|
||
presigned_url = "https://example.com/upload" | ||
headers = {"Content-Type": "text/plain"} | ||
|
||
with patch.object( | ||
cursor.connection.http_client, "request" | ||
) as mock_http_request: | ||
mock_response = MagicMock() | ||
mock_response.status = 200 | ||
mock_response.data = b"success" | ||
mock_http_request.return_value = mock_response | ||
|
||
test_stream = io.BytesIO(b"test data") | ||
cursor._handle_staging_put_stream( | ||
presigned_url=presigned_url, stream=test_stream, headers=headers | ||
) | ||
|
||
# Verify the HTTP client was called correctly | ||
mock_http_request.assert_called_once() | ||
call_args = mock_http_request.call_args | ||
# Check positional arguments: (method, url, body=..., headers=...) | ||
assert call_args[0][0].value == "PUT" # First positional arg is method | ||
assert call_args[0][1] == presigned_url # Second positional arg is url | ||
# Check keyword arguments | ||
assert call_args[1]["body"] == b"test data" | ||
assert call_args[1]["headers"] == headers | ||
|
||
def test_handle_staging_put_stream_http_error(self, cursor): | ||
"""Test streaming PUT operation with HTTP error.""" | ||
|
||
presigned_url = "https://example.com/upload" | ||
|
||
with patch.object( | ||
cursor.connection.http_client, "request" | ||
) as mock_http_request: | ||
mock_response = MagicMock() | ||
mock_response.status = 500 | ||
mock_response.data = b"Internal Server Error" | ||
mock_http_request.return_value = mock_response | ||
|
||
test_stream = io.BytesIO(b"test data") | ||
with pytest.raises(client.OperationalError) as excinfo: | ||
cursor._handle_staging_put_stream( | ||
presigned_url=presigned_url, stream=test_stream | ||
) | ||
|
||
# Check for the actual error message format | ||
assert "500" in str(excinfo.value) |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.