Adding cockroach-db intergration support for EvaDB

.gitignore

@@ -101,6 +101,7 @@ env.bak/
 venv.bak/
 env38/
 env_eva/
+evadb-venv/
 test_eva_db/
 
 # Spyder project settings

docs/source/reference/databases/cockroachdb.rst

@@ -0,0 +1,36 @@
+CockroachDB
+==========
+
+The connection to CockroackDB is based on the `psycopg2 <https://pypi.org/project/psycopg2/>`_ library.
+
+Dependency
+----------
+
+* psycopg2 
+
+
+
+Parameters
+----------
+
+Required:
+
+* `user` is the database user.
+* `password` is the database password.
+* `host` is the host name or IP address.
+* `port` is the port used to make TCP/IP connection to the CockroachDB server.
+* `database` is the database name.
+
+
+Create Connection
+-----------------
+
+.. code-block:: text
+
+   CREATE DATABASE cockroachdb_data WITH ENGINE = 'cockroachdb', PARAMETERS = {
+        "user": "eva", 
+        "password": "password",
+        "host": "localhost",
+        "port": "26257", 
+        "database": "evadb"
+   };

evadb/third_party/databases/cockroachdb/__init__.py

@@ -0,0 +1,15 @@
+# coding=utf-8
+# Copyright 2018-2023 EvaDB
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""cockroach db integrations"""

evadb/third_party/databases/cockroachdb/cockroachdb_handler.py

@@ -0,0 +1,200 @@
+# coding=utf-8
+# Copyright 2018-2023 EvaDB
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+import psycopg2
+
+from evadb.third_party.databases.types import (
+    DBHandler,
+    DBHandlerResponse,
+    DBHandlerStatus,
+)
+
+class CockroachDBHandler(DBHandler):
+    def __init__(self, name: str, **kwargs):
+        """
+        Initialize the handler.
+        Args:
+            name (str): name of the DB handler instance
+            **kwargs: arbitrary keyword arguments for establishing the connection.
+        """
+        super().__init__(name)
+        self.host = kwargs.get("host")
+        self.port = kwargs.get("port")
+        self.user = kwargs.get("user")
+        self.password = kwargs.get("password")
+        self.database = kwargs.get("database")
+        self.uri = kwargs.get("uri")
+        self.connection = None
+
+    def connect(self) -> DBHandlerStatus:
+        """
+        Set up the connection required by the handler.
+        Returns:
+            DBHandlerStatus
+        """
+        try:
+            if self.uri:
+                self.connection = psycopg2.connect(self.uri, 
+                                                cursor_factory=psycopg2.extras.RealDictCursor)
+            else:
+                self.connection = psycopg2.connect(
+                    host=self.host,
+                    port=self.port,
+                    user=self.user,
+                    password=self.password,
+                    database=self.database,
+                )
+            self.connection.autocommit = True
+            return DBHandlerStatus(status=True)
+        except psycopg2.Error as e:
+            return DBHandlerStatus(status=False, error=str(e))
+
+    def disconnect(self):
+        """
+        Close any existing connections.
+        """
+        if self.connection:
+            self.connection.close()
+
+    def check_connection(self) -> DBHandlerStatus:
+        """
+        Check connection to the handler.
+        Returns:
+            DBHandlerStatus
+        """
+        if self.connection:
+            return DBHandlerStatus(status=True)
+        else:
+            return DBHandlerStatus(status=False, error="Not connected to the database.")
+
+    def get_tables(self) -> DBHandlerResponse:
+        """
+        Return the list of tables in the database.
+        Returns:
+            DBHandlerResponse
+        """
+        if not self.connection:
+            return DBHandlerResponse(data=None, error="Not connected to the database.")
+
+        try:
+            query = "SELECT table_name FROM information_schema.tables WHERE table_schema NOT IN ('information_schema', 'pg_catalog')"
+            cursor = self.connection.cursor()
+            cursor.execute(query)
+            res = cursor.fetchall()
+            res_df = pd.DataFrame(
+                res, columns=[desc[0].lower() for desc in cursor.description]
+            )
+            return DBHandlerResponse(data=res_df)
+        except psycopg2.Error as e:
+            return DBHandlerResponse(data=None, error=str(e))
+
+    def get_columns(self, table_name: str) -> DBHandlerResponse:
+        """
+        Returns the list of columns for the given table.
+        Args:
+            table_name (str): name of the table whose columns are to be retrieved.
+        Returns:
+            DBHandlerResponse
+        """
+        if not self.connection:
+            return DBHandlerResponse(data=None, error="Not connected to the database.")
+
+        try:
+            query = f"SELECT column_name as name, data_type as dtype, udt_name FROM information_schema.columns WHERE table_name='{table_name}'"
+            self.connection.cursor().execute(query)
+            cursor = self.connection.cursor()
+
+            cursor.execute(query)
+            res = cursor.fetchall()
+            columns_df = pd.DataFrame(
+                res, columns=[desc[0].lower() for desc in cursor.description]
+            )
+            columns_df["dtype"] = columns_df.apply(
+                lambda x: self._pg_to_python_types(x["dtype"], x["udt_name"]), axis=1
+            )
+            return DBHandlerResponse(data=columns_df)
+        except psycopg2.Error as e:
+            return DBHandlerResponse(data=None, error=str(e))
+
+    def _fetch_results_as_df(self, cursor):
+        """
+        This is currently the only clean solution that we have found so far.
+        Reference to Postgres API: https://www.psycopg.org/docs/cursor.html#fetch
+        In short, currently there is no very clean programming way to differentiate
+        CREATE, INSERT, SELECT. CREATE and INSERT do not return any result, so calling
+        fetchall() on those will yield a programming error. Cursor has an attribute
+        rowcount, but it indicates # of rows that are affected. In that case, for both
+        INSERT and SELECT rowcount is not 0, so we also cannot use this API to
+        differentiate INSERT and SELECT.
+        """
+        try:
+            res = cursor.fetchall()
+            res_df = pd.DataFrame(
+                res, columns=[desc[0].lower() for desc in cursor.description]
+            )
+            return res_df
+        except psycopg2.ProgrammingError as e:
+            if str(e) == "no results to fetch":
+                return pd.DataFrame({"status": ["success"]})
+            raise e
+
+    def execute_native_query(self, query_string: str) -> DBHandlerResponse:
+        """
+        Executes the native query on the database.
+        Args:
+            query_string (str): query in native format
+        Returns:
+            DBHandlerResponse
+        """
+        if not self.connection:
+            return DBHandlerResponse(data=None, error="Not connected to the database.")
+
+        try:
+            cursor = self.connection.cursor()
+            cursor.execute(query_string)
+            return DBHandlerResponse(data=self._fetch_results_as_df(cursor))
+        except psycopg2.Error as e:
+            return DBHandlerResponse(data=None, error=str(e))
+
+    def _pg_to_python_types(self, pg_type: str, udt_name: str):
+        primitive_type_mapping = {
+            "integer": int,
+            "bigint": int,
+            "smallint": int,
+            "numeric": float,
+            "real": float,
+            "double precision": float,
+            "character": str,
+            "character varying": str,
+            "text": str,
+            "boolean": bool,
+            # Add more mappings as needed
+        }
+
+        user_defined_type_mapping = {
+            "vector": np.ndarray
+            # Handle user defined types constructed by Postgres extension.
+        }
+
+        if pg_type in primitive_type_mapping:
+            return primitive_type_mapping[pg_type]
+        elif pg_type == "USER-DEFINED" and udt_name in user_defined_type_mapping:
+            return user_defined_type_mapping[udt_name]
+        else:
+            raise Exception(
+                f"Unsupported column {pg_type} encountered in the postgres table. Please raise a feature request!"
+            )

evadb/third_party/databases/interface.py

@@ -42,6 +42,8 @@ def _get_database_handler(engine: str, **kwargs):
         return mod.MysqlHandler(engine, **kwargs)
     elif engine == "mariadb":
         return mod.MariaDbHandler(engine, **kwargs)
+    elif engine == "cockroachdb":
+        return mod.CockroachDBHandler(engine, **kwargs)
     elif engine == "clickhouse":
         return mod.ClickHouseHandler(engine, **kwargs)
     elif engine == "snowflake":