SNOW-143457 import sqlalchemy

Author: sfc-gh-stakeda

__init__.py

@@ -7,7 +7,14 @@
 from . import base
 from . import snowdialect
 from .custom_commands import (
-    MergeInto, CSVFormatter, JSONFormatter, PARQUETFormatter, CopyIntoStorage, AWSBucket, AzureContainer
+    MergeInto,
+    CSVFormatter,
+    JSONFormatter,
+    PARQUETFormatter,
+    CopyIntoStorage,
+    AWSBucket,
+    AzureContainer,
+    ExternalStage
 )
 from .util import _url as URL
 from .version import VERSION
@@ -95,5 +102,6 @@
     'PARQUETFormatter',
     'CopyIntoStorage',
     'AWSBucket',
-    'AzureContainer'
+    'AzureContainer',
+    'ExternalStage',
 )

base.py

@@ -224,6 +224,9 @@ def visit_copy_formatter(self, formatter, **kw):
                                                         value))
                                                     for name, value in options_list])) if formatter.options else "")
 
+    def visit_external_stage(self, stage, **kw):
+        return "@{}{}{}".format(stage.namespace, stage.name, stage.path)
+
     def visit_aws_bucket(self, aws_bucket, **kw):
         credentials_list = list(aws_bucket.credentials_used.items())
         if kw.get('deterministic', False):

custom_commands.py

@@ -289,6 +289,27 @@ def snappy_compression(self, comp):
         return self
 
 
+class ExternalStage(ClauseElement):
+    """External Stage descriptor"""
+    __visit_name__ = "external_stage"
+
+    @staticmethod
+    def prepare_namespace(namespace):
+        return "{}.".format(namespace) if not namespace.endswith(".") else namespace
+
+    @staticmethod
+    def prepare_path(path):
+        return "/{}".format(path) if not path.startswith("/") else path
+
+    def __init__(self, name, path=None, namespace=None):
+        self.name = name
+        self.path = self.prepare_path(path) if path else ""
+        self.namespace = self.prepare_namespace(namespace) if namespace else ""
+
+    def __repr__(self):
+        return "@{}{}{}".format(self.namespace, self.name, self.path)
+
+
 class AWSBucket(ClauseElement):
     """AWS S3 bucket descriptor"""
     __visit_name__ = 'aws_bucket'

test/test_copy.py

@@ -9,13 +9,27 @@
     AzureContainer,
     CopyIntoStorage,
     CSVFormatter,
+    ExternalStage,
     JSONFormatter,
     PARQUETFormatter,
 )
 from sqlalchemy import Column, Integer, MetaData, Sequence, String, Table
 from sqlalchemy.sql import select
 
 
+def test_external_stage(sql_compiler):
+    assert ExternalStage.prepare_namespace("something") == "something."
+    assert ExternalStage.prepare_path("prefix") == "/prefix"
+
+    # All arguments are handled
+    assert (
+        sql_compiler(ExternalStage(name="name", path="prefix/path", namespace="namespace")) == "@namespace.name/prefix/path"
+    )
+
+    # defaults don't ruin things
+    assert sql_compiler(ExternalStage(name="name", path=None, namespace=None)) == "@name"
+
+
 def test_copy_into_location(engine_testaccount, sql_compiler):
     meta = MetaData()
     conn = engine_testaccount.connect()
@@ -71,6 +85,12 @@ def test_copy_into_location(engine_testaccount, sql_compiler):
                                          "FIELD_DELIMITER=',') ENCRYPTION="
                                          "(KMS_KEY_ID='1234abcd-12ab-34cd-56ef-1234567890ab' TYPE='AWS_SSE_KMS')")
 
+    copy_stmt_6 = CopyIntoStorage(from_=food_items, into=ExternalStage(name="stage_name"), formatter=CSVFormatter())
+    assert sql_compiler(copy_stmt_6) == "COPY INTO @stage_name FROM python_tests_foods FILE_FORMAT=(TYPE=csv)"
+
+    copy_stmt_7 = CopyIntoStorage(from_=food_items, into=ExternalStage(name="stage_name", path="prefix/file", namespace="name"), formatter=CSVFormatter())
+    assert sql_compiler(copy_stmt_7) == "COPY INTO @name.stage_name/prefix/file FROM python_tests_foods FILE_FORMAT=(TYPE=csv)"
+
     # NOTE Other than expect known compiled text, submit it to RegressionTests environment and expect them to fail, but
     # because of the right reasons
     try: