Repository: piskvorky/smart_open
Pull request #826 (2024-07-02 00:45:37)

Updates readline logic for azure to match s3
#### Title

Updates readline buffer management for azure, improving performance.

#### Motivation

Observed slow read times when iterating using `readlines()` from a 100MB csv file in Azure Blob Storage. Simple copies and reads were fine, slowness only observed with some code that was explicitly calling `readlines()`.

Borrows from the logic found in the S3 implementation, which was updated to be more efficient and does not perform as many scans on the local chunk of the remote file.

#### Tests

Confirmed clean run with `pytest -k test_azure`

#### Checklist

Before you create the PR, please make sure you have:

- [x] Picked a concise, informative and complete title
- [x] Clearly explained the motivation behind the PR
- [x] Linked to any existing issues that your PR will be solving
- [x] Included tests for any new functionality
- [x] Checked that all unit tests pass

#### Workflow

Please avoid rebasing and force-pushing to the branch of the PR once a review is in progress.
Rebasing can make your commits look a bit cleaner, but it also makes life more difficult from the reviewer, because they are no longer able to distinguish between code that has already been reviewed, and unreviewed code.


ddelange: @mpenkov can you review and merge?

Patch:
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 5863597..cb0d124 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -1,5 +1,8 @@
 name: Test
 on: [push, pull_request]
+concurrency:  # https://stackoverflow.com/questions/66335225#comment133398800_72408109
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 jobs:
   linters:
     runs-on: ubuntu-latest
@@ -12,10 +15,10 @@ jobs:
           python-version: "3.11"
 
       - name: Update pip
-        run: python -m pip install -U pip
+        run: pip install -U pip
 
       - name: Install dependencies
-        run: python -m pip install flake8
+        run: pip install flake8
 
       - name: Run flake8 linter (source)
         run: flake8 --show-source smart_open
@@ -26,17 +29,17 @@ jobs:
     strategy:
       matrix:
         include:
-          - {python: '3.8', os: ubuntu-20.04}
-          - {python: '3.9', os: ubuntu-20.04}
-          - {python: '3.10', os: ubuntu-20.04}
-          - {python: '3.11', os: ubuntu-20.04}
-          - {python: '3.12', os: ubuntu-20.04}
-
-          - {python: '3.8', os: windows-2019}
-          - {python: '3.9', os: windows-2019}
-          - {python: '3.10', os: windows-2019}
-          - {python: '3.11', os: windows-2019}
-          - {python: '3.12', os: windows-2019}
+          - {python-version: '3.8', os: ubuntu-20.04}
+          - {python-version: '3.9', os: ubuntu-20.04}
+          - {python-version: '3.10', os: ubuntu-20.04}
+          - {python-version: '3.11', os: ubuntu-20.04}
+          - {python-version: '3.12', os: ubuntu-20.04}
+
+          - {python-version: '3.8', os: windows-2019}
+          - {python-version: '3.9', os: windows-2019}
+          - {python-version: '3.10', os: windows-2019}
+          - {python-version: '3.11', os: windows-2019}
+          - {python-version: '3.12', os: windows-2019}
     steps:
       - uses: actions/checkout@v2
 
@@ -45,13 +48,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Update pip
-        run: python -m pip install -U pip
-
-      #
-      # https://askubuntu.com/questions/1428181/module-lib-has-no-attribute-x509-v-flag-cb-issuer-check
-      #
-      - name: Upgrade PyOpenSSL
-        run: python -m pip install pyOpenSSL --upgrade
+        run: pip install -U pip
 
       - name: Install smart_open without dependencies
         run: pip install -e .
@@ -71,20 +68,20 @@ jobs:
     strategy:
       matrix:
         include:
-          - {python: '3.8', os: ubuntu-20.04}
-          - {python: '3.9', os: ubuntu-20.04}
-          - {python: '3.10', os: ubuntu-20.04}
-          - {python: '3.11', os: ubuntu-20.04}
-          - {python: '3.12', os: ubuntu-20.04}
+          - {python-version: '3.8', os: ubuntu-20.04}
+          - {python-version: '3.9', os: ubuntu-20.04}
+          - {python-version: '3.10', os: ubuntu-20.04}
+          - {python-version: '3.11', os: ubuntu-20.04}
+          - {python-version: '3.12', os: ubuntu-20.04}
 
           #
           # Some of the doctests don't pass on Windows because of Windows-specific
           # character encoding issues.
           #
