gpextprotocal.c

  1. #include "postgres.h"
  2. #include "fmgr.h"
  3. #include "funcapi.h" 
  4. #include "access/extprotocol.h"
  5. #include "catalog/pg_proc.h"
  6. #include "utils/array.h"
  7. #include "utils/builtins.h"
  8. #include "utils/memutils.h" 
  10. /* Our chosen URI format. We can change it however needed */
  11. typedef struct DemoUri 
  12. { 
  13.    char     *protocol;
  14.    char     *path;
  15. }  DemoUri; 
  16. static DemoUri *ParseDemoUri(const char *uri_str);
  17. static void FreeDemoUri(DemoUri* uri); 
  19. /* Do the module magic dance */ 
  20. PG_MODULE_MAGIC; 
  21. PG_FUNCTION_INFO_V1(demoprot_export); 
  22. PG_FUNCTION_INFO_V1(demoprot_import); 
  23. PG_FUNCTION_INFO_V1(demoprot_validate_urls); 
  25. Datum demoprot_export(PG_FUNCTION_ARGS); 
  26. Datum demoprot_import(PG_FUNCTION_ARGS); 
  27. Datum demoprot_validate_urls(PG_FUNCTION_ARGS); 
  29. /* A user context that persists across calls. Can be 
  30. declared in any other way */
  31. typedef struct { 
  32.   char    *url; 
  33.   char    *filename; 
  34.   FILE    *file; 
  35. } extprotocol_t; 
  36. /* 
  37. * The read function - Import data into GPDB.
  38. */ 
  39. Datum 
  40. myprot_import(PG_FUNCTION_ARGS) 
  41. { 
  42.   extprotocol_t   *myData; 
  43.   char            *data; 
  44.   int             datlen; 
  45.   size_t          nread = 0; 
  47.   /* Must be called via the external table format manager */ 
  48.   if (!CALLED_AS_EXTPROTOCOL(fcinfo)) 
  49.     elog(ERROR, "myprot_import: not called by external
  50.        protocol manager"); 
  52.   /* Get our internal description of the protocol */ 
  53.   myData = (extprotocol_t *) EXTPROTOCOL_GET_USER_CTX(fcinfo); 
  55.   if(EXTPROTOCOL_IS_LAST_CALL(fcinfo)) 
  56.   { 
  57.     /* we're done receiving data. close our connection */ 
  58.     if(myData && myData->file) 
  59.       if(fclose(myData->file)) 
  60.         ereport(ERROR, 
  61.           (errcode_for_file_access(), 
  62.            errmsg("could not close file \"%s\": %m", 
  63.                myData->filename))); 
  65.     PG_RETURN_INT32(0); 
  66.   }
  68.   if (myData == NULL) 
  69.   { 
  70.     /* first call. do any desired init */ 
  72.     const char    *p_name = "myprot"; 
  73.     DemoUri       *parsed_url; 
  74.     char          *url = EXTPROTOCOL_GET_URL(fcinfo); 
  75.     myData        = palloc(sizeof(extprotocol_t)); 
  77.     myData->url   = pstrdup(url); 
  78.     parsed_url    = ParseDemoUri(myData->url); 
  79.     myData->filename = pstrdup(parsed_url->path); 
  81.     if(strcasecmp(parsed_url->protocol, p_name) != 0) 
  82.       elog(ERROR, "internal error: myprot called with a
  83.           different protocol (%s)", 
  84.             parsed_url->protocol); 
  86.     FreeDemoUri(parsed_url); 
  88.     /* open the destination file (or connect to remote server in
  89.        other cases) */ 
  90.     myData->file = fopen(myData->filename, "r"); 
  92.     if (myData->file == NULL) 
  93.       ereport(ERROR, 
  94.           (errcode_for_file_access(), 
  95.            errmsg("myprot_import: could not open file \"%s\"
  96.              for reading: %m", 
  97.              myData->filename), 
  98.            errOmitLocation(true))); 
  100.     EXTPROTOCOL_SET_USER_CTX(fcinfo, myData); 
  101.   }
  102.   /* ========================================== 
  103.    *          DO THE IMPORT 
  104.    * ========================================== */ 
  105.   data    = EXTPROTOCOL_GET_DATABUF(fcinfo); 
  106.   datlen  = EXTPROTOCOL_GET_DATALEN(fcinfo); 
  108.   /* read some bytes (with fread in this example, but normally
  109.      in some other method over the network) */
  110.   if(datlen > 0) 
  111.   { 
  112.     nread = fread(data, 1, datlen, myData->file); 
  113.     if (ferror(myData->file)) 
  114.       ereport(ERROR, 
  115.         (errcode_for_file_access(), 
  116.           errmsg("myprot_import: could not write to file
  117.             \"%s\": %m", 
  118.             myData->filename))); 
  119.   }
  120.   PG_RETURN_INT32((int)nread); 
  121. }
  122. /* 
  123.  * Write function - Export data out of GPDB 
  124.  */ 
  125. Datum  
  126. myprot_export(PG_FUNCTION_ARGS) 
  127. { 
  128.   extprotocol_t  *myData; 
  129.   char           *data; 
  130.   int            datlen; 
  131.   size_t         wrote = 0; 
  133.   /* Must be called via the external table format manager */ 
  134.   if (!CALLED_AS_EXTPROTOCOL(fcinfo)) 
  135.     elog(ERROR, "myprot_export: not called by external
  136.        protocol manager"); 
  138.   /* Get our internal description of the protocol */ 
  139.   myData = (extprotocol_t *) EXTPROTOCOL_GET_USER_CTX(fcinfo); 
  140.   if(EXTPROTOCOL_IS_LAST_CALL(fcinfo)) 
  141.   { 
  142.     /* we're done sending data. close our connection */ 
  143.     if(myData && myData->file) 
  144.       if(fclose(myData->file)) 
  145.         ereport(ERROR, 
  146.             (errcode_for_file_access(), 
  147.               errmsg("could not close file \"%s\": %m", 
  148.                  myData->filename))); 
  150.     PG_RETURN_INT32(0); 
  151.   }
  152.   if (myData == NULL) 
  153.   { 
  154.     /* first call. do any desired init */ 
  155.     const char *p_name = "myprot"; 
  156.     DemoUri    *parsed_url; 
  157.     char       *url = EXTPROTOCOL_GET_URL(fcinfo); 
  159.     myData           = palloc(sizeof(extprotocol_t)); 
  161.     myData->url      = pstrdup(url); 
  162.     parsed_url       = ParseDemoUri(myData->url); 
  163.     myData->filename = pstrdup(parsed_url->path); 
  165.     if(strcasecmp(parsed_url->protocol, p_name) != 0) 
  166.       elog(ERROR, "internal error: myprot called with a 
  167.          different protocol (%s)", 
  168.          parsed_url->protocol); 
  170.     FreeDemoUri(parsed_url); 
  172.     /* open the destination file (or connect to remote server in
  173.     other cases) */ 
  174.     myData->file = fopen(myData->filename, "a"); 
  175.     if (myData->file == NULL) 
  176.       ereport(ERROR, 
  177.         (errcode_for_file_access(), 
  178.            errmsg("myprot_export: could not open file \"%s\"
  179.              for writing: %m", 
  180.              myData->filename), 
  181.          errOmitLocation(true))); 
  183.     EXTPROTOCOL_SET_USER_CTX(fcinfo, myData); 
  184.   } 
  185.   /* ======================================== 
  186.    *      DO THE EXPORT 
  187.    * ======================================== */ 
  188.   data   = EXTPROTOCOL_GET_DATABUF(fcinfo); 
  189.   datlen   = EXTPROTOCOL_GET_DATALEN(fcinfo); 
  191.   if(datlen > 0) 
  192.   { 
  193.     wrote = fwrite(data, 1, datlen, myData->file); 
  195.     if (ferror(myData->file)) 
  196.       ereport(ERROR, 
  197.         (errcode_for_file_access(), 
  198.          errmsg("myprot_import: could not read from file
  199.             \"%s\": %m", 
  200.             myData->filename))); 
  201.   } 
  202.   PG_RETURN_INT32((int)wrote); 
  203. } 
  204. Datum  
  205. myprot_validate_urls(PG_FUNCTION_ARGS) 
  206. { 
  207.   List         *urls; 
  208.   int          nurls; 
  209.   int          i; 
  210.   ValidatorDirection  direction; 
  212.   /* Must be called via the external table format manager */ 
  213.   if (!CALLED_AS_EXTPROTOCOL_VALIDATOR(fcinfo)) 
  214.     elog(ERROR, "myprot_validate_urls: not called by external
  215.        protocol manager");
  217.   nurls       = EXTPROTOCOL_VALIDATOR_GET_NUM_URLS(fcinfo); 
  218.   urls        = EXTPROTOCOL_VALIDATOR_GET_URL_LIST(fcinfo); 
  219.   direction   = EXTPROTOCOL_VALIDATOR_GET_DIRECTION(fcinfo); 
  220.   /* 
  221.    * Dumb example 1: search each url for a substring  
  222.    * we don't want to be used in a url. in this example 
  223.    * it's 'secured_directory'. 
  224.    */ 
  225.   for (i = 1 ; i <= nurls ; i++) 
  226.   { 
  227.     char *url = EXTPROTOCOL_VALIDATOR_GET_NTH_URL(fcinfo, i); 
  229.     if (strstr(url, "secured_directory") != 0) 
  230.     { 
  231.       ereport(ERROR, 
  232.        (errcode(ERRCODE_PROTOCOL_VIOLATION), 
  233.           errmsg("using 'secured_directory' in a url
  234.             isn't allowed "))); 
  235.     } 
  236.   } 
  237.   /* 
  238.    * Dumb example 2: set a limit on the number of urls  
  239.    * used. In this example we limit readable external 
  240.    * tables that use our protocol to 2 urls max. 
  241.    */ 
  242.   if(direction == EXT_VALIDATE_READ && nurls > 2) 
  243.   { 
  244.     ereport(ERROR, 
  245.       (errcode(ERRCODE_PROTOCOL_VIOLATION), 
  246.         errmsg("more than 2 urls aren't allowed in this protocol "))); 
  247.   }
  248.   PG_RETURN_VOID(); 
  249. }
  250. /* --- utility functions --- */ 
  251. static  
  252. DemoUri *ParseDemoUri(const char *uri_str) 
  253. { 
  254.   DemoUri *uri = (DemoUri *) palloc0(sizeof(DemoUri)); 
  255.   int     protocol_len; 
  257.    uri->path = NULL; 
  258.    uri->protocol = NULL; 
  259.   /* 
  260.    * parse protocol 
  261.    */ 
  262.   char *post_protocol = strstr(uri_str, "://"); 
  264.   if(!post_protocol) 
  265.   { 
  266.     ereport(ERROR, 
  267.       (errcode(ERRCODE_SYNTAX_ERROR), 
  268.        errmsg("invalid protocol URI \'%s\'", uri_str), 
  269.        errOmitLocation(true))); 
  270.   }
  272.   protocol_len = post_protocol - uri_str; 
  273.   uri->protocol = (char *)palloc0(protocol_len + 1); 
  274.   strncpy(uri->protocol, uri_str, protocol_len); 
  276.   /* make sure there is more to the uri string */ 
  277.   if (strlen(uri_str) <= protocol_len) 
  278.     ereport(ERROR, 
  279.       (errcode(ERRCODE_SYNTAX_ERROR), 
  280.        errmsg("invalid myprot URI \'%s\' : missing path",
  281.          uri_str), 
  282.       errOmitLocation(true))); 
  284.   /* parse path */ 
  285.   uri->path = pstrdup(uri_str + protocol_len + strlen("://"));
  287.   return uri; 
  288. }
  289. static 
  290. void FreeDemoUri(DemoUri *uri) 
  291. { 
  292.   if (uri->path) 
  293.     pfree(uri->path); 
  294.   if (uri->protocol) 
  295.     pfree(uri->protocol); 
  297.   pfree(uri); 
  298. }

上级主题： 安装外部表协议