Fix official connector to ingest caplena-com/stage-experiment DB snaposhot

Makefile

@@ -1,7 +1,7 @@
 .PHONY: test lint autoformat run ftest install release docker-build docker-run docker-push
 
-PYTHON=python3.10
-ARCH=$(shell uname -m)
+PYTHON=python
+ARCH=x86_64
 PERF8?=no
 SLOW_TEST_THRESHOLD=1 # seconds
 VERSION=$(shell cat connectors/VERSION)
@@ -14,13 +14,14 @@ bin/python:
 install: bin/python bin/elastic-ingest
 
 bin/elastic-ingest: bin/python
+	echo $(ARCH)
 	bin/pip install -r requirements/$(ARCH).txt
 	bin/python setup.py develop
 
 bin/black: bin/python
 	bin/pip install -r requirements/$(ARCH).txt
 	bin/pip install -r requirements/tests.txt
-	
+
 
 bin/pytest: bin/python
 	bin/pip install -r requirements/$(ARCH).txt
@@ -71,7 +72,7 @@ default-config: install
 	bin/elastic-ingest --action config --service-type $(SERVICE_TYPE)
 
 docker-build:
-	docker build -t docker.elastic.co/enterprise-search/elastic-connectors:$(VERSION)-SNAPSHOT .
+	docker build --platform linux/amd64 -t docker.elastic.co/enterprise-search/elastic-connectors:$(VERSION)-SNAPSHOT .
 
 docker-run:
 	docker run -v $(PWD):/config docker.elastic.co/enterprise-search/elastic-connectors:$(VERSION)-SNAPSHOT /app/bin/elastic-ingest -c /config/config.yml --log-level=DEBUG

connectors/protocol/connectors.py

@@ -923,7 +923,7 @@ def _extra(self):
         }
 
 
-IDLE_JOBS_THRESHOLD = 60  # 60 seconds
+IDLE_JOBS_THRESHOLD = 3600  # 3600 seconds
 
 
 class SyncJobIndex(ESIndex):

connectors/sources/generic_database.py

@@ -13,7 +13,7 @@
 
 WILDCARD = "*"
 
-DEFAULT_FETCH_SIZE = 50
+DEFAULT_FETCH_SIZE = 50 * 100
 DEFAULT_RETRY_COUNT = 3
 DEFAULT_WAIT_MULTIPLIER = 2
 

connectors/sources/postgresql.py

@@ -4,6 +4,7 @@
 # you may not use this file except in compliance with the Elastic License 2.0.
 #
 """Postgresql source module is responsible to fetch documents from PostgreSQL."""
+import asyncio
 import ssl
 from functools import cached_property, partial
 from urllib.parse import quote
@@ -114,7 +115,13 @@ async def get_cursor(self, query):
         """
         try:
             async with self.engine.connect() as connection:  # pyright: ignore
+                self._logger.info(
+                    f"PostgreSQLClient.get_cursor; connection.execute; query: {query}"
+                )
                 cursor = await connection.execute(text(query))
+                self._logger.info(
+                    f"PostgreSQLClient.get_cursor; cursor: {cursor}"
+                )
                 return cursor
         except Exception as exception:
             self._logger.warning(
@@ -221,19 +228,52 @@ async def data_streamer(self, schema, table):
         Yields:
             list: It will first yield the column names, then data in each row
         """
-        async for data in fetch(
-            cursor_func=partial(
-                self.get_cursor,
-                self.queries.table_data(
-                    schema=schema,
-                    table=table,
-                ),
-            ),
-            fetch_columns=True,
-            fetch_size=self.fetch_size,
-            retry_count=self.retry_count,
-        ):
-            yield data
+        try:
+            async with self.engine.connect() as connection:  # pyright: ignore
+                query = self.queries.table_data(schema=schema, table=table)
+                self._logger.info(
+                    f"PostgreSQLClient.data_streamer; connection.stream; query: {query}"
+                )
+                async with connection.stream(text(query)) as batch:
+                    self._logger.info(
+                        f"PostgreSQLClient.data_streamer; got batch; type(batch): {type(batch)}"
+                    )
+                    keys = batch.keys()
+                    self._logger.info(
+                        f"PostgreSQLClient.data_streamer; yielding keys; keys: {keys}"
+                    )
+                    yield keys
+                    async for rows in batch.partitions(size=self.fetch_size):
+                        for row in rows:
+                            yield row
+                        await asyncio.sleep(0)
+                    # while True:
+                    #     batch_size = 0
+                    #     async for row in batch.fetchmany(size=self.fetch_size):
+                    #         batch_size += 1
+                    #         yield row
+                    #     if batch_size < self.fetch_size:
+                    #         break
+                    #     await asyncio.sleep(0)
+        except Exception as exception:
+            self._logger.warning(
+                f"PostgreSQLClient.data_streamer; Something went wrong while getting cursor. Exception: {exception}"
+            )
+            raise
+
+        # async for data in fetch(
+        #     cursor_func=partial(
+        #         self.get_cursor,
+        #         self.queries.table_data(
+        #             schema=schema,
+        #             table=table,
+        #         ),
+        #     ),
+        #     fetch_columns=True,
+        #     fetch_size=self.fetch_size,
+        #     retry_count=self.retry_count,
+        # ):
+        #     yield data
 
     def _get_connect_args(self):
         """Convert string to pem format and create an SSL context
@@ -374,9 +414,15 @@ async def fetch_documents(self, table, schema):
             Dict: Document to be indexed
         """
         try:
+            self._logger.info(
+                f"PostgreSQLDataSource.fetch_documents: table: {table}; getting row count"
+            )
             row_count = await self.postgresql_client.get_table_row_count(
                 schema=schema, table=table
             )
+            self._logger.info(
+                f"PostgreSQLDataSource.fetch_documents: table: {table}; row count: {row_count}"
+            )
             if row_count > 0:
                 # Query to get the table's primary key
                 keys = await self.postgresql_client.get_table_primary_key(
@@ -385,6 +431,9 @@ async def fetch_documents(self, table, schema):
                 keys = map_column_names(
                     column_names=keys, schema=schema, tables=[table]
                 )
+                self._logger.info(
+                    f"PostgreSQLDataSource.fetch_documents: table: {table}; keys: {keys}"
+                )
                 if keys:
                     try:
                         last_update_time = (
@@ -397,6 +446,9 @@ async def fetch_documents(self, table, schema):
                             f"Unable to fetch last_updated_time for {table}"
                         )
                         last_update_time = None
+                    self._logger.info(
+                        f"PostgreSQLDataSource.fetch_documents: table: {table}; last_update_time: {last_update_time}"
+                    )
                     streamer = self.postgresql_client.data_streamer(
                         schema=schema, table=table
                     )