test.py

from contextlib import contextmanager
from ctypes import cdll
from ctypes.util import find_library
import datetime
import functools
import hashlib
import hmac
import os
import socket
import sqlite3
import tempfile
import threading
import unittest
import urllib.parse
import uuid

import httpx
from httpx import HTTPStatusError

from sqlite_s3_query import (
    VersioningNotEnabledError,
    SQLiteError,
    QueryContextClosedError,
    sqlite_s3_query,
    sqlite_s3_query_multi,
)


class TestSqliteS3Query(unittest.TestCase):

    def test_sqlite3_installed_on_ci(self):
        ci = os.environ.get('CI', '')
        sqlite3_version = os.environ.get('SQLITE3_VERSION', 'default')
        if ci and sqlite3_version != 'default':
            libsqlite3 = get_libsqlite3()
            self.assertEqual(libsqlite3.sqlite3_libversion_number(), int(sqlite3_version))

    def test_without_versioning(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
        ]) as db:
            put_object_without_versioning('bucket-without-versioning', 'my.db', db)

        with self.assertRaisesRegex(VersioningNotEnabledError, 'The bucket must have versioning enabled'):
            sqlite_s3_query('http://localhost:9000/bucket-without-versioning/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3).__enter__()

    def test_select(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),())
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with query('SELECT my_col_a FROM my_table') as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [('some-text-a',)] * 500)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with \
                    query('SELECT my_col_a FROM my_table') as (columns_a, rows_a), \
                    query('SELECT my_col_b FROM my_table') as (columns_b, rows_b):

                rows = [
                    (next(rows_a)[0], next(rows_b)[0])
                    for i in range(0, 500)
                ]

        self.assertEqual(rows, [('some-text-a','some-text-b')] * 500)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with get_db([
                ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ()),
            ] + [
                ("INSERT INTO my_table VALUES " + ','.join(["('some-new-a', 'some-new-b')"] * 500), ()),
            ]) as db:
                put_object_with_versioning('my-bucket', 'my.db', db)

            with query('SELECT my_col_a FROM my_table') as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [('some-text-a',)] * 500)

        with self.assertRaisesRegex(QueryContextClosedError, 'Attempting to use finalized statement'):
            with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                with query('SELECT my_col_a FROM my_table') as (columns, rows):
                    for row in rows:
                        break
                next(rows)

        with self.assertRaisesRegex(QueryContextClosedError, 'Attempting to use finalized statement'):
            with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                with query('SELECT my_col_a FROM my_table') as (columns, rows):
                    pass
                next(rows)

    def test_select_with_named_params(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ())
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500), ()),
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-c', 'some-text-d')"] * 100), ()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with query('SELECT COUNT(*) FROM my_table WHERE my_col_a = :first', named_params=((':first', 'some-text-a'),)) as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [(500,)])

    def test_select_with_named_params_public_bucket(self):
        create_bucket('my-public-bucket')
        disable_auth('my-public-bucket')
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ())
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500), ()),
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-c', 'some-text-d')"] * 100), ()),
        ]) as db:
            put_object_with_versioning('my-public-bucket', 'my.db', db)

        with sqlite_s3_query(
                'http://localhost:9000/my-public-bucket/my.db',
                get_credentials=None,
                get_libsqlite3=get_libsqlite3
        ) as query:
            with query('SELECT COUNT(*) FROM my_table WHERE my_col_a = :first', named_params=((':first', 'some-text-a'),)) as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [(500,)])

    def test_select_large(self):
        empty = (bytes(4050),)

        def sqls():
            yield ("CREATE TABLE foo(content BLOB);",())
            for _ in range(0, 1200000):
                yield ("INSERT INTO foo VALUES (?);", empty)

        with get_db(sqls()) as db:
            length = 0
            for chunk in db():
                length += len(chunk)
            self.assertGreater(length, 4294967296)
            put_object_with_versioning('my-bucket', 'my.db', db)

        count = 0
        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with query('SELECT content FROM foo ORDER BY rowid LIMIT 1') as (columns, rows):
                for _ in rows:
                    count += 1

            self.assertEqual(count, 1)

            count = 0
            with query('SELECT content FROM foo ORDER BY rowid DESC LIMIT 1') as (columns, rows):
                for _ in rows:
                    count += 1

            self.assertEqual(count, 1)

    def test_select_multi(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ())
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500), ()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            rows_list = [
                list(rows)
                for (columns, rows) in query('''
                    SELECT my_col_a FROM my_table;
                    SELECT my_col_a FROM my_table LIMIT 10;
                ''')
            ]

        self.assertEqual(rows_list, [[('some-text-a',)] * 500, [('some-text-a',)] * 10])

        with self.assertRaisesRegex(Exception, 'Just after creating context'):
            with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                raise Exception('Just after creating context')

        with self.assertRaisesRegex(Exception, 'Just after iterating statements'):
            with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                for (columns, rows) in query('''
                    SELECT my_col_a FROM my_table;
                    SELECT my_col_a FROM my_table LIMIT 10;
                '''):
                    raise Exception('Just after iterating statements')

        with self.assertRaisesRegex(Exception, 'Just after iterating first row'):
            with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                for (columns, rows) in query('''
                    SELECT my_col_a FROM my_table;
                    SELECT my_col_a FROM my_table LIMIT 10;
                '''):
                    rows_it = iter(rows)
                    next(rows_it)
                    raise Exception('Just after iterating first row')

        with self.assertRaisesRegex(Exception, 'Multiple open statements'):
            with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                it = iter(query('''
                    SELECT my_col_a FROM my_table;
                    SELECT my_col_a FROM my_table LIMIT 10;
                '''))
                columns_1, rows_1 = next(it)
                rows_1_it = iter(rows_1)
                next(rows_1_it)

                columns_2, rows_2 = next(it)
                rows_2_it = iter(rows_2)
                next(rows_2_it)
                raise Exception('Multiple open statements')

        with self.assertRaisesRegex(QueryContextClosedError, 'Attempting to use finalized statement'):
            with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_libsqlite3=get_libsqlite3) as query:
                for columns, rows in query('''
                    SELECT my_col_a FROM my_table;
                    SELECT my_col_a FROM my_table LIMIT 10;
                '''):
                    pass

                rows_list = list(rows)

    def test_select_multi_with_named_params(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ())
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500), ()),
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-c', 'some-text-d')"] * 100), ()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            rows_list = [
                list(rows)
                for (columns, rows) in query('''
                    SELECT COUNT(*) FROM my_table WHERE my_col_a = :first;
                    SELECT COUNT(*) FROM my_table WHERE my_col_a = :second;
                ''', named_params=(((':first', 'some-text-a'),),((':second', 'some-text-c'),)))
            ]

        self.assertEqual(rows_list, [[(500,)], [(100,)]])

    def test_select_multi_with_positional_params(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);", ())
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500), ()),
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-c', 'some-text-d')"] * 100), ()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query_multi('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            rows_list = [
                list(rows)
                for (columns, rows) in query('''
                    SELECT COUNT(*) FROM my_table WHERE my_col_a = ?;
                    SELECT COUNT(*) FROM my_table WHERE my_col_a = ?;
                ''', params=(('some-text-a',), ('some-text-c',),))
            ]

        self.assertEqual(rows_list, [[(500,)], [(100,)]])

    def test_placeholder(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES ('a','b'),('c','d')",()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with query("SELECT my_col_a FROM my_table WHERE my_col_b = ?", params=(('d',))) as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [('c',)])

    def test_partial(self):
        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES ('a','b'),('c','d')",()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        query_my_db = functools.partial(sqlite_s3_query,
            url='http://localhost:9000/my-bucket/my.db',
            get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ),
            get_libsqlite3=get_libsqlite3,
        )

        with query_my_db() as query:
            with query("SELECT my_col_a FROM my_table WHERE my_col_b = ?", params=(('d',))) as (columns, rows):
                rows = list(rows)

        self.assertEqual(rows, [('c',)])

    def test_time_and_non_python_identifier(self):
        with get_db([("CREATE TABLE my_table (my_col_a text, my_col_b text);",())]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            now = datetime.datetime.utcnow()
            with query("SELECT date('now'), time('now')") as (columns, rows):
                rows = list((datetime.datetime.strptime(row[0] + ' ' + row[1], '%Y-%m-%d %H:%M:%S'),) for row in rows)

        self.assertTrue(all((row[0] - now) < datetime.timedelta(seconds=3) for row in rows))
        self.assertEqual(columns, ("date('now')", "time('now')"))

    def test_non_existant_table(self):
        with get_db([("CREATE TABLE my_table (my_col_a text, my_col_b text);",())]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(SQLiteError, 'no such table: non_table'):
                query("SELECT * FROM non_table").__enter__()

    def test_empty_object(self):
        put_object_with_versioning('my-bucket', 'my.db', lambda: (b'',))

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(HTTPStatusError, r"\b416\b"):
                query('SELECT 1').__enter__()

    def test_incorrect_permission_on_context_enter(self):
        with get_db([("CREATE TABLE my_table (my_col_a text, my_col_b text);",())]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with self.assertRaisesRegex(HTTPStatusError, r"\b403\b"):
            sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'not-the-right-key',
                None,
            ), get_libsqlite3=get_libsqlite3).__enter__()

    def test_incorrect_permission_on_run_query(self):
        with get_db([("CREATE TABLE my_table (my_col_a text, my_col_b text);",())]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        creds = (
            (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'not-the-right-key',
                None,
            )
        )
        creds_it = iter(creds)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: next(creds_it), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(HTTPStatusError, r"\b403\b"):
                query('SELECT 1').__enter__()

    def test_short_db_header(self):
        put_object_with_versioning('my-bucket', 'my.db', lambda: (b'*' * 99,))

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(SQLiteError, 'not a database'):
                query("SELECT * FROM non_table").__enter__()

    def test_bad_db_header(self):
        put_object_with_versioning('my-bucket', 'my.db', lambda: (b'*' * 100,))

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(SQLiteError, 'not a database'):
                query("SELECT * FROM non_table").__enter__()

    def test_bad_db_first_page(self):
        put_object_with_versioning('my-bucket', 'my.db', lambda: (b'*' * 4096,))

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(SQLiteError, 'not a database'):
                query("SELECT * FROM non_table").__enter__()

    def test_bad_db_second_half(self):
        with get_db([("CREATE TABLE my_table (my_col_a text, my_col_b text);",())] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
        ] * 10) as db_full:
            db = b''.join(db_full())
            half_len = int(len(db) / 2)
            db = db[:half_len] + len(db[half_len:]) * b'-'
            put_object_with_versioning('my-bucket', 'my.db', lambda: (db,))

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_libsqlite3=get_libsqlite3) as query:
            with self.assertRaisesRegex(SQLiteError, 'database disk image is malformed'):
                with query("SELECT * FROM my_table") as (columns, rows):
                    list(rows)

    def test_num_connections(self):
        num_connections = 0

        @contextmanager
        def server():
            nonlocal num_connections
            def _run(server_sock):
                nonlocal num_connections

                while True:
                    try:
                        downstream_sock, _ = server_sock.accept()
                    except Exception:
                        break
                    num_connections += 1
                    connection_t = threading.Thread(target=handle_downstream, args=(downstream_sock,))
                    connection_t.start()

            with shutdown(get_new_socket()) as server_sock:
                server_sock.bind(('127.0.0.1', 9001))
                server_sock.listen(socket.IPPROTO_TCP)
                threading.Thread(target=_run, args=(server_sock,)).start()
                yield server_sock

        def get_http_client():
            @contextmanager
            def client():
                with httpx.Client() as original_client:
                    class Client():
                        def stream(self, method, url, params, headers):
                            parsed_url = urllib.parse.urlparse(url)
                            url = urllib.parse.urlunparse(parsed_url._replace(netloc='localhost:9001'))
                            headers_proxy_host = tuple((key, value) for key, value in headers if key != 'host') + (('host', 'localhost:9000'),)
                            return original_client.stream(method, url, params=params, headers=headers_proxy_host)
                    yield Client()
            return client()

        with server() as server_sock:
            with get_db([
                ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
            ] + [
                ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
            ]) as db:
                put_object_with_versioning('my-bucket', 'my.db', db)

            with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                'us-east-1',
                'AKIAIOSFODNN7EXAMPLE',
                'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                None,
            ), get_http_client=get_http_client, get_libsqlite3=get_libsqlite3) as query:
                with query('SELECT my_col_a FROM my_table') as (columns, rows):
                    rows = list(rows)

            self.assertEqual(rows, [('some-text-a',)] * 500)
            self.assertEqual(num_connections, 1)

    def test_streaming(self):
        rows_count = 0
        rows_yielded_at_request = []

        def get_http_client():
            @contextmanager
            def client():
                with httpx.Client() as original_client:
                    class Client():
                        @contextmanager
                        def stream(self, method, url, params, headers):
                            rows_yielded_at_request.append(
                                (rows_count, dict(headers).get('range'))
                            )
                            with original_client.stream(method, url,
                                params=params, headers=headers
                            ) as response:
                                yield response
                    yield Client()
            return client()

        with get_db([
            ("PRAGMA page_size = 4096;",()),
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_http_client=get_http_client, get_libsqlite3=get_libsqlite3) as query:
            with query('SELECT my_col_a FROM my_table') as (cols, rows):
                for row in rows:
                    rows_count += 1

        self.assertIn(rows_yielded_at_request, ([
            (0, None),
            (0, 'bytes=0-99'),
            (0, 'bytes=0-4095'),
            (0, 'bytes=24-39'),  # For older SQLite that doesn't support immutable files
            (0, 'bytes=4096-8191'),
            (0, 'bytes=8192-12287'),
            (140, 'bytes=12288-16383'),
            (276, 'bytes=16384-20479'),
            (412, 'bytes=20480-24575'),
        ], [
            (0, None),
            (0, 'bytes=0-99'),
            (0, 'bytes=0-4095'),
            (0, 'bytes=4096-8191'),
            (0, 'bytes=8192-12287'),
            (140, 'bytes=12288-16383'),
            (276, 'bytes=16384-20479'),
            (412, 'bytes=20480-24575'),
        ]))

        # Documenting the difference with the above and a query that is not streaming. In this
        # case, a query with an ORDER BY on a column that does not have an index requires SQLite to
        # fetch all the pages before yielding any rows to client code
        rows_count = 0
        rows_yielded_at_request.clear()
        with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
            'us-east-1',
            'AKIAIOSFODNN7EXAMPLE',
            'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            None,
        ), get_http_client=get_http_client, get_libsqlite3=get_libsqlite3) as query:
            with query('SELECT my_col_a FROM my_table ORDER BY my_col_a') as (cols, rows):
                for row in rows:
                    rows_count += 1

        self.assertIn(rows_yielded_at_request, ([
            (0, None),
            (0, 'bytes=0-99'),
            (0, 'bytes=0-4095'),
            (0, 'bytes=24-39'),  # For older SQLite that doesn't support immutable files
            (0, 'bytes=4096-8191'),
            (0, 'bytes=8192-12287'),
            (0, 'bytes=12288-16383'),
            (0, 'bytes=16384-20479'),
            (0, 'bytes=20480-24575'),
        ], [
            (0, None),
            (0, 'bytes=0-99'),
            (0, 'bytes=0-4095'),
            (0, 'bytes=4096-8191'),
            (0, 'bytes=8192-12287'),
            (0, 'bytes=12288-16383'),
            (0, 'bytes=16384-20479'),
            (0, 'bytes=20480-24575'),
        ]))

    def test_too_many_bytes(self):
        @contextmanager
        def server(port):
            def _run(server_sock):
                while True:
                    try:
                        downstream_sock, _ = server_sock.accept()
                    except Exception:
                        break
                    connection_t = threading.Thread(target=handle_downstream, args=(downstream_sock,))
                    connection_t.start()

            with shutdown(get_new_socket()) as server_sock:
                server_sock.bind(('127.0.0.1', port))
                server_sock.listen(socket.IPPROTO_TCP)
                threading.Thread(target=_run, args=(server_sock,)).start()
                yield server_sock

        def get_http_client(port):
            @contextmanager
            def client():
                with httpx.Client() as original_client:
                    class Client():
                        @contextmanager
                        def stream(self, method, url, params, headers):
                            parsed_url = urllib.parse.urlparse(url)
                            url = urllib.parse.urlunparse(parsed_url._replace(netloc=f'localhost:{port}'))
                            range_query = dict(headers).get('range')
                            yield_extra = not only_after_header or (range_query and range_query != 'bytes=0-99')
                            headers_proxy_host = tuple((key, value) for key, value in headers if key != 'host') + (('host', 'localhost:9000'),)
                            with original_client.stream(method, url,
                                params=params, headers=headers_proxy_host
                            ) as response:
                                chunks = response.iter_bytes()
                                def iter_bytes(chunk_size=None):
                                    yield from chunks
                                    if yield_extra:
                                        yield b'e'
                                response.iter_bytes = iter_bytes
                                yield response
                    yield Client()
            return client()

        for only_after_header, port in [(False, 9001), (True, 9002)]:
            with self.subTest((only_after_header, port)):
                with server(port) as server_sock:
                    with get_db([
                        ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
                    ] + [
                        ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
                    ]) as db:
                        put_object_with_versioning('my-bucket', 'my.db', db)

                    with sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                        'us-east-1',
                        'AKIAIOSFODNN7EXAMPLE',
                        'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                        None,
                    ), get_http_client=functools.partial(get_http_client, port), get_libsqlite3=get_libsqlite3) as query:
                        with self.assertRaisesRegex(SQLiteError, 'disk I/O error'):
                            query('SELECT my_col_a FROM my_table').__enter__()

    def test_disconnection(self):
        @contextmanager
        def server():
            def _run(server_sock):
                while True:
                    try:
                        downstream_sock, _ = server_sock.accept()
                    except Exception:
                        break
                    downstream_sock.close()
                    connection_t = threading.Thread(target=handle_downstream, args=(downstream_sock,))
                    connection_t.start()

            with shutdown(get_new_socket()) as server_sock:
                server_sock.bind(('127.0.0.1', 9003))
                server_sock.listen(socket.IPPROTO_TCP)
                threading.Thread(target=_run, args=(server_sock,)).start()
                yield server_sock

        def get_http_client():
            @contextmanager
            def client():
                with httpx.Client() as original_client:
                    class Client():
                        def stream(self, method, url, headers, params):
                            parsed_url = urllib.parse.urlparse(url)
                            url = urllib.parse.urlunparse(parsed_url._replace(netloc='localhost:9003'))
                            headers_proxy_host = tuple((key, value) for key, value in headers if key != 'host') + (('host', 'localhost:9000'),)
                            return original_client.stream(method, url, headers=headers_proxy_host)
                    yield Client()
            return client()

        with get_db([
            ("CREATE TABLE my_table (my_col_a text, my_col_b text);",()),
        ] + [
            ("INSERT INTO my_table VALUES " + ','.join(["('some-text-a', 'some-text-b')"] * 500),()),
        ]) as db:
            put_object_with_versioning('my-bucket', 'my.db', db)

        with server() as server_sock:
            with self.assertRaisesRegex(Exception, 'Server disconnected|Connection reset|WinError 10053|WinError 10054'):
                sqlite_s3_query('http://localhost:9000/my-bucket/my.db', get_credentials=lambda now: (
                    'us-east-1',
                    'AKIAIOSFODNN7EXAMPLE',
                    'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                    None,
                ), get_http_client=get_http_client, get_libsqlite3=get_libsqlite3).__enter__()

