mseep-txtai 9.1.1__py3-none-any.whl

txtai/database/duckdb.py

@@ -0,0 +1,150 @@
+"""
+DuckDB module
+"""
+
+import os
+import re
+
+from tempfile import TemporaryDirectory
+
+# Conditional import
+try:
+    import duckdb
+
+    DUCKDB = True
+except ImportError:
+    DUCKDB = False
+
+from .embedded import Embedded
+from .schema import Statement
+
+
+class DuckDB(Embedded):
+    """
+    Database instance backed by DuckDB.
+    """
+
+    # Delete single document and object
+    DELETE_DOCUMENT = "DELETE FROM documents WHERE id = ?"
+    DELETE_OBJECT = "DELETE FROM objects WHERE id = ?"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not DUCKDB:
+            raise ImportError('DuckDB is not available - install "database" extra to enable')
+
+    def execute(self, function, *args):
+        # Call parent method with DuckDB compatible arguments
+        return super().execute(function, *self.formatargs(args))
+
+    def insertdocument(self, uid, data, tags, entry):
+        # Delete existing document
+        self.cursor.execute(DuckDB.DELETE_DOCUMENT, [uid])
+
+        # Call parent method
+        super().insertdocument(uid, data, tags, entry)
+
+    def insertobject(self, uid, data, tags, entry):
+        # Delete existing object
+        self.cursor.execute(DuckDB.DELETE_OBJECT, [uid])
+
+        # Call parent method
+        super().insertobject(uid, data, tags, entry)
+
+    def connect(self, path=":memory:"):
+        # Create connection and start a transaction
+        # pylint: disable=I1101
+        connection = duckdb.connect(path)
+        connection.begin()
+
+        return connection
+
+    def getcursor(self):
+        return self.connection
+
+    def jsonprefix(self):
+        # Return json column prefix
+        return "json_extract_string(data"
+
+    def jsoncolumn(self, name):
+        # Generate json column using json_extract function
+        return f"json_extract_string(data, '$.{name}')"
+
+    def rows(self):
+        # Iteratively retrieve and yield rows
+        batch = 256
+        rows = self.cursor.fetchmany(batch)
+        while rows:
+            yield from rows
+            rows = self.cursor.fetchmany(batch)
+
+    def addfunctions(self):
+        # DuckDB doesn't currently support scalar functions
+        return
+
+    def copy(self, path):
+        # Delete existing file, if necessary
+        if os.path.exists(path):
+            os.remove(path)
+
+        # Create database connection
+        # pylint: disable=I1101
+        connection = duckdb.connect(path)
+
+        # List of tables
+        tables = ["documents", "objects", "sections"]
+
+        with TemporaryDirectory() as directory:
+            # Export existing tables
+            for table in tables:
+                self.connection.execute(f"COPY {table} TO '{directory}/{table}.parquet' (FORMAT parquet)")
+
+            # Create initial schema
+            for schema in [Statement.CREATE_DOCUMENTS, Statement.CREATE_OBJECTS, Statement.CREATE_SECTIONS % "sections"]:
+                connection.execute(schema)
+
+            # Import tables into new schema
+            for table in tables:
+                connection.execute(f"COPY {table} FROM '{directory}/{table}.parquet' (FORMAT parquet)")
+
+            # Create indexes and sync data to database file
+            connection.execute(Statement.CREATE_SECTIONS_INDEX)
+            connection.execute("CHECKPOINT")
+
+        # Start transaction
+        connection.begin()
+
+        return connection
+
+    def formatargs(self, args):
+        """
+        DuckDB doesn't support named parameters. This method replaces named parameters with question marks
+        and makes parameters a list.
+
+        Args:
+            args: input arguments
+
+        Returns:
+            DuckDB compatible args
+        """
+
+        if args and len(args) > 1:
+            # Unpack query args
+            query, parameters = args
+
+            # Iterate over parameters
+            #   - Replace named parameters with ?'s
+            #   - Build list of value with position indexes
+            params = []
+            for key, value in parameters.items():
+                pattern = rf"\:{key}(?=\s|$)"
+                match = re.search(pattern, query)
+                if match:
+                    query = re.sub(pattern, "?", query, count=1)
+                    params.append((match.start(), value))
+
+            # Repack query and parameter list
+            args = (query, [value for _, value in sorted(params, key=lambda x: x[0])])
+
+        return args

txtai/database/embedded.py

@@ -0,0 +1,76 @@
+"""
+Embedded module
+"""
+
+from .rdbms import RDBMS
+
+
+class Embedded(RDBMS):
+    """
+    Base class for embedded relational databases. An embedded relational database stores all content in a local file.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new Database.
+
+        Args:
+            config: database configuration parameters
+        """
+
+        super().__init__(config)
+
+        # Path to database file
+        self.path = None
+
+    def load(self, path):
+        # Call parent logic
+        super().load(path)
+
+        # Store path reference
+        self.path = path
+
+    def save(self, path):
+        # Temporary database
+        if not self.path:
+            # Save temporary database
+            self.connection.commit()
+
+            # Copy data from current to new
+            connection = self.copy(path)
+
+            # Close temporary database
+            self.connection.close()
+
+            # Point connection to new connection
+            self.session(connection=connection)
+            self.path = path
+
+        # Paths are equal, commit changes
+        elif self.path == path:
+            self.connection.commit()
+
+        # New path is different from current path, copy data and continue using current connection
+        else:
+            self.copy(path).close()
+
+    def jsonprefix(self):
+        # Return json column prefix
+        return "json_extract(data"
+
+    def jsoncolumn(self, name):
+        # Generate json column using json_extract function
+        return f"json_extract(data, '$.{name}')"
+
+    def copy(self, path):
+        """
+        Copies the current database into path.
+
+        Args:
+            path: path to write database
+
+        Returns:
+            new connection with data copied over
+        """
+
+        raise NotImplementedError