-          # - {python: '3.7', os: windows-2019}
-          # - {python: '3.8', os: windows-2019}
-          # - {python: '3.9', os: windows-2019}
-          # - {python: '3.10', os: windows-2019}
+          # - {python-version: '3.7', os: windows-2019}
+          # - {python-version: '3.8', os: windows-2019}
+          # - {python-version: '3.9', os: windows-2019}
+          # - {python-version: '3.10', os: windows-2019}
 
     steps:
       - uses: actions/checkout@v2
@@ -94,10 +91,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Update pip
-        run: python -m pip install -U pip
-
-      - name: Upgrade PyOpenSSL
-        run: python -m pip install pyOpenSSL --upgrade
+        run: pip install -U pip
 
       - name: Install smart_open and its dependencies
         run: pip install -e .[test]
@@ -114,18 +108,18 @@ jobs:
     strategy:
       matrix:
         include:
-          - {python: '3.8', os: ubuntu-20.04}
-          - {python: '3.9', os: ubuntu-20.04}
-          - {python: '3.10', os: ubuntu-20.04}
-          - {python: '3.11', os: ubuntu-20.04}
-          - {python: '3.12', os: ubuntu-20.04}
+          - {python-version: '3.8', os: ubuntu-20.04}
+          - {python-version: '3.9', os: ubuntu-20.04}
+          - {python-version: '3.10', os: ubuntu-20.04}
+          - {python-version: '3.11', os: ubuntu-20.04}
+          - {python-version: '3.12', os: ubuntu-20.04}
 
           # Not sure why we exclude these, perhaps for historical reasons?
           #
-          # - {python: '3.7', os: windows-2019}
-          # - {python: '3.8', os: windows-2019}
-          # - {python: '3.9', os: windows-2019}
-          # - {python: '3.10', os: windows-2019}
+          # - {python-version: '3.7', os: windows-2019}
+          # - {python-version: '3.8', os: windows-2019}
+          # - {python-version: '3.9', os: windows-2019}
+          # - {python-version: '3.10', os: windows-2019}
 
     steps:
       - uses: actions/checkout@v2
@@ -135,12 +129,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Update pip
-        run: python -m pip install -U pip
-
-      - name: Upgrade PyOpenSSL
-        run: python -m pip install pyOpenSSL --upgrade
-
-      - run: python -m pip install numpy
+        run: pip install -U pip
 
       - name: Install smart_open and its dependencies
         run: pip install -e .[test]
@@ -148,7 +137,9 @@ jobs:
       - run: bash ci_helpers/helpers.sh enable_moto_server
         if: ${{ matrix.moto_server }}
 
-      - run: |
+      - name: Start vsftpd
+        timeout-minutes: 2
+        run: |
           sudo apt-get install vsftpd
           sudo bash ci_helpers/helpers.sh create_ftp_ftps_servers
 
@@ -169,16 +160,16 @@ jobs:
     strategy:
       matrix:
         include:
-          - {python: '3.8', os: ubuntu-20.04}
-          - {python: '3.9', os: ubuntu-20.04}
-          - {python: '3.10', os: ubuntu-20.04}
-          - {python: '3.11', os: ubuntu-20.04}
-          - {python: '3.12', os: ubuntu-20.04}
+          - {python-version: '3.8', os: ubuntu-20.04}
+          - {python-version: '3.9', os: ubuntu-20.04}
+          - {python-version: '3.10', os: ubuntu-20.04}
+          - {python-version: '3.11', os: ubuntu-20.04}
+          - {python-version: '3.12', os: ubuntu-20.04}
 
-          # - {python: '3.7', os: windows-2019}
-          # - {python: '3.8', os: windows-2019}
-          # - {python: '3.9', os: windows-2019}
-          # - {python: '3.10', os: windows-2019}
+          # - {python-version: '3.7', os: windows-2019}
+          # - {python-version: '3.8', os: windows-2019}
+          # - {python-version: '3.9', os: windows-2019}
+          # - {python-version: '3.10', os: windows-2019}
 
     steps:
       - uses: actions/checkout@v2
@@ -188,16 +179,11 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Update pip
-        run: python -m pip install -U pip
-
-      - name: Upgrade PyOpenSSL
-        run: python -m pip install pyOpenSSL --upgrade
+        run: pip install -U pip
 
       - name: Install smart_open and its dependencies
         run: pip install -e .[test]
 