def get_libsqlite3():
    return cdll.LoadLibrary(os.environ.get('LIBSQLITE3_PATH', find_library('sqlite3')))

def put_object_without_versioning(bucket, key, content):
    create_bucket(bucket)

    url = f'http://127.0.0.1:9000/{bucket}/{key}'
    sha = hashlib.sha256()
    length = 0
    for chunk in content():
        length += len(chunk)
        sha.update(chunk)
    body_hash = sha.hexdigest()
    parsed_url = urllib.parse.urlsplit(url)

    headers = aws_sigv4_headers(
        'AKIAIOSFODNN7EXAMPLE', 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
        (), 's3', 'us-east-1', parsed_url.netloc, 'PUT', parsed_url.path, (), body_hash,
    ) + ((b'content-length', str(length).encode()),)
    response = httpx.put(url, content=content(), headers=headers)
    response.raise_for_status()

def put_object_with_versioning(bucket, key, content):
    create_bucket(bucket)
    enable_versioning(bucket)

    url = f'http://127.0.0.1:9000/{bucket}/{key}'
    sha = hashlib.sha256()
    length = 0
    for chunk in content():
        length += len(chunk)
        sha.update(chunk)
    body_hash = sha.hexdigest()
    parsed_url = urllib.parse.urlsplit(url)

    headers = aws_sigv4_headers(
        'AKIAIOSFODNN7EXAMPLE', 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
        (), 's3', 'us-east-1', parsed_url.netloc, 'PUT', parsed_url.path, (), body_hash,
    ) + ((b'content-length', str(length).encode()),)

    response = httpx.put(url, content=content(), headers=headers)
    response.raise_for_status()