txtai/database/encoder/__init__.py

@@ -0,0 +1,8 @@
+"""
+Encoder imports
+"""
+
+from .base import Encoder
+from .factory import EncoderFactory
+from .image import ImageEncoder
+from .serialize import SerializeEncoder

txtai/database/encoder/base.py

@@ -0,0 +1,37 @@
+"""
+Encoder module
+"""
+
+from io import BytesIO
+
+
+class Encoder:
+    """
+    Encodes and decodes object content. The base encoder works only with byte arrays. It can be extended to encode different datatypes.
+    """
+
+    def encode(self, obj):
+        """
+        Encodes an object to a byte array using the encoder.
+
+        Args:
+            obj: object to encode
+
+        Returns:
+            encoded object as a byte array
+        """
+
+        return obj
+
+    def decode(self, data):
+        """
+        Decodes input byte array into an object using this encoder.
+
+        Args:
+            data: encoded data
+
+        Returns:
+            decoded object
+        """
+
+        return BytesIO(data) if data else None

txtai/database/encoder/factory.py

@@ -0,0 +1,56 @@
+"""
+Encoder factory module
+"""
+
+from ...util import Resolver
+
+from .base import Encoder
+from .serialize import SerializeEncoder
+
+
+class EncoderFactory:
+    """
+    Encoder factory. Creates new Encoder instances.
+    """
+
+    @staticmethod
+    def get(encoder):
+        """
+        Gets a new instance of encoder class.
+
+        Args:
+            encoder: Encoder instance class
+
+        Returns:
+            Encoder class
+        """
+
+        # Local task if no package
+        if "." not in encoder:
+            # Get parent package
+            encoder = ".".join(__name__.split(".")[:-1]) + "." + encoder.capitalize() + "Encoder"
+
+        return Resolver()(encoder)
+
+    @staticmethod
+    def create(encoder):
+        """
+        Creates a new Encoder instance.
+
+        Args:
+            encoder: Encoder instance class
+
+        Returns:
+            Encoder
+        """
+
+        # Return default encoder
+        if encoder is True:
+            return Encoder()
+
+        # Supported serialization methods
+        if encoder in ["messagepack", "pickle"]:
+            return SerializeEncoder(encoder)
+
+        # Get Encoder instance
+        return EncoderFactory.get(encoder)()

txtai/database/encoder/image.py

@@ -0,0 +1,43 @@
+"""
+ImageEncoder module
+"""
+
+from io import BytesIO
+
+# Conditional import
+try:
+    from PIL import Image
+
+    PIL = True
+except ImportError:
+    PIL = False
+
+from .base import Encoder
+
+
+class ImageEncoder(Encoder):
+    """
+    Encodes and decodes Image objects as compressed binary content, using the original image's algorithm.
+    """
+
+    def __init__(self):
+        """
+        Creates a new ImageEncoder.
+        """
+
+        if not PIL:
+            raise ImportError('ImageEncoder is not available - install "database" extra to enable')
+
+    def encode(self, obj):
+        # Create byte stream
+        output = BytesIO()
+
+        # Write image to byte stream
+        obj.save(output, format=obj.format, quality="keep")
+
+        # Return byte array
+        return output.getvalue()
+
+    def decode(self, data):
+        # Return a PIL image
+        return Image.open(BytesIO(data)) if data else None

txtai/database/encoder/serialize.py

@@ -0,0 +1,28 @@
+"""
+SerializeEncoder module
+"""
+
+from ...serialize import SerializeFactory
+
+from .base import Encoder
+
+
+class SerializeEncoder(Encoder):
+    """
+    Encodes and decodes objects using the internal serialize package.
+    """
+
+    def __init__(self, method):
+        # Parent constructor
+        super().__init__()
+
+        # Pickle serialization
+        self.serializer = SerializeFactory.create(method)
+
+    def encode(self, obj):
+        # Pickle object
+        return self.serializer.savebytes(obj)
+
+    def decode(self, data):
+        # Unpickle to object
+        return self.serializer.loadbytes(data)

txtai/database/factory.py

@@ -0,0 +1,77 @@
+"""
+Factory module
+"""
+
+from urllib.parse import urlparse
+
+from ..util import Resolver
+
+from .client import Client
+from .duckdb import DuckDB
+from .sqlite import SQLite
+
+
+class DatabaseFactory:
+    """
+    Methods to create document databases.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Create a Database.
+
+        Args:
+            config: database configuration parameters
+
+        Returns:
+            Database
+        """
+
+        # Database instance
+        database = None
+
+        # Enables document database
+        content = config.get("content")
+
+        # Standardize content name
+        if content is True:
+            content = "sqlite"
+
+        # Create document database instance
+        if content == "duckdb":
+            database = DuckDB(config)
+        elif content == "sqlite":
+            database = SQLite(config)
+        elif content:
+            # Check if content is a URL
+            url = urlparse(content)
+            if content == "client" or url.scheme:
+                # Connect to database server URL
+                database = Client(config)
+            else:
+                # Resolve custom database if content is not a URL
+                database = DatabaseFactory.resolve(content, config)
+
+        # Store config back
+        config["content"] = content
+
+        return database
+
+    @staticmethod
+    def resolve(backend, config):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: index configuration parameters
+
+        Returns:
+            Database
+        """
+
+        try:
+            return Resolver()(backend)(config)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve database backend: '{backend}'") from e