-      - run: pip install awscli pytest_benchmark
-
       - name: Run benchmarks
         run: python ci_helpers/run_benchmarks.py
         env:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 85fc7d3..575f1f5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -482,3 +482,4 @@ The old `smart_open.smart_open` function is deprecated, but continues to work as
 
   - support for multistream bzip files (PR #9, @pombredanne)
   - introduce this CHANGELOG
+
diff --git a/ci_helpers/helpers.sh b/ci_helpers/helpers.sh
index 4533b33..e0accc7 100644
--- a/ci_helpers/helpers.sh
+++ b/ci_helpers/helpers.sh
@@ -20,6 +20,9 @@ create_ftp_ftps_servers(){
   mkdir $home_dir
   useradd -p $(echo $pass | openssl passwd -1 -stdin) -d $home_dir $user
   chown $user:$user $home_dir
+  openssl req -x509 -nodes -new -sha256 -days 10240 -newkey rsa:2048 -keyout /etc/vsftpd.key -out /etc/vsftpd.pem -subj "/C=ZA/CN=localhost"
+  chmod 755 /etc/vsftpd.key
+  chmod 755 /etc/vsftpd.pem
 
   server_setup='''
 listen=YES
@@ -32,6 +35,8 @@ chroot_local_user=YES
 allow_writeable_chroot=YES'''
 
   additional_ssl_setup='''
+rsa_cert_file=/etc/vsftpd.pem
+rsa_private_key_file=/etc/vsftpd.key
 ssl_enable=YES
 allow_anon_ssl=NO
 force_local_data_ssl=NO
diff --git a/setup.py b/setup.py
index a9a4fc5..9e738be 100644
--- a/setup.py
+++ b/setup.py
@@ -47,9 +47,12 @@ all_deps = aws_deps + gcs_deps + azure_deps + http_deps + ssh_deps + zst_deps
 tests_require = all_deps + [
     'moto[server]',
     'responses',
-    'boto3',
     'pytest',
     'pytest-rerunfailures',
+    'pytest_benchmark',
+    'awscli',
+    'pyopenssl',
+    'numpy',
 ]
 
 setup(
diff --git a/smart_open/azure.py b/smart_open/azure.py
index 5cac221..1c991f0 100644
--- a/smart_open/azure.py
+++ b/smart_open/azure.py
@@ -325,24 +325,22 @@ class Reader(io.BufferedIOBase):
         """Read up to and including the next newline.  Returns the bytes read."""
         if limit != -1:
             raise NotImplementedError('limits other than -1 not implemented yet')
-        the_line = io.BytesIO()
+
+        #
+        # A single line may span multiple buffers.
+        #
+        line = io.BytesIO()
         while not (self._position == self._size and len(self._current_part) == 0):
-            #
-            # In the worst case, we're reading the unread part of self._current_part
-            # twice here, once in the if condition and once when calling index.
-            #
-            # This is sub-optimal, but better than the alternative: wrapping
-            # .index in a try..except, because that is slower.
-            #
-            remaining_buffer = self._current_part.peek()
-            if self._line_terminator in remaining_buffer:
-                next_newline = remaining_buffer.index(self._line_terminator)
-                the_line.write(self._read_from_buffer(next_newline + 1))
+            line_part = self._current_part.readline(self._line_terminator)
+            line.write(line_part)
+            self._position += len(line_part)
+
+            if line_part.endswith(self._line_terminator):
                 break
             else:
-                the_line.write(self._read_from_buffer())
                 self._fill_buffer()
-        return the_line.getvalue()
+
+        return line.getvalue()
 
     #
     # Internal methods.
diff --git a/smart_open/ftp.py b/smart_open/ftp.py
index 40d54ca..a7212ec 100644
--- a/smart_open/ftp.py
+++ b/smart_open/ftp.py
@@ -86,7 +86,7 @@ def convert_transport_params_to_args(transport_params):
 def _connect(hostname, username, port, password, secure_connection, transport_params):
     kwargs = convert_transport_params_to_args(transport_params)
     if secure_connection:
-        ssl_context = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
+        ssl_context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
         ftp = FTP_TLS(context=ssl_context, **kwargs)
     else:
         ftp = FTP(**kwargs)
diff --git a/smart_open/s3.py b/smart_open/s3.py
index 09433ad..60ae2a9 100644
--- a/smart_open/s3.py
+++ b/smart_open/s3.py
@@ -1209,8 +1209,8 @@ class SinglepartWriter(io.BufferedIOBase):
         return length
 
     def terminate(self):
-        """Nothing to cancel in single-part uploads."""
-        return
+        self._buf = None
+        logger.debug('%s: terminated singlepart upload', self)
 
     #
     # Internal methods.
diff --git a/smart_open/utils.py b/smart_open/utils.py
index 2be57d1..c06e2cf 100644
--- a/smart_open/utils.py
+++ b/smart_open/utils.py
@@ -208,6 +208,8 @@ class TextIOWrapper(io.TextIOWrapper):
 
 
 class FileLikeProxy(wrapt.ObjectProxy):
+    __inner = ...  # initialized before wrapt disallows __setattr__ on certain objects
+
     def __init__(self, outer, inner):
         super().__init__(outer)
         self.__inner = inner


Test patch:
diff --git a/integration-tests/test_ftp.py b/integration-tests/test_ftp.py
index 000faae..94b4e03 100644
--- a/integration-tests/test_ftp.py
+++ b/integration-tests/test_ftp.py
@@ -1,6 +1,11 @@
 from __future__ import unicode_literals
 import pytest
 from smart_open import open
+import ssl
+from functools import partial
+
+# localhost has self-signed cert, see ci_helpers/helpers.sh:create_ftp_ftps_servers
+ssl.create_default_context = partial(ssl.create_default_context, cafile="/etc/vsftpd.pem")
 
 
 @pytest.fixture(params=[("ftp", 21), ("ftps", 90)])
@@ -81,4 +86,4 @@ def test_line_endings_binary(server_info):
     
     with open(f"{server_type}://user:123@localhost:{port_num}/file4", "rb") as f:    
         for line in f:
-            assert B_CLRF in line
\ No newline at end of file
+            assert B_CLRF in line
diff --git a/smart_open/tests/test_http.py b/smart_open/tests/test_http.py
index f4f8338..9433043 100644
--- a/smart_open/tests/test_http.py
+++ b/smart_open/tests/test_http.py
@@ -21,24 +21,22 @@ BYTES = b'i tried so hard and got so far but in the end it doesn\'t even matter'
 URL = 'http://localhost'
 HTTPS_URL = 'https://localhost'
 HEADERS = {
-    'Content-Length': str(len(BYTES)),
     'Accept-Ranges': 'bytes',
 }
 
 
 def request_callback(request, headers=HEADERS, data=BYTES):
-    try:
-        range_string = request.headers['range']
-    except KeyError:
-        return (200, headers, data)
+    headers = headers.copy()
+    range_string = request.headers.get('range', 'bytes=0-')
 
-    start, end = range_string.replace('bytes=', '').split('-', 1)
+    start, end = range_string.replace('bytes=', '', 1).split('-', 1)
     start = int(start)
-    if end:
-        end = int(end)
-    else:
-        end = len(data)
-    return (200, headers, data[start:end])
+    end = int(end) if end else len(data)
+
+    data = data[start:end]
+    headers['Content-Length'] = str(len(data))
+
+    return (200, headers, data)
 
 
 @unittest.skipIf(os.environ.get('TRAVIS'), 'This test does not work on TravisCI for some reason')
diff --git a/smart_open/tests/test_s3.py b/smart_open/tests/test_s3.py
index ff44ad4..61aec98 100644
--- a/smart_open/tests/test_s3.py
+++ b/smart_open/tests/test_s3.py
@@ -795,6 +795,26 @@ class SinglepartWriterTest(unittest.TestCase):
         with smart_open.s3.open(BUCKET_NAME, 'key', 'wb', multipart_upload=False) as fout:
             assert str(fout) == "smart_open.s3.SinglepartWriter('test-smartopen', 'key')"
 
+    def test_ensure_no_side_effects_on_exception(self):
+        class WriteError(Exception):
+            pass
+
+        s3_resource = _resource("s3")
+        obj = s3_resource.Object(BUCKET_NAME, KEY_NAME)
+
+        # wrap in closure to ease writer dereferencing
+        def _run():
+            with smart_open.s3.open(BUCKET_NAME, obj.key, "wb", multipart_upload=False) as fout:
+                fout.write(b"this should not be written")
+                raise WriteError
+
+        try:
+            _run()
+        except WriteError:
+            pass
+        finally:
+            self.assertRaises(s3_resource.meta.client.exceptions.NoSuchKey, obj.get)
+
 
 ARBITRARY_CLIENT_ERROR = botocore.client.ClientError(error_response={}, operation_name='bar')
 
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index e31f48b..789f44c 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -607,6 +607,22 @@ class RealFileSystemTests(unittest.TestCase):
         self.assertEqual(text, SAMPLE_TEXT * 2)
 
 
+class CompressionRealFileSystemTests(RealFileSystemTests):
+    """Same as RealFileSystemTests but with a compressed file."""
+
+    def setUp(self):
+        with named_temporary_file(prefix='test', suffix='.zst', delete=False) as fout:
+            self.temp_file = fout.name
+        with smart_open.open(self.temp_file, 'wb') as fout:
+            fout.write(SAMPLE_BYTES)
+
+    def test_aplus(self):
+        pass  # transparent (de)compression unsupported for mode 'ab+'
+
+    def test_atplus(self):
+        pass  # transparent (de)compression unsupported for mode 'ab+'
+
+
 #
 # What exactly to patch here differs on _how_ we're opening the file.
 # See the _shortcut_open function for details.