def create_bucket(bucket):
    url = f'http://127.0.0.1:9000/{bucket}/'
    content = b''
    body_hash = hashlib.sha256(content).hexdigest()
    parsed_url = urllib.parse.urlsplit(url)

    headers = aws_sigv4_headers(
        'AKIAIOSFODNN7EXAMPLE', 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
        (), 's3', 'us-east-1', parsed_url.netloc, 'PUT', parsed_url.path, (), body_hash,
    )
    response = httpx.put(url, content=content, headers=headers)

def enable_versioning(bucket):
    content = '''
        <VersioningConfiguration xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
            <Status>Enabled</Status>
        </VersioningConfiguration>
    '''.encode()
    url = f'http://127.0.0.1:9000/{bucket}/?versioning'
    body_hash = hashlib.sha256(content).hexdigest()
    parsed_url = urllib.parse.urlsplit(url)

    headers = aws_sigv4_headers(
        'AKIAIOSFODNN7EXAMPLE', 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
        (), 's3', 'us-east-1', parsed_url.netloc, 'PUT', parsed_url.path, (('versioning', ''),), body_hash,
    )
    response = httpx.put(url, content=content, headers=headers)
    response.raise_for_status()

def disable_auth(bucket):
    content = f'''
        {{
            "Version": "2012-10-17",
            "Statement": [
                {{
                    "Sid": "Stmt1405592139000",
                    "Effect": "Allow",
                    "Principal": "*",
                    "Action": [
                        "s3:GetObject",
                        "s3:GetObjectVersion"
                    ],
                    "Resource": [
                        "arn:aws:s3:::{bucket}/*"
                    ]
                }}
            ]
        }}
    '''.encode()
    url = f'http://127.0.0.1:9000/{bucket}/?policy'
    body_hash = hashlib.sha256(content).hexdigest()
    parsed_url = urllib.parse.urlsplit(url)

    headers = aws_sigv4_headers(
        'AKIAIOSFODNN7EXAMPLE', 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
        (), 's3', 'us-east-1', parsed_url.netloc, 'PUT', parsed_url.path, (('policy', ''),), body_hash,
    )
    response = httpx.put(url, content=content, headers=headers)
    response.raise_for_status()

