Use connection URI in SqliteHook (#28721)

* Use connection URI in SqliteHook

This allows the user to define more sqlite args such as mode. See https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#uri-connections for details.
- remove unsupported schema, login and password fields in docs
- add info about host field to docs

airflow/providers/sqlite/hooks/sqlite.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import sqlite3
+from urllib.parse import unquote
 
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
@@ -33,13 +34,23 @@ class SqliteHook(DbApiHook):
 
     def get_conn(self) -> sqlite3.dbapi2.Connection:
         """Returns a sqlite connection object"""
-        conn_id = getattr(self, self.conn_name_attr)
-        airflow_conn = self.get_connection(conn_id)
-        conn = sqlite3.connect(airflow_conn.host)
+        sqlalchemy_uri = self.get_uri()
+        # The sqlite3 connection does not use the sqlite scheme.
+        # See https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#uri-connections for details.
+        sqlite_uri = sqlalchemy_uri.replace("sqlite:///", "file:")
+        conn = sqlite3.connect(sqlite_uri, uri=True)
         return conn
 
     def get_uri(self) -> str:
         """Override DbApiHook get_uri method for get_sqlalchemy_engine()"""
         conn_id = getattr(self, self.conn_name_attr)
         airflow_conn = self.get_connection(conn_id)
-        return f"sqlite:///{airflow_conn.host}"
+        if airflow_conn.conn_type is None:
+            airflow_conn.conn_type = self.conn_type
+        airflow_uri = unquote(airflow_conn.get_uri())
+        # For sqlite, there is no schema in the connection URI. So we need to drop the trailing slash.
+        airflow_sqlite_uri = airflow_uri.replace("/?", "?")
+        # The sqlite connection has one more slash for path specification.
+        # See https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#connect-strings for details.
+        sqlalchemy_uri = airflow_sqlite_uri.replace("sqlite://", "sqlite:///")
+        return sqlalchemy_uri

docs/apache-airflow-providers-sqlite/connections/sqlite.rst

@@ -23,14 +23,51 @@ The SQLite connection type provides connection to a SQLite database.
 
 Configuring the Connection
 --------------------------
-Host (required)
-    The host to connect to.
+Host (optional)
+    The host to connect to. This can either be a file on disk or an in-memory database. If not set, an in-memory database is being used.
 
-Schema (optional)
-    Specify the schema name to be used in the database.
+Extra (optional)
+    Specify the extra parameters (as json dictionary) that can be used in the sqlite connection.
+    See `Recognized Query Parameters <https://www.sqlite.org/uri.html>`_ for all supported parameters.
 
-Login (required)
-    Specify the user name to connect.
+URI format example
+^^^^^^^^^^^^^^^^^^
 
-Password (required)
-    Specify the password to connect.
+If serializing with Airflow URI:
+
+.. code-block:: bash
+
+   export AIRFLOW_CONN_SQLITE_DEFAULT='sqlite://relative/path/to/db?mode=ro'
+
+or using an absolute path:
+
+.. code-block:: bash
+
+   export AIRFLOW_CONN_SQLITE_DEFAULT='sqlite:///absolute/path/to/db?mode=ro'
+
+Note the **three** slashes after the connection type.
+
+Or using an in-memory database:
+
+.. code-block:: bash
+
+   export AIRFLOW_CONN_SQLITE_DEFAULT='sqlite://?mode=ro'
+
+When specifying the connection as an environment variable in Airflow versions prior to 2.3.0, you need to specify the connection using the URI format.
+
+Note that all components of the URI should be URL-encoded.
+
+JSON format example
+^^^^^^^^^^^^^^^^^^^
+
+If serializing with JSON:
+
+.. code-block:: bash
+
+    export AIRFLOW_CONN_SQLITE_DEFAULT='{
+        "conn_type": "sqlite",
+        "host": "relative/path/to/db",
+        "extra": {
+            "mode": "ro"
+        }
+    }'

tests/providers/sqlite/hooks/test_sqlite.py

@@ -20,6 +20,7 @@
 from unittest import mock
 from unittest.mock import patch
 
+import pytest
 import sqlalchemy
 
 from airflow.models import Connection
@@ -28,26 +29,35 @@
 
 class TestSqliteHookConn:
     def setup_method(self):
-
-        self.connection = Connection(host="host")
-
         class UnitTestSqliteHook(SqliteHook):
             conn_name_attr = "test_conn_id"
 
         self.db_hook = UnitTestSqliteHook()
-        self.db_hook.get_connection = mock.Mock()
-        self.db_hook.get_connection.return_value = self.connection
 
+    @pytest.mark.parametrize(
+        "connection, uri",
+        [
+            (Connection(host="host"), "file:host"),
+            (Connection(host="host", extra='{"mode":"ro"}'), "file:host?mode=ro"),
+            (Connection(host=":memory:"), "file::memory:"),
+            (Connection(), "file:"),
+            (Connection(uri="sqlite://relative/path/to/db?mode=ro"), "file:relative/path/to/db?mode=ro"),
+            (Connection(uri="sqlite:///absolute/path/to/db?mode=ro"), "file:/absolute/path/to/db?mode=ro"),
+            (Connection(uri="sqlite://?mode=ro"), "file:?mode=ro"),
+        ],
+    )
     @patch("airflow.providers.sqlite.hooks.sqlite.sqlite3.connect")
-    def test_get_conn(self, mock_connect):
+    def test_get_conn(self, mock_connect, connection, uri):
+        self.db_hook.get_connection = mock.Mock(return_value=connection)
         self.db_hook.get_conn()
-        mock_connect.assert_called_once_with("host")
+        mock_connect.assert_called_once_with(uri, uri=True)
 
     @patch("airflow.providers.sqlite.hooks.sqlite.sqlite3.connect")
     def test_get_conn_non_default_id(self, mock_connect):
+        self.db_hook.get_connection = mock.Mock(return_value=Connection(host="host"))
         self.db_hook.test_conn_id = "non_default"
         self.db_hook.get_conn()
-        mock_connect.assert_called_once_with("host")
+        mock_connect.assert_called_once_with("file:host", uri=True)
         self.db_hook.get_connection.assert_called_once_with("non_default")