copy to returning a stream

NEWS

@@ -1,5 +1,6 @@
 xxx-xx-2019: version 1.5.0
  - max line length from 80 to 120
+ - Adds a copyto_stream method in sql (#103)
 
 Nov-29-2018: version 1.4.0
  - Adds a custom client_id argument for authentication that optionally overwrites the default (#99)

carto/sql.py

@@ -16,6 +16,7 @@
 
 from .exceptions import CartoException
 from requests import HTTPError
+from .utils import ResponseStream
 
 SQL_API_URL = 'api/{api_version}/sql'
 SQL_BATCH_API_URL = 'api/{api_version}/sql/job/'
@@ -454,3 +455,17 @@ def copyto_file_path(self, query, path, append=False):
         file_mode = 'wb' if not append else 'ab'
         with open(path, file_mode) as f:
             self.copyto_file_object(query, f)
+
+    def copyto_stream(self, query):
+        """
+        Gets data from a table into a stream
+        :param query: The "COPY { table_name [(column_name[, ...])] | (query) }
+                           TO STDOUT [WITH(option[,...])]" query to execute
+        :type query: str
+
+        :return: the data from COPY TO query
+        :rtype: raw binary (text stream)
+
+        :raise: CartoException
+        """
+        return ResponseStream(self.copyto(query))

carto/utils.py

@@ -0,0 +1,20 @@
+from io import RawIOBase
+
+
+class ResponseStream(RawIOBase):
+    def __init__(self, response):
+        self.it = response.iter_content(8 * 1024)
+        self.leftover = None
+
+    def readable(self):
+        return True
+
+    def readinto(self, b):
+        try:
+            length = len(b)
+            chunk = self.leftover or next(self.it)
+            output, self.leftover = chunk[:length], chunk[length:]
+            b[:len(output)] = output
+            return len(output)
+        except StopIteration:
+            return 0

examples/copy_and_pandas_example.py

@@ -0,0 +1,73 @@
+import argparse
+import os
+import sys
+import logging
+import pandas as pd
+
+from carto.auth import APIKeyAuthClient
+from carto.sql import SQLClient
+from carto.sql import CopySQLClient
+
+# Logger (better than print)
+logging.basicConfig(
+    level=logging.INFO,
+    format=' %(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%I:%M:%S %p')
+logger = logging.getLogger()
+
+
+# set input arguments
+parser = argparse.ArgumentParser(description='Example of CopySQLClient usage with COPY feature and pandas (file-based interface)')
+
+parser.add_argument('--base_url', type=str, dest='CARTO_BASE_URL',
+                    default=os.environ.get('CARTO_API_URL', ''),
+                    help='Set the base URL. For example:' +
+                    ' https://username.carto.com/ ' +
+                    '(defaults to env variable CARTO_API_URL)')
+
+parser.add_argument('--api_key', dest='CARTO_API_KEY',
+                    default=os.environ.get('CARTO_API_KEY', ''),
+                    help='Api key of the account' +
+                    ' (defaults to env variable CARTO_API_KEY)')
+
+
+args = parser.parse_args()
+
+# Set authentification to CARTO
+if args.CARTO_BASE_URL and args.CARTO_API_KEY:
+    auth_client = APIKeyAuthClient(
+        args.CARTO_BASE_URL, args.CARTO_API_KEY)
+else:
+    logger.error('You need to provide valid credentials, run with '
+                 '-h parameter for details')
+    sys.exit(1)
+
+# Create and cartodbfy a table
+sqlClient = SQLClient(auth_client)
+sqlClient.send("""
+    CREATE TABLE IF NOT EXISTS copy_example (
+      the_geom geometry(Geometry,4326),
+      name text,
+      age integer
+    )
+    """)
+sqlClient.send("SELECT CDB_CartodbfyTable(current_schema, 'copy_example')")
+
+copyClient = CopySQLClient(auth_client)
+
+# COPY FROM example
+logger.info("COPY'ing FROM file...")
+query = ('COPY copy_example (the_geom, name, age) '
+         'FROM stdin WITH (FORMAT csv, HEADER true)')
+result = copyClient.copyfrom_file_path(query, 'files/copy_from.csv')
+logger.info('result = %s' % result)
+
+# COPY TO example with pandas DataFrame
+logger.info("COPY'ing TO pandas DataFrame...")
+query = 'COPY copy_example TO stdout WITH (FORMAT csv, HEADER true)'
+result = copyClient.copyto_stream(query)
+df = pd.read_csv(result)
+logger.info(df.head())
+
+# Truncate the table to make this example repeatable
+sqlClient.send('TRUNCATE TABLE copy_example RESTART IDENTITY')

examples/copy_example.py

@@ -2,7 +2,6 @@
 import os
 import sys
 import logging
-import time
 
 from carto.auth import APIKeyAuthClient
 from carto.sql import SQLClient
@@ -40,7 +39,6 @@
 else:
     logger.error('You need to provide valid credentials, run with '
                  '-h parameter for details')
-    import sys
     sys.exit(1)
 
 # Create and cartodbfy a table

examples/requirements.txt

@@ -1,5 +1,6 @@
 requests>=2.7.0
 pyrestcli>=0.6.4
 configparser>=3.5.0
+pandas>=0.24.2
 ../
 prettytable

tests/test_sql_copy.py

@@ -172,3 +172,10 @@ def test_copyto_file_path(copy_client, copyto_sample_query,
     target_path = tmpdir.join('carto-python-sdk-copy-test.dump')
     copy_client.copyto_file_path(copyto_sample_query, target_path.strpath)
     assert target_path.read() == copyto_expected_result.decode('utf-8')
+
+
+def test_copyto_stream(copy_client, copyto_sample_query,
+                       copyto_expected_result):
+    response = copy_client.copyto_stream(copyto_sample_query)
+
+    assert response.read() == copyto_expected_result