def aws_sigv4_headers(access_key_id, secret_access_key, pre_auth_headers,
                      service, region, host, method, path, params, body_hash):
    algorithm = 'AWS4-HMAC-SHA256'

    now = datetime.datetime.utcnow()
    amzdate = now.strftime('%Y%m%dT%H%M%SZ')
    datestamp = now.strftime('%Y%m%d')
    credential_scope = f'{datestamp}/{region}/{service}/aws4_request'

    pre_auth_headers_lower = tuple((
        (header_key.lower(), ' '.join(header_value.split()))
        for header_key, header_value in pre_auth_headers
    ))
    required_headers = (
        ('host', host),
        ('x-amz-content-sha256', body_hash),
        ('x-amz-date', amzdate),
    )
    headers = sorted(pre_auth_headers_lower + required_headers)
    signed_headers = ';'.join(key for key, _ in headers)

    def signature():
        def canonical_request():
            canonical_uri = urllib.parse.quote(path, safe='/~')
            quoted_params = sorted(
                (urllib.parse.quote(key, safe='~'), urllib.parse.quote(value, safe='~'))
                for key, value in params
            )
            canonical_querystring = '&'.join(f'{key}={value}' for key, value in quoted_params)
            canonical_headers = ''.join(f'{key}:{value}\n' for key, value in headers)

            return f'{method}\n{canonical_uri}\n{canonical_querystring}\n' + \
                   f'{canonical_headers}\n{signed_headers}\n{body_hash}'

        def sign(key, msg):
            return hmac.new(key, msg.encode('ascii'), hashlib.sha256).digest()

        string_to_sign = f'{algorithm}\n{amzdate}\n{credential_scope}\n' + \
                         hashlib.sha256(canonical_request().encode('ascii')).hexdigest()

        date_key = sign(('AWS4' + secret_access_key).encode('ascii'), datestamp)
        region_key = sign(date_key, region)
        service_key = sign(region_key, service)
        request_key = sign(service_key, 'aws4_request')
        return sign(request_key, string_to_sign).hex()

    return (
        (b'authorization', (
            f'{algorithm} Credential={access_key_id}/{credential_scope}, '
            f'SignedHeaders={signed_headers}, Signature=' + signature()).encode('ascii')
         ),
        (b'x-amz-date', amzdate.encode('ascii')),
        (b'x-amz-content-sha256', body_hash.encode('ascii')),
    ) + pre_auth_headers


@contextmanager
def get_db(sqls):

    @contextmanager
    def temporary_directory_ignore_cleanup_errors():
        # A backport of the ignore_cleanup_errors=True parameter in
        # TemporaryDirectory added in Python 3.10
        try:
            with tempfile.TemporaryDirectory() as directory_name:
                yield directory_name
        except (PermissionError, NotADirectoryError):
            pass

    with temporary_directory_ignore_cleanup_errors() as directory_name:
        db_path = os.path.join(directory_name, 'sqlite-s3-query-test.db')
        with sqlite3.connect(db_path, isolation_level=None) as con:
            cur = con.cursor()
            cur.execute('BEGIN')
            for sql, params in sqls:
                cur.execute(sql, params)
            cur.execute('COMMIT')

        # Really close the file, especially on Windows
        del cur, con

        def db():
            with open(db_path, 'rb') as f:
                while True:
                    chunk = f.read(65536)
                    if not chunk:
                        break
                    yield chunk

        yield db


def get_new_socket():
    sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM,
                         proto=socket.IPPROTO_TCP)
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    return sock

def upstream_connect():
    upstream_sock = socket.create_connection(('127.0.0.1', 9000))
    upstream_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    return upstream_sock

@contextmanager
def shutdown(sock):
    try:
        yield sock
    finally:
        try:
            sock.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        finally:
            sock.close()

def proxy(done, source, target):
    try:
        chunk = source.recv(1)
        while chunk:
            target.sendall(chunk)
            chunk = source.recv(1)
    except OSError:
        pass
    finally:
        done.set()

def handle_downstream(downstream_sock):
    with \
            shutdown(upstream_connect()) as upstream_sock, \
            shutdown(downstream_sock) as downstream_sock:

        done = threading.Event()
        threading.Thread(target=proxy, args=(done, upstream_sock, downstream_sock)).start()
        threading.Thread(target=proxy, args=(done, downstream_sock, upstream_sock)).start()
        done.